from transformers import T5ForConditionalGeneration, T5Tokenizer
import gradio as gr


model = T5ForConditionalGeneration.from_pretrained("Baldezo313/t5-finetuned-squad-mamadou")
tokenizer = T5Tokenizer.from_pretrained("Baldezo313/t5-finetuned-squad-mamadou")

def generate_answer(context, question):
    inputs = tokenizer(context, question, return_tensors="pt", padding="max_length", truncation=True, max_length=512)
    outputs = model.generate(**inputs, max_length=64, num_beams=4, early_stopping=True)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

demo = gr.Interface(
    fn=generate_answer,
    inputs=[
        gr.Textbox(label="Contexte", lines=6, placeholder="Entrez ici le contexte..."),
        gr.Textbox(label="Question", placeholder="Entrez votre question...")
    ],
    outputs=gr.Textbox(label="Réponse générée"),
    title="🧠 Générateur de Réponses T5",
    description="Donnez un contexte et une question — le modèle T5 fine-tuné génère une réponse cohérente."
)

demo.launch()