Update app.py

app.py

@@ -554,71 +554,73 @@ with st.container(border=None, width="stretch", height="content", horizontal=Fal
 # Handle the submit button click
 if submit_button:
     with st.spinner("Translating...", show_time=True):
-        if model_name.startswith('Helsinki-NLP'):
-            # input_ids = tokenizer.encode(input_text, return_tensors='pt')
-            # # Perform translation
-            # output_ids = model.generate(input_ids)
-            # # Decode the translated text
-            # translated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-            # Use a pipeline as a high-level helper
-            try:
-                model_name = f"Helsinki-NLP/opus-mt-{sl}-{tl}"
-                tokenizer = AutoTokenizer.from_pretrained(model_name)
-                model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-                pipe = pipeline("translation", model=model, tokenizer=tokenizer)
-            except (EnvironmentError, OSError):
-                model_name = f"Helsinki-NLP/opus-tatoeba-{sl}-{tl}"
-                tokenizer = AutoTokenizer.from_pretrained(model_name)
-                model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-                pipe = pipeline("translation", model=model, tokenizer=tokenizer)
-            translation = pipe(input_text)
-            translated_text = translation[0]['translation_text']
+        translated_text, message = translate_text(model_name, sselected_language, sselected_language, input_text: str)
+        # if model_name.startswith('Helsinki-NLP'):
+        #     # input_ids = tokenizer.encode(input_text, return_tensors='pt')
+        #     # # Perform translation
+        #     # output_ids = model.generate(input_ids)
+        #     # # Decode the translated text
+        #     # translated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+        #     # Use a pipeline as a high-level helper
+        #     try:
+        #         model_name = f"Helsinki-NLP/opus-mt-{sl}-{tl}"
+        #         tokenizer = AutoTokenizer.from_pretrained(model_name)
+        #         model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+        #         pipe = pipeline("translation", model=model, tokenizer=tokenizer)
+        #     except (EnvironmentError, OSError):
+        #         model_name = f"Helsinki-NLP/opus-tatoeba-{sl}-{tl}"
+        #         tokenizer = AutoTokenizer.from_pretrained(model_name)
+        #         model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+        #         pipe = pipeline("translation", model=model, tokenizer=tokenizer)
+        #     translation = pipe(input_text)
+        #     translated_text = translation[0]['translation_text']
     
-        elif model_name.startswith('Google'): 
-            url = os.environ['GCLIENT'] + f'sl={sl}&tl={tl}&q={input_text}'
-            response = httpx.get(url)
-            translated_text = response.json()[0][0][0]
-            print(response.json()[0][0])
-        elif model_name.startswith('t5'):
-            tokenizer = T5Tokenizer.from_pretrained(model_name)
-            model = T5ForConditionalGeneration.from_pretrained(model_name).to(device)
-            prompt = f'translate {sselected_language} to {tselected_language}: {input_text}'
-            input_ids = tokenizer.encode(prompt, return_tensors='pt').to(device)
-            # Perform translation
-            output_ids = model.generate(input_ids)
-            # Decode the translated text
-            translated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-        elif 'Unbabel' in model_name:   
-            pipe = pipeline("text-generation", model=model_name, torch_dtype=torch.bfloat16, device_map="auto")
-            # We use the tokenizer’s chat template to format each message - see https://huggingface.co/docs/transformers/main/en/chat_templating
-            messages = [{"role": "user",
-                         "content": f"Translate the following text from {sselected_language} into {tselected_language}.\n{sselected_language}: {input_text}.\n{tselected_language}:"}]
-            prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
-            outputs = pipe(prompt, max_new_tokens=256, do_sample=False)
-            translated_text = outputs[0]["generated_text"]
-            start_marker = "<end_of_turn>"
-            if start_marker in translated_text:
-                translated_text = translated_text.split(start_marker)[1].strip()
-            translated_text = translated_text.replace('Answer:', '').strip() if translated_text.startswith('Answer:') else translated_text
-        elif 'Argos' in model_name:
-            import argostranslate.translate       
-            # Translate
-            try:
-                download_argos_model(sl, tl)
-                translated_text = argostranslate.translate.translate(input_text, sl, tl)
-            except StopIteration:
-                translated_text = f"No Argos model for {sselected_language} to {tselected_language}. Try other model or languages combination!"
-            except Exception as error:
-                translated_text = error
-        elif model_name == "winninghealth/WiNGPT-Babel-2":      
-            translated_text = wingpt(model_name, sselected_language, tselected_language, input_text)
+        # elif model_name.startswith('Google'): 
+        #     url = os.environ['GCLIENT'] + f'sl={sl}&tl={tl}&q={input_text}'
+        #     response = httpx.get(url)
+        #     translated_text = response.json()[0][0][0]
+        #     print(response.json()[0][0])
+        # elif model_name.startswith('t5'):
+        #     tokenizer = T5Tokenizer.from_pretrained(model_name)
+        #     model = T5ForConditionalGeneration.from_pretrained(model_name).to(device)
+        #     prompt = f'translate {sselected_language} to {tselected_language}: {input_text}'
+        #     input_ids = tokenizer.encode(prompt, return_tensors='pt').to(device)
+        #     # Perform translation
+        #     output_ids = model.generate(input_ids)
+        #     # Decode the translated text
+        #     translated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+        # elif 'Unbabel' in model_name:   
+        #     pipe = pipeline("text-generation", model=model_name, torch_dtype=torch.bfloat16, device_map="auto")
+        #     # We use the tokenizer’s chat template to format each message - see https://huggingface.co/docs/transformers/main/en/chat_templating
+        #     messages = [{"role": "user",
+        #                  "content": f"Translate the following text from {sselected_language} into {tselected_language}.\n{sselected_language}: {input_text}.\n{tselected_language}:"}]
+        #     prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
+        #     outputs = pipe(prompt, max_new_tokens=256, do_sample=False)
+        #     translated_text = outputs[0]["generated_text"]
+        #     start_marker = "<end_of_turn>"
+        #     if start_marker in translated_text:
+        #         translated_text = translated_text.split(start_marker)[1].strip()
+        #     translated_text = translated_text.replace('Answer:', '').strip() if translated_text.startswith('Answer:') else translated_text
+        # elif 'Argos' in model_name:
+        #     import argostranslate.translate       
+        #     # Translate
+        #     try:
+        #         download_argos_model(sl, tl)
+        #         translated_text = argostranslate.translate.translate(input_text, sl, tl)
+        #     except StopIteration:
+        #         translated_text = f"No Argos model for {sselected_language} to {tselected_language}. Try other model or languages combination!"
+        #     except Exception as error:
+        #         translated_text = error
+        # elif model_name == "winninghealth/WiNGPT-Babel-2":      
+        #     translated_text = wingpt(model_name, sselected_language, tselected_language, input_text)
         
     # Display the translated text
     print(f"Translated from {sselected_language} to {tselected_language} using {model_name}.", input_text, translated_text)
-    message = f"Translated from {sselected_language} to {tselected_language} using {model_name}."
+    # message = f"Translated from {sselected_language} to {tselected_language} using {model_name}."
     # translated_textarea.text_area(":green[Translation:]", value=translated_text)
     # message_textarea.text_input(":blue[Message:]", value=message)
     st.text_area(":green[Translation:]", placeholder="Translation area", value=translated_text)
-    st.success(message)
+    # st.success(message, icon=":material/check:") st.info(message, icon="ℹ️"), st.warning(message, icon=":material/warning:"), error(message, icon=":material/error:"), st.exception
+    st.info(message, icon=":material/info:")
     # st.text_input(":blue[Messages:]", value=message)
     # st.rerun()