| from transformers import AutoTokenizer, AutoModelForSeq2SeqLM |
| import gradio as gr |
|
|
| model_name = "facebook/nllb-200-distilled-600M" |
| tokenizer = AutoTokenizer.from_pretrained(model_name) |
| model = AutoModelForSeq2SeqLM.from_pretrained(model_name) |
|
|
| lang_readable = { |
| "aka_Latn": "Akan", "bam_Latn": "Bambara", "dyu_Latn": "Dioula", |
| "ewe_Latn": "Éwé", "fon_Latn": "Fon", "fuv_Latn": "Peul (Nigéria)", |
| "hau_Latn": "Haoussa", "ibo_Latn": "Igbo", "kab_Latn": "Kabyle", |
| "knc_Arab": "Kanouri (arabe)", "knc_Latn": "Kanouri (latin)", |
| "mos_Latn": "Mooré", "som_Latn": "Somali", "twi_Latn": "Twi", |
| "wol_Latn": "Wolof", "yor_Latn": "Yoruba", "kbp_Latn": "Kabiye", |
| "kon_Latn": "Kikongo", "lin_Latn": "Lingala", "lua_Latn": "Luba-Lulua", |
| "lug_Latn": "Luganda", "luo_Latn": "Luo", "run_Latn": "Kirundi", |
| "sag_Latn": "Sango", "amh_Ethi": "Amharique", "gaz_Latn": "Oromo", |
| "kam_Latn": "Kamba", "kik_Latn": "Kikuyu", "kin_Latn": "Kinyarwanda", |
| "nus_Latn": "Nuer", "swh_Latn": "Swahili", "tir_Ethi": "Tigrigna", |
| "afr_Latn": "Afrikaans", "bem_Latn": "Bemba", "cjk_Latn": "Chokwe", |
| "dik_Latn": "Dinka du Sud-Ouest", "kmb_Latn": "Kimbundu", |
| "nso_Latn": "Sepedi", "nya_Latn": "Chichewa", "sna_Latn": "Shona", |
| "sot_Latn": "Sesotho", "ssw_Latn": "Swati", "tsn_Latn": "Tswana", |
| "tso_Latn": "Tsonga", "tum_Latn": "Tumbuka", "umb_Latn": "Umbundu", |
| "xho_Latn": "Xhosa", "zul_Latn": "Zoulou", "aeb_Arab": "Arabe tunisien", |
| "ary_Arab": "Arabe marocain", "arz_Arab": "Arabe égyptien", |
| "arb_Arab": "Arabe standard moderne", "tzm_Tfng": "Tamazight (Tifinagh)", |
| "taq_Latn": "Tamasheq (latin)", "taq_Tfng": "Tamasheq (Tifinagh)", |
| "kea_Latn": "Créole capverdien", "plt_Latn": "Malgache (Plateau)", |
| "fra_Latn": "Français", "eng_Latn": "Anglais", |
| "por_Latn": "Portugais", "spa_Latn": "Espagnol", "ara_Arab": "Arabe" |
| } |
|
|
| lang_names = list(lang_readable.values()) |
|
|
| def translate(text, src_lang_name, tgt_lang_name): |
| src_lang = [k for k, v in lang_readable.items() if v == src_lang_name][0] |
| tgt_lang = [k for k, v in lang_readable.items() if v == tgt_lang_name][0] |
|
|
| |
| tokenizer.src_lang = src_lang |
|
|
| inputs = tokenizer(text, return_tensors="pt") |
|
|
| |
| tgt_lang_id = tokenizer.convert_tokens_to_ids(tgt_lang) |
|
|
| translated_tokens = model.generate( |
| **inputs, |
| forced_bos_token_id=tgt_lang_id |
| ) |
| return tokenizer.decode(translated_tokens[0], skip_special_tokens=True) |
|
|
| demo = gr.Interface( |
| fn=translate, |
| inputs=[ |
| gr.Textbox(label="Texte à traduire"), |
| gr.Dropdown(choices=lang_names, label="Langue source", filterable=True), |
| gr.Dropdown(choices=lang_names, label="Langue cible", filterable=True) |
| ], |
| outputs="text", |
| title="NLLB Translator", |
| description="Traduction multilingue avec NLLB" |
| ) |
| demo.launch() |