from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Charger le tokenizer et le modèle
tokenizer = AutoTokenizer.from_pretrained("papasega/wolo_generation_t5")
model = AutoModelForSeq2SeqLM.from_pretrained("papasega/wolo_generation_t5")

def generate_text(input_text, max_length=50):
    # Tokeniser l'entrée
    input_ids = tokenizer.encode(input_text, return_tensors="pt")
    
    # Générer le texte
    output_ids = model.generate(input_ids, max_length=max_length, num_return_sequences=1)
    
    # Décoder le texte généré
    output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    
    return output_text


import gradio as gr

# Définir l'interface Gradio
iface = gr.Interface(
    fn=generate_text, 
    inputs=[gr.inputs.Textbox(lines=2, placeholder="Entrez le texte en wolof ici..."), gr.inputs.Slider(10, 100, step=1, default=50)], 
    outputs="text",
    title="Générateur de Texte en Wolof",
    description="Entrez un texte en wolof pour générer une continuation à l'aide du modèle T5 pré-entraîné."
)

# Lancer l'interface
iface.launch(debug=True)