Update app.py

app.py

@@ -6,45 +6,46 @@ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 model_name = "jsbeaudry/creole-translation-nllb-600M"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+# Assuming the necessary imports and model loading from the previous code are in place
 
-# Supported languages
-language_codes = {
-    "English": "eng_Latn",
-    "Haitian Creole": "hat_Latn"
-}
+def translate_text(text):
+    # Set the source and target language codes
+    src_lang = "eng_Latn"
+    tgt_lang = "hat_Latn"
+    tokenizer_ = tokenizer
+    model_ = model
 
-# Translation function
-def translate_text(text, src_lang_name, tgt_lang_name):
-    src_lang = language_codes[src_lang_name]
-    tgt_lang = language_codes[tgt_lang_name]
+    # Set tokenizer to source language
+    tokenizer_.src_lang = src_lang
 
-    tokenizer.src_lang = src_lang
-    inputs = tokenizer(text, return_tensors="pt")
-    forced_bos_token_id = tokenizer.convert_tokens_to_ids(tgt_lang)
+    # Tokenize the input
+    inputs = tokenizer_(text, return_tensors="pt")
 
+    # Find the BOS token ID for the target language
+    forced_bos_token_id = tokenizer_.convert_tokens_to_ids(tgt_lang)
+
+    # Move model to GPU if available
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    model.to(device)
+    model_ = model_.to(device)
     inputs = inputs.to(device)
 
-    generated_tokens = model.generate(
+    # Generate translation
+    generated_tokens = model_.generate(
         **inputs,
         forced_bos_token_id=forced_bos_token_id,
         max_length=100
     )
-    translated = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
+
+    # Decode and print
+    translated = tokenizer_.batch_decode(generated_tokens, skip_special_tokens=True)
     return translated[0]
 
-# Gradio interface
 iface = gr.Interface(
     fn=translate_text,
-    inputs=[
-        gr.Textbox(lines=5, placeholder="Enter text to translate", label="Input Text"),
-        gr.Dropdown(choices=list(language_codes.keys()), value="English", label="Source Language"),
-        gr.Dropdown(choices=list(language_codes.keys()), value="Haitian Creole", label="Target Language")
-    ],
-    outputs=gr.Textbox(label="Translated Text"),
-    title="Multilingual Translation (English ↔ Haitian Creole)",
-    description="Translate text between English and Haitian Creole using a fine-tuned NLLB model."
+    inputs=gr.Textbox(lines=5, placeholder="Enter text to translate"),
+    outputs="text",
+    title="English to Haitian Creole Translation",
+    description="Translate English text to Haitian Creole using a fine-tuned NLLB model."
 )
 
-iface.launch()
+iface.launch()