Update app.py

app.py

@@ -6,8 +6,8 @@ st.header("Text Machine Translation")
 input_text = st.text_input("Enter text to translate:")
 # Create a list of options for the select box
 options = ["German", "Romanian", "English", "French", "Spanish"]
-langs = {"English":"en", "Romanian":"ro", "German":"de", "French":"fr", "Spanish":"es"}
-models = ["Helsinki-NLP", "t5-base", "t5-small", "t5-large"]
+langs = {"English":"en", "Romanian":"ro", "German":"de", "French":"fr", "Spanish":"es", "Italian":"it"}
+models = ["Helsinki-NLP", "t5-base", "t5-small", "t5-large", "Unbabel/TowerInstruct-7B-v0.2"]
 
 # Create two columns
 scol, tcol = st.columns(2)
@@ -34,9 +34,10 @@ if model_name == 'Helsinki-NLP':
         model_name = f"Helsinki-NLP/opus-tatoeba-{sl}-{tl}"
         tokenizer = AutoTokenizer.from_pretrained(model_name)
         model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-else:
+if model_name.startswith('t5'):
     tokenizer = T5Tokenizer.from_pretrained(model_name)
     model = T5ForConditionalGeneration.from_pretrained(model_name)
+
 st.write("Selected language combination:", sselected_language, " - ", tselected_language, "Selected model:", model_name)
 submit_button = st.button("Translate")
 translated_textarea = st.text("")
@@ -45,14 +46,30 @@ translated_textarea = st.text("")
 if submit_button:
     if model_name.startswith('Helsinki-NLP'):
         prompt = input_text
-    else:
+        print(prompt)
+        input_ids = tokenizer.encode(prompt, 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)
+    elif model_name.startswith('t5'):
         prompt = f'translate {sselected_language} to {tselected_language}: {input_text}'
-    print(prompt)
-    input_ids = tokenizer.encode(prompt, 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)
+        print(prompt)
+        input_ids = tokenizer.encode(prompt, 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)
+    else:
+        pipe = pipeline("text-generation", model="Unbabel/TowerInstruct-7B-v0.2", 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=True)
+        outputs = pipe(prompt, max_new_tokens=256, do_sample=False)
+        translated_text = outputs[0]["generated_text"]
+
     # Display the translated text
     print(translated_text)
     st.write(f"Translated text from {sselected_language} to {tselected_language} using {model_name}")