Update app.py

app.py

@@ -144,17 +144,8 @@ class Translators:
         else:
             prompt = f"Translate the following segment into {self.tl}, without additional explanation.\n\n{self.input_text}."
         tokenizer = AutoTokenizer.from_pretrained(self.model_name)
-        model = AutoModelForCausalLM.from_pretrained(self.model_name, device_map="auto", torch_dtype=torch.bfloat16)
+        model = AutoModelForCausalLM.from_pretrained(self.model_name, device_map="auto", dtype=torch.bfloat16)
         messages = [{"role": "user", "content": prompt}]
-        # tokenized_chat = tokenizer.apply_chat_template(
-        #     messages,
-        #     tokenize=True,
-        #     add_generation_prompt=True,
-        #     return_tensors="pt"
-        # )
-        # outputs = model.generate(tokenized_chat.to(model.device), max_new_tokens=512, top_k=20, top_p=0.6, repetition_penalty=1.05, temperature=0.7)
-        # output_text = tokenizer.decode(outputs[0])
-        # return output_text
         # Tokenize the conversation
         tokenized_chat = tokenizer.apply_chat_template(
             messages,
@@ -169,14 +160,18 @@ class Translators:
                 tokenized_chat.to(model.device),
                 max_new_tokens=512,
                 temperature=temperature,
-                top_p=0.6,
+                top_k=20
+                top_p=0.95,
+                repetition_penalty=1.05,
                 do_sample=True if temperature > 0 else False,
                 pad_token_id=tokenizer.eos_token_id
             )
-        
+  
+        # outputs = model.generate(tokenized_chat.to(model.device), max_new_tokens=512, top_k=20, top_p=0.6, repetition_penalty=1.05, temperature=0.7)
+        output_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
         # Decode only the new tokens
-        response = tokenizer.decode(outputs[0][tokenized_chat.shape[-1]:], skip_special_tokens=True)
-        return response
+        # output_text = tokenizer.decode(outputs[0][tokenized_chat.shape[-1]:], skip_special_tokens=True)
+        return output_text
     
     def HelsinkiNLP_mulroa(self):
         try: