Update app.py

app.py

@@ -89,13 +89,19 @@ def translate_text(text, source_lang, target_lang):
     src_code = LANGUAGE_CODES.get(source_lang, "eng_Latn")
     tgt_code = LANGUAGE_CODES.get(target_lang, "ara_Arab")
     
+    # Format target language token with double underscores according to NLLB format
+    tgt_token = f"__{tgt_code}__"
+    
     # Tokenize
     inputs = translator_tokenizer(text, return_tensors="pt", padding=True)
     
-    # Translate
+    # Get the token ID for the target language
+    forced_bos_token_id = translator_tokenizer.convert_tokens_to_ids(tgt_token)
+    
+    # Generate translation with the target language token
     translated_tokens = translator_model.generate(
         **inputs,
-        forced_bos_token_id=translator_tokenizer.lang_code_to_id[tgt_code],
+        forced_bos_token_id=forced_bos_token_id,
         max_length=128
     )