import gradio as gr
from transformers import pipeline
import joblib
# importer la liste des langues
languages = joblib.load('languages.joblib')
# Fonction de traduction
def translate_text(text, model_choice, src_lang, tgt_lang):
    if model_choice == "NLLB-200 (facebook)":
        nllb_translator = pipeline("translation", model="facebook/nllb-200-distilled-600M", src_lang=src_lang, tgt_lang=tgt_lang)
        result = nllb_translator(text)[0]["translation_text"]
    elif model_choice == "Opus-MT-EN-FR (Helsinki-NLP)":
        en_fr_translator = pipeline("translation", model="Helsinki-NLP/opus-mt-en-fr")
        result = en_fr_translator(text)[0]["translation_text"]
    else:
        fr_en_translator = pipeline("translation", model="Helsinki-NLP/opus-mt-fr-en")
        result = fr_en_translator(text)[0]["translation_text"]
    return result

# Interface Gradio
demo = gr.Interface(fn=translate_text,
       inputs=[gr.Textbox(label="Texte à traduire", placeholder="Entrez du texte en français ici..."),
        gr.Radio(choices = ["NLLB-200 (facebook)", "Opus-MT-EN-FR (Helsinki-NLP)", "Opus-MT-FR-EN (Helsinki-NLP)"],label="Choisissez le modèle"),
        gr.Dropdown(choices = languages, label = 'Source Language'),
        gr.Dropdown(choices = languages, label = 'Target Language')],
    outputs=gr.Textbox(label="Texte traduit"),
    title="Traduction dans plusieurs langues",
    theme='shivi/calm_seafoam', 
    description = """Cette application permet de traduire un texte entre plusieurs langues
en utilisant différents modèles de traduction automatique.

Vous pouvez choisir parmi trois modèles populaires,
dont NLLB-200 de Facebook pour une large couverture linguistique,
ou les modèles Opus-MT spécialisés en anglais-français et français-anglais.

Sélectionnez la langue source et la langue cible,
saisissez votre texte, et obtenez rapidement la traduction correspondante."""
)

demo.launch()