Create nllb.py

nllb.py

@@ -0,0 +1,62 @@
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+from flores200_codes import flores_codes
+
+model_dict = {}
+
+
+def load_models(model_name: str):
+    # build model and tokenizer
+    model_name_dict = {
+        "nllb-1.3B": "facebook/nllb-200-1.3B",
+        "nllb-distilled-1.3B": "facebook/nllb-200-distilled-1.3B",
+        "nllb-3.3B": "facebook/nllb-200-3.3B",
+    }[model_name]
+
+    print("\tLoading model: %s" % model_name)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name_dict)
+    tokenizer = AutoTokenizer.from_pretrained(model_name_dict)
+    model_dict[model_name + "_model"] = model
+    model_dict[model_name + "_tokenizer"] = tokenizer
+
+    return model_dict
+
+
+def translation(model_name: str, source, target, text: str):
+
+    model_dict = load_models(model_name)
+
+    source = flores_codes[source]
+    target = flores_codes[target]
+
+    model = model_dict[model_name + "_model"]
+    tokenizer = model_dict[model_name + "_tokenizer"]
+
+    translator = pipeline(
+        "translation",
+        model=model,
+        tokenizer=tokenizer,
+        src_lang=source,
+        tgt_lang=target,
+    )
+    output = translator(text, max_length=400)
+
+    output = output[0]["translation_text"]
+    result = {
+        "source": source,
+        "target": target,
+        "result": output,
+    }
+
+    return result
+
+
+NLLB_EXAMPLES = [
+    ["nllb-distilled-1.3B", "English", "Shan", "Hello, how are you today?"],
+    ["nllb-distilled-1.3B", "Shan", "English", "မႂ်ႇသုင်ၶႃႈ ယူႇလီယူႇၶႃႈၼေႃႈ"],
+    [
+        "nllb-distilled-1.3B",
+        "English",
+        "Shan",
+        "Forming Myanmar’s New Political System Will Remain an Ideal but Never in Practicality",
+    ],
+]