Update app.py

app.py

@@ -7,19 +7,18 @@ import httpx
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 # Language options and mappings
-options = ["German", "Romanian", "English", "French", "Spanish", "Italian", "Hindi"]
-favourite_langs = {"German": "de", "Romanian": "ro", "English": "en", "-----": "-----"}
+favourite_langs = {"Romanian": "ro", "German": "de", "English": "en", "-----": "-----"}
+# langs = ["German", "Romanian", "English", "French", "Spanish", "Italian",]
 df = pl.read_parquet("isolanguages.parquet")
 non_empty_isos = df.slice(1).filter(pl.col("ISO639-1") != "").rows()
 # all_langs = languagecodes.iso_languages_byname
 all_langs = {iso[0]: (iso[1], iso[2], iso[3]) for iso in non_empty_isos} # {'Romanian': ('ro', 'rum', 'ron')}
-# langs = list(favourite_langs.keys())
-# langs.extend(list(all_langs.keys())) # Language options as list, add favourite languages first
+langs = list(favourite_langs.keys())
+langs.extend(list(all_langs.keys())) # Language options as list, add favourite languages first
 # iso1_to_name = {codes[0]: lang for entry in all_langs for lang, codes in entry.items()} # {'ro': 'Romanian', 'de': 'German'}
 iso1_to_name = {iso[1]: iso[0] for iso in non_empty_isos} # {'ro': 'Romanian', 'de': 'German'}
 langs = {iso[0]: iso[1] for iso in non_empty_isos} # {'Romanian': 'ro', 'German': 'de'}
 
-
 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",
@@ -139,9 +138,26 @@ class Translators:
         return translated_text
 
     def hunyuan(self):
-        pipe = pipeline("translation", model=self.model_name, device=self.device)
-        translated_text = pipe(self.input_text, src_lang=self.sl, tgt_lang=self.tl)
-        return translation
+        ZH_CODES = {"Chinese": "zh", "Traditional Chinese": "zh-Hant", "Cantonese": "yue"}
+        if self.sl in ZH_CODES.keys() or self.tl in ZH_CODES.keys():
+            prompt = f"把下面的文本翻译成{self.tl}，不要额外解释。\n\n{self.input_text}"
+        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")
+        messages = [{"role": "user", "content": prompt}]
+        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)
+        return output_text = tokenizer.decode(outputs[0])
     
     def HelsinkiNLP_mulroa(self):
         try:
@@ -564,9 +580,9 @@ input_text = st.text_area("Enter text to translate:", placeholder="Enter text to
 
 # Initialize session state if not already set
 if "sselected_language" not in st.session_state:
-    st.session_state["sselected_language"] = options[0]
+    st.session_state["sselected_language"] = langs[0]
 if "tselected_language" not in st.session_state:
-    st.session_state["tselected_language"] = options[1]
+    st.session_state["tselected_language"] = langs[1]
 if "model_name" not in st.session_state:
     st.session_state["model_name"] = models[1]
 
@@ -578,8 +594,8 @@ model_name = st.selectbox("Select a model:", models,
 scol, swapcol, tcol = st.columns([3, 1, 3])
 
 with scol:
-    sselected_language = st.selectbox("Source language:", options, 
-                                      index=options.index(st.session_state["sselected_language"]))
+    sselected_language = st.selectbox("Source language:", langs, 
+                                      index=langs.index(st.session_state["sselected_language"]))
 with swapcol:
     if st.button("🔄 Swap"):
         st.session_state["model_name"] = model_name  # Preserve model
@@ -587,8 +603,8 @@ with swapcol:
             st.session_state["tselected_language"], st.session_state["sselected_language"]
         st.rerun()
 with tcol:
-    tselected_language = st.selectbox("Target language:", options, 
-                                      index=options.index(st.session_state["tselected_language"]))
+    tselected_language = st.selectbox("Target language:", langs, 
+                                      index=langs.index(st.session_state["tselected_language"]))
 
 # Language codes
 sl = langs[st.session_state["sselected_language"]]