Update app.py

app.py

@@ -76,6 +76,15 @@ def bigscience(model_name, sl, tl, input_text):
     translation = translation.replace('<pad> ', '').replace('</s>', '')
     return translation
 
+def bloomz(model_name, sl, tl, input_text):  
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+    inputs = tokenizer.encode(f"Translate to {tl}: {input_text}", return_tensors="pt")
+    outputs = model.generate(inputs)
+    translation = tokenizer.decode(outputs[0])
+    translation = translation.replace('<pad> ', '').replace('</s>', '')
+    return translation
+
 def eurollm(model_name, sl, tl, input_text):
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForCausalLM.from_pretrained(model_name)  
@@ -119,9 +128,13 @@ def translate_text(input_text, sselected_language, tselected_language, model_nam
         translated_text = teuken(model_name, sselected_language, tselected_language, input_text)
         return translated_text, message_text
 
-    if 'bigscience' in model_name.lower():
+    if 'mt0' in model_name.lower():
         translated_text = bigscience(model_name, sselected_language, tselected_language, input_text)
         return translated_text, message_text
+
+    if 'bloomz' in model_name.lower():
+        translated_text = bloomz(model_name, sselected_language, tselected_language, input_text)
+        return translated_text, message_text
         
     if 'nllb' in model_name.lower():
         nnlbsl, nnlbtl = languagecodes.nllb_language_codes[sselected_language], languagecodes.nllb_language_codes[tselected_language]