Update app.py

app.py

@@ -56,6 +56,7 @@ class Translators:
         self.sl, self.tl = sl, tl
         self.input_text = input_text
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.max_new_tokens = 512
         
     def google(self):
         # for rep in ('\r\n', '\r', '\n', '  '):
@@ -242,6 +243,35 @@ class Translators:
         except Exception as error:
             return f"Error translating with model: {self.model_name}! Try other available language combination or model.", error
     
+     def rosetta(self):
+        model = AutoModelForCausalLM.from_pretrained(
+            self.model_name,
+            dtype=torch.bfloat16,
+            device_map="auto",
+            max_memory={0: "16GB"},
+        )
+        tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+        context = {"tone": "Formal and precise"}        
+        system = f"Translate the user's text to {self.tl}. Provide the final translation in a formal tone immediately immediately without any other text."
+        messages = [
+            {"role": "system", "content": system},
+            {"role": "user", "content": self.input_text},
+        ]
+        
+        prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        print(prompt)
+        inputs = tokenizer(prompt, return_tensors="pt").to(self.device)
+        input_length = inputs["input_ids"].shape[1]       
+        with torch.inference_mode():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=self.max_new_tokens,
+            )
+        
+        generated_tokens = outputs[0][input_length:]
+        translation = tokenizer.decode(generated_tokens, skip_special_tokens=True)
+        return translation    
+    
     def salamandratapipe(self):
         pipe = pipeline("text-generation", model=self.model_name)
         messages = [{"role": "user", "content": f"Translate the following text from {self.sl} into {self.tl}.\n{self.sl}: {self.input_text} \n{self.tl}:"}]
@@ -631,7 +661,7 @@ def translate_text(model_name: str, s_language: str, t_language: str, input_text
             translated_text = Translators(model_name, s_language, t_language, input_text).salamandratapipe()
         
         elif "rosetta" in model_name.lower():
-            translated_text = Translators(model_name, s_language, t_language, input_text).simplepipe()
+            translated_text = Translators(model_name, s_language, t_language, input_text).rosetta()
         
         elif "small100" in model_name.lower():
             translated_text = Translators(model_name, sl, tl, input_text).smallonehundred()