Upload 9 files

.gitattributes

@@ -2,20 +2,27 @@
 *.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
 *.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
 *.ftz filter=lfs diff=lfs merge=lfs -text
 *.gz filter=lfs diff=lfs merge=lfs -text
 *.h5 filter=lfs diff=lfs merge=lfs -text
 *.joblib filter=lfs diff=lfs merge=lfs -text
 *.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
 *.model filter=lfs diff=lfs merge=lfs -text
 *.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
 *.onnx filter=lfs diff=lfs merge=lfs -text
 *.ot filter=lfs diff=lfs merge=lfs -text
 *.parquet filter=lfs diff=lfs merge=lfs -text
 *.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
 *.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
 saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.tar.* filter=lfs diff=lfs merge=lfs -text
 *.tflite filter=lfs diff=lfs merge=lfs -text
@@ -23,5 +30,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.wasm filter=lfs diff=lfs merge=lfs -text
 *.xz filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
-*.zstandard filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text

.gitignore.txt

@@ -0,0 +1,2 @@
+.devcontainer/*
+ui.cpython-310.pyc

README.md

@@ -1,13 +1,13 @@
 ---
-title: Nllb Translation Demo
-emoji: 👀
+title: NLLB200 Translate Distill 600
+emoji: 🐢
 colorFrom: indigo
-colorTo: green
+colorTo: yellow
 sdk: gradio
-sdk_version: 3.0.26
+sdk_version: 3.18.0
 app_file: app.py
 pinned: false
-duplicated_from: Geonmo/nllb-translation-demo
+license: mit
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -1,106 +1,55 @@
-import os
-import torch
 import gradio as gr
-import time
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
-from flores200_codes import flores_codes
-
-
-def load_models():
-    # build model and tokenizer
-    model_name_dict = {#'nllb-distilled-600M': 'facebook/nllb-200-distilled-600M',
-                  '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_dict = {}
-
-    for call_name, real_name in model_name_dict.items():
-        print('\tLoading model: %s' % call_name)
-        model = AutoModelForSeq2SeqLM.from_pretrained(real_name)
-        tokenizer = AutoTokenizer.from_pretrained(real_name)
-        model_dict[call_name+'_model'] = model
-        model_dict[call_name+'_tokenizer'] = tokenizer
-
-    return model_dict
-
-
-def load_models():
-    # build model and tokenizer
-    model_name_dict = {'nllb-distilled-600M': 'facebook/nllb-200-distilled-600M',
-                  #'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_dict = {}
-
-    for call_name, real_name in model_name_dict.items():
-        print('\tLoading model: %s' % call_name)
-        model = AutoModelForSeq2SeqLM.from_pretrained(real_name)
-        tokenizer = AutoTokenizer.from_pretrained(real_name)
-        model_dict[call_name+'_model'] = model
-        model_dict[call_name+'_tokenizer'] = tokenizer
-
-    return model_dict
-
-
-def translation(source, target, text):
-    if len(model_dict) == 2:
-        model_name = 'nllb-distilled-600M'
-
-    start_time = time.time()
-    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)
-
-    end_time = time.time()
-
-    output = output[0]['translation_text']
-    result = {'inference_time': end_time - start_time,
-              'source': source,
-              'target': target,
-              'result': output}
-    return result
-
-
-
-if __name__ == '__main__':
-    print('\tinit models')
-
-    global model_dict
-
-    model_dict = load_models()
+import torch
+from ui import title, description, examples
+from langs import LANGS
+#from langs_all import LANGS  ##for 200+ languages
+
+TASK = "translation"
+CKPT = "facebook/nllb-200-distilled-1.3B"
+#CKPT = "facebook/nllb-200-distilled-600M"
+
+model = AutoModelForSeq2SeqLM.from_pretrained(CKPT)
+tokenizer = AutoTokenizer.from_pretrained(CKPT)
+
+# device = 0 if torch.cuda.is_available() else -1
+
+
+def translate(text, src_lang, tgt_lang, max_length=512):
+    """
+    Translate the text from source lang to target lang
+    """
+    translation_pipeline = pipeline(TASK,
+                                    model=model,
+                                    tokenizer=tokenizer,
+                                    src_lang=src_lang,
+                                    tgt_lang=tgt_lang,
+                                    max_length=max_length)
     
-    # define gradio demo
-    lang_codes = list(flores_codes.keys())
-    #inputs = [gr.inputs.Radio(['nllb-distilled-600M', 'nllb-1.3B', 'nllb-distilled-1.3B'], label='NLLB Model'),
-    inputs = [gr.inputs.Dropdown(lang_codes, default='English', label='Source'),
-              gr.inputs.Dropdown(lang_codes, default='Korean', label='Target'),
-              gr.inputs.Textbox(lines=5, label="Input text"),
-              ]
-
-    outputs = gr.outputs.JSON()
-
-    title = "NLLB distilled 1.3B demo"
-
-    demo_status = "Demo is running on CPU"
-    description = f"Details: https://github.com/facebookresearch/fairseq/tree/nllb. {demo_status}"
-    examples = [
-    ['English', 'Korean', 'Hi. nice to meet you']
-    ]
-
-    gr.Interface(translation,
-                 inputs,
-                 outputs,
-                 title=title,
-                 description=description,
-                 ).launch()
-
-
+    # translation_pipeline = pipeline(TASK,
+    #                                 model=model,
+    #                                 tokenizer=tokenizer,
+    #                                 src_lang=src_lang,
+    #                                 tgt_lang=tgt_lang,
+    #                                 max_length=max_length,
+    #                                 device=device)
+
+    result = translation_pipeline(text)
+    return result[0]['translation_text']
+
+
+gr.Interface(
+    translate,
+    [
+        gr.components.Textbox(label="Text"),
+        gr.components.Dropdown(label="Source Language", choices=LANGS),
+        gr.components.Dropdown(label="Target Language", choices=LANGS),
+        gr.components.Slider(8, 512, value=512, step=8, label="Max Length")
+    ],
+    ["text"],
+    examples=examples,
+    # article=article,
+    cache_examples=False,
+    title=title,
+    description=description
+).launch()

langs.py

@@ -0,0 +1,8 @@
+LANGS = [
+    "bod_Tibt",
+    "khk_Cyrl",
+    "uig_Arab",
+    "yue_Hant",
+    "zho_Hans",
+    "zho_Hant"
+]

langs_all.py

@@ -0,0 +1,204 @@
+LANGS = [
+    "ace_Arab",
+    "ace_Latn",
+    "acm_Arab",
+    "acq_Arab",
+    "aeb_Arab",
+    "afr_Latn",
+    "ajp_Arab",
+    "aka_Latn",
+    "amh_Ethi",
+    "apc_Arab",
+    "arb_Arab",
+    "ars_Arab",
+    "ary_Arab",
+    "arz_Arab",
+    "asm_Beng",
+    "ast_Latn",
+    "awa_Deva",
+    "ayr_Latn",
+    "azb_Arab",
+    "azj_Latn",
+    "bak_Cyrl",
+    "bam_Latn",
+    "ban_Latn",
+    "bel_Cyrl",
+    "bem_Latn",
+    "ben_Beng",
+    "bho_Deva",
+    "bjn_Arab",
+    "bjn_Latn",
+    "bod_Tibt",
+    "bos_Latn",
+    "bug_Latn",
+    "bul_Cyrl",
+    "cat_Latn",
+    "ceb_Latn",
+    "ces_Latn",
+    "cjk_Latn",
+    "ckb_Arab",
+    "crh_Latn",
+    "cym_Latn",
+    "dan_Latn",
+    "deu_Latn",
+    "dik_Latn",
+    "dyu_Latn",
+    "dzo_Tibt",
+    "ell_Grek",
+    "eng_Latn",
+    "epo_Latn",
+    "est_Latn",
+    "eus_Latn",
+    "ewe_Latn",
+    "fao_Latn",
+    "pes_Arab",
+    "fij_Latn",
+    "fin_Latn",
+    "fon_Latn",
+    "fra_Latn",
+    "fur_Latn",
+    "fuv_Latn",
+    "gla_Latn",
+    "gle_Latn",
+    "glg_Latn",
+    "grn_Latn",
+    "guj_Gujr",
+    "hat_Latn",
+    "hau_Latn",
+    "heb_Hebr",
+    "hin_Deva",
+    "hne_Deva",
+    "hrv_Latn",
+    "hun_Latn",
+    "hye_Armn",
+    "ibo_Latn",
+    "ilo_Latn",
+    "ind_Latn",
+    "isl_Latn",
+    "ita_Latn",
+    "jav_Latn",
+    "jpn_Jpan",
+    "kab_Latn",
+    "kac_Latn",
+    "kam_Latn",
+    "kan_Knda",
+    "kas_Arab",
+    "kas_Deva",
+    "kat_Geor",
+    "knc_Arab",
+    "knc_Latn",
+    "kaz_Cyrl",
+    "kbp_Latn",
+    "kea_Latn",
+    "khm_Khmr",
+    "kik_Latn",
+    "kin_Latn",
+    "kir_Cyrl",
+    "kmb_Latn",
+    "kon_Latn",
+    "kor_Hang",
+    "kmr_Latn",
+    "lao_Laoo",
+    "lvs_Latn",
+    "lij_Latn",
+    "lim_Latn",
+    "lin_Latn",
+    "lit_Latn",
+    "lmo_Latn",
+    "ltg_Latn",
+    "ltz_Latn",
+    "lua_Latn",
+    "lug_Latn",
+    "luo_Latn",
+    "lus_Latn",
+    "mag_Deva",
+    "mai_Deva",
+    "mal_Mlym",
+    "mar_Deva",
+    "min_Latn",
+    "mkd_Cyrl",
+    "plt_Latn",
+    "mlt_Latn",
+    "mni_Beng",
+    "khk_Cyrl",
+    "mos_Latn",
+    "mri_Latn",
+    "zsm_Latn",
+    "mya_Mymr",
+    "nld_Latn",
+    "nno_Latn",
+    "nob_Latn",
+    "npi_Deva",
+    "nso_Latn",
+    "nus_Latn",
+    "nya_Latn",
+    "oci_Latn",
+    "gaz_Latn",
+    "ory_Orya",
+    "pag_Latn",
+    "pan_Guru",
+    "pap_Latn",
+    "pol_Latn",
+    "por_Latn",
+    "prs_Arab",
+    "pbt_Arab",
+    "quy_Latn",
+    "ron_Latn",
+    "run_Latn",
+    "rus_Cyrl",
+    "sag_Latn",
+    "san_Deva",
+    "sat_Beng",
+    "scn_Latn",
+    "shn_Mymr",
+    "sin_Sinh",
+    "slk_Latn",
+    "slv_Latn",
+    "smo_Latn",
+    "sna_Latn",
+    "snd_Arab",
+    "som_Latn",
+    "sot_Latn",
+    "spa_Latn",
+    "als_Latn",
+    "srd_Latn",
+    "srp_Cyrl",
+    "ssw_Latn",
+    "sun_Latn",
+    "swe_Latn",
+    "swh_Latn",
+    "szl_Latn",
+    "tam_Taml",
+    "tat_Cyrl",
+    "tel_Telu",
+    "tgk_Cyrl",
+    "tgl_Latn",
+    "tha_Thai",
+    "tir_Ethi",
+    "taq_Latn",
+    "taq_Tfng",
+    "tpi_Latn",
+    "tsn_Latn",
+    "tso_Latn",
+    "tuk_Latn",
+    "tum_Latn",
+    "tur_Latn",
+    "twi_Latn",
+    "tzm_Tfng",
+    "uig_Arab",
+    "ukr_Cyrl",
+    "umb_Latn",
+    "urd_Arab",
+    "uzn_Latn",
+    "vec_Latn",
+    "vie_Latn",
+    "war_Latn",
+    "wol_Latn",
+    "xho_Latn",
+    "ydd_Hebr",
+    "yor_Latn",
+    "yue_Hant",
+    "zho_Hans",
+    "zho_Hant",
+    "zul_Latn"
+]

requirements.txt

@@ -1,7 +1,3 @@
 git+https://github.com/huggingface/transformers
-
-gradio==3.8
-
-torch
-
-httpx==0.24.1
+gradio
+torch

ui.py

@@ -0,0 +1,12 @@
+title = "NLLB-200 Traslation Demo"
+description = """
+<p>
+<center>
+Translator using <a href='https://ai.facebook.com/research/no-language-left-behind/' target='_blank'>Facebook's NLLB</a> models.
+Codes Using <a href='https://github.com/facebookresearch/fairseq/tree/nllb' target=blank'>Facebook's fairseq NLLB</a>. 
+Demo is running on CPU.
+</center>
+</p>
+"""
+
+examples = [["我非常喜欢这个地方", "zho_Hans", "yue_Hant", 512]]