Update app.py

app.py

@@ -553,72 +553,72 @@ with st.container(border=None, width="stretch", height="content", horizontal=Fal
 
 # Handle the submit button click
 if submit_button:
-    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)
+    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']
+    
+        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}."
     # translated_textarea.text_area(":green[Translation:]", value=translated_text)
     # message_textarea.text_input(":blue[Message:]", value=message)
-    with st.spinner("Translating...", show_time=True):
-        st.success(message)
-        st.text_area(":green[Translation:]", placeholder="Translation area", value=translated_text)
-        st.text_input(":blue[Messages:]", placeholder="Messages area", value=message)
+    st.text_area(":green[Translation:]", placeholder="Translation area", value=translated_text)
+    st.success(message)
+    st.text_input(":blue[Messages:]", value=message)
     # st.rerun()