Update app.py

app.py

@@ -22,7 +22,7 @@ iso1_to_name = {iso[1]: iso[0] for iso in non_empty_isos} # {'ro': 'Romanian', '
 models = ["Helsinki-NLP", "QUICKMT", "Argos", "Google", "HPLT", "t5-base", "t5-small", "t5-large",
           "utter-project/EuroLLM-1.7B", "utter-project/EuroLLM-1.7B-Instruct",
           "Unbabel/Tower-Plus-2B", "Unbabel/TowerInstruct-7B-v0.2", "Unbabel/TowerInstruct-Mistral-7B-v0.2",
-          "HuggingFaceTB/SmolLM3-3B", "winninghealth/WiNGPT-Babel-2", "tencent/Hunyuan-MT-7B-fp8",
+          "HuggingFaceTB/SmolLM3-3B", "winninghealth/WiNGPT-Babel-2", "tencent/Hunyuan-MT-7B",
          "openGPT-X/Teuken-7B-instruct-commercial-v0.4", "openGPT-X/Teuken-7B-instruct-v0.6"]
 allmodels = ["Helsinki-NLP",
           "Helsinki-NLP/opus-mt-tc-bible-big-mul-mul", "Helsinki-NLP/opus-mt-tc-bible-big-mul-deu_eng_nld",
@@ -144,22 +144,39 @@ 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")
-        model.tie_weights() # fp8
+        model = AutoModelForCausalLM.from_pretrained(self.model_name, device_map="auto", torch_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,
             tokenize=True,
-            add_generation_prompt=False,
-            return_tensors="pt",
-            top_k=20,
-            top_p=0.6,
-            repetition_penalty=1.05,
-            temperature=0.7
-        )      
-        outputs = model.generate(tokenized_chat.to(model.device), max_new_tokens=512)
-        output_text = tokenizer.decode(outputs[0])
-        return output_text
+            add_generation_prompt=True,
+            return_tensors="pt"
+        )
+        # Generate response
+        temperature = 0.7
+        with torch.no_grad():
+            outputs = model.generate(
+                tokenized_chat.to(model.device),
+                max_new_tokens=512,
+                temperature=temperature,
+                top_p=0.6,
+                do_sample=True if temperature > 0 else False,
+                pad_token_id=tokenizer.eos_token_id
+            )
+        
+        # Decode only the new tokens
+        response = tokenizer.decode(outputs[0][tokenized_chat.shape[-1]:], skip_special_tokens=True)
+        return response
     
     def HelsinkiNLP_mulroa(self):
         try: