Update app.py

app.py

@@ -11,6 +11,7 @@ blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image
 translator_tokenizer = NllbTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")
 translator_model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M")
 lang_code = "vie_Latn"
+lang_token_id = translator_tokenizer.lang_code_to_id[lang_code]
 
 def caption_translate(image):
     # Step 1: Get English caption
@@ -19,11 +20,15 @@ def caption_translate(image):
     eng_caption = blip_processor.decode(out[0], skip_special_tokens=True)
 
     # Step 2: Translate to Vietnamese
-    inputs = translator_tokenizer(eng_caption, return_tensors="pt", src_lang="eng_Latn", tgt_lang=lang_code)
-    translated = translator_model.generate(**inputs, max_length=100)
+    inputs = translator_tokenizer(eng_caption, return_tensors="pt", src_lang="eng_Latn")
+    translated = translator_model.generate(
+        **inputs,
+        forced_bos_token_id=lang_token_id,
+        max_length=100
+    )
     vi_caption = translator_tokenizer.decode(translated[0], skip_special_tokens=True)
 
-    return f"📷 Mô tả: {vi_caption}\n\n(English: {eng_caption})"
+    return vi_caption
 
 iface = gr.Interface(
     fn=caption_translate,