import gradio as gr
from transformers import MarianMTModel, MarianTokenizer

# Charger le modèle et le tokenizer depuis Hugging Face
model_name = "MrFrijo/LiAPI"  # Nom du modèle sur Hugging Face

model = MarianMTModel.from_pretrained(model_name)
tokenizer = MarianTokenizer.from_pretrained(model_name)

def translate_text(text, src_lang, target_lang):
    # Préparer les entrées pour le modèle
    tokenized_text = tokenizer(text, return_tensors="pt")

    # Effectuer la traduction
    translated = model.generate(**tokenized_text)

    # Convertir et retourner le texte traduit
    translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)
    return translated_text

# Créer l'interface Gradio
with gr.Blocks() as interface:
    # Section d'en-tête
    gr.Markdown("<h2>Traduction automatique Lingala-Français</h2>")
    gr.Markdown("""
    <p>Ceci est la version d'essai et nous comptons sur vous pour améliorer les performances du modèle pour notre langue Lingala.</p>
    <p>Oyo e za version ya komeka pe to zo talela bino pôna ko kolisa bokasi ya modeli pôna nkota nà biso Lingala.</p>
    """)

    # Entrée de traduction
    with gr.Row():
        text_input = gr.Textbox(label="Entrez le texte à traduire", placeholder="Entrez le texte à traduire ici...")
    
    with gr.Row():
        source_lang = gr.Dropdown(choices=["fr", "li"], label="Langue Source")
        target_lang = gr.Dropdown(choices=["li", "fr"], label="Langue Cible")
    
    # Sortie de traduction
    translation_output = gr.Textbox(label="Traduction", placeholder="Le texte traduit s'affichera ici...")

    # Bouton pour soumettre
    translate_button = gr.Button("Traduire")
    translate_button.click(
        translate_text, 
        inputs=[text_input, source_lang, target_lang], 
        outputs=translation_output
    )

# Lancer l'interface
interface.launch()