Update app.py

app.py

@@ -28,7 +28,7 @@ input_text = st.text_input("Enter text to translate:")
 options = ["German", "Romanian", "English", "French", "Spanish", "Italian"]
 langs = {"English": "en", "Romanian": "ro", "German": "de", "French": "fr", "Spanish": "es", "Italian": "it"}
 models = ["Helsinki-NLP", "Argos", "t5-base", "t5-small", "t5-large", "Unbabel/Tower-Plus-2B", 
-          "Unbabel/TowerInstruct-Mistral-7B-v0.2", "Google"]
+          "Unbabel/TowerInstruct-Mistral-7B-v0.2", "winninghealth/WiNGPT-Babel-2", "Google"]
 
 # Initialize session state if not already set
 if "sselected_language" not in st.session_state:
@@ -125,7 +125,7 @@ if submit_button:
         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
-    if 'Argos' in model_name:   
+    if 'Argos' in model_name:
         import argostranslate.translate       
         # Translate
         try:
@@ -135,6 +135,10 @@ if submit_button:
             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
+    if model_name == "winninghealth/WiNGPT-Babel-2":
+        pipe = pipeline("translation", model="winninghealth/WiNGPT-Babel-2")
+        translation = pipe(input_text)
+        translated_text = translation[0]['translation_text']
     # Display the translated text
     print(translated_text)
     st.write(f"Translated text from {sselected_language} to {tselected_language} using {model_name}:")