import gradio as gr
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

# Charger le modèle Flan-T5-large et le tokenizer
model_name = "google/flan-t5-large"
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Fonction pour générer le programme alimentaire
def generer_programme(prompt):
    # Préparer le prompt pour le modèle
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
    # Générer le texte avec le modèle Flan-T5-large
    output = model.generate(input_ids, max_length=500, num_return_sequences=1)
    # Décoder le texte généré
    programme_alimentaire = tokenizer.decode(output[0], skip_special_tokens=True)
    return programme_alimentaire

# Créer une interface Gradio
with gr.Blocks() as demo:
    gr.Markdown("# Générateur de Programme Alimentaire avec Flan-T5-large")
    
    # Champ pour entrer le prompt de l'utilisateur
    prompt_input = gr.Textbox(label="Prompt", placeholder="Décrivez les besoins nutritionnels et préférences ici")
    
    # Champ de sortie pour afficher le programme généré
    result_output = gr.Textbox(label="Programme Alimentaire Généré")
    
    # Bouton pour déclencher la génération
    generer_button = gr.Button("Générer le Programme")

    # Associer le bouton à la fonction de génération
    generer_button.click(generer_programme, inputs=prompt_input, outputs=result_output)

# Lancer l'application Gradio
demo.launch()