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
interface = gr.Interface(
    fn=translate_text, 
    inputs=[
        gr.Textbox(label="Texte à traduire"), 
        gr.Dropdown(choices=["fr", "li"], label="Langue Source"),
        gr.Dropdown(choices=["fr", "li"], label="Langue Cible")
    ], 
    outputs="text"
)

# Lancer l'interface
interface.launch()