Update app.py

app.py

@@ -7,23 +7,28 @@ import os
 os.environ["HF_TOKEN"] = os.getenv('hf')
 
 # Load the tokenizer and model from Hugging Face
-tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")
-model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M")
+model_name = "facebook/nllb-200-distilled-600M"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+
+# Mapping for language codes
+lang_codes = {
+    "English to Iban": ("eng_Latn", "iba_Latn"),
+    "Iban to English": ("iba_Latn", "eng_Latn")
+}
 
 # Translation function
 def translate_text(text, translation_direction):
-    if translation_direction == "English to Iban":
-        src_lang = "eng_Latn"
-        tgt_lang = "iba_Latn"
-    elif translation_direction == "Iban to English":
-        src_lang = "iba_Latn"
-        tgt_lang = "eng_Latn"
-    else:
+    if translation_direction not in lang_codes:
         return json.dumps({"error": "Invalid translation direction"})
 
-    # Tokenize and translate
-    inputs = tokenizer(text, return_tensors="pt", src_lang=src_lang)
-    outputs = model.generate(**inputs, forced_bos_token_id=tokenizer.lang_code_to_id[tgt_lang])
+    src_lang, tgt_lang = lang_codes[translation_direction]
+
+    # Tokenize input text
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
+
+    # Generate translation
+    outputs = model.generate(**inputs, forced_bos_token_id=tokenizer.convert_tokens_to_ids(f"[{tgt_lang}]"))
     translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
     # Return JSON response
@@ -45,5 +50,5 @@ iface = gr.Interface(
     description="This app translates between English and Iban. Choose the translation direction and enter your text."
 )
 
-# Launch the app with shareable link
+# Launch the app with a shareable link
 iface.launch(share=True)