add nllb translate codes

app.py

@@ -24,7 +24,7 @@ from gradio.themes.utils import colors, fonts, sizes
 import argparse
 
 import langid
-from transformers import pipeline
+from transformers import pipeline, AutoModelForSeq2SeqLM, AutoTokenizer
 
 
 class myTheme(Base):
@@ -112,8 +112,6 @@ def opus_trans(article, target_language):
         target_lang = "en"
     elif target_language == "Chinese":
         target_lang = "zh"
-    elif target_language == "Spanish":
-        target_lang = "es"
 
     if result_lang != target_lang:
         task_name = f"translation_{result_lang}_to_{target_lang}"
@@ -129,15 +127,31 @@ def opus_trans(article, target_language):
 
 
 def nllb_trans(article, target_language):
-    pass
+    tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")
+    model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M")
+    inputs = tokenizer(article, return_tensors="pt")
+
+    if target_language == "English":
+        target_lang = "Eng_Latn"
+    elif target_language == "Chinese":
+        target_lang = "zho_Hans"
+
+    translated_tokens = model.generate(
+        **inputs,
+        forced_bos_token_id=tokenizer.lang_code_to_id[target_lang],
+        max_length=30,
+    )
+
+    return tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]
 
 
 def translate(article, toolkit, target_language):
     if toolkit == "OPUS":
         translated = opus_trans(article, target_language)
-        return translated
     elif toolkit == "NLLB":
-        pass
+        translated = nllb_trans(article, target_language)
+
+    return translated
 
 
 myTheme = myTheme()