from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import gradio as gr

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

# Fonction de génération de texte
def generate_text(input_text, max_length=50):
    input_ids = tokenizer.encode(input_text, return_tensors="pt")
    output_ids = model.generate(input_ids, max_length=max_length, num_return_sequences=1)
    output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    return output_text

# Définir l'interface Gradio
iface = gr.Interface(
    fn=generate_text, 
    inputs=[gr.Textbox(lines=2, placeholder="Entrez le texte en wolof ici..."), gr.Slider(10, 100, step=1, value=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()