import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# Charger le modèle et le tokenizer depuis Hugging Face
model_name = "GueuleDange/UnloTeach"  # Remplacez par le nom de votre modèle
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Définir une fonction pour interagir avec le modèle
def generate_text(prompt, max_tokens=100, temperature=0.7, top_p=0.95):
    """
    Génère une réponse à partir du modèle.
    :param prompt: Le texte d'entrée.
    :param max_tokens: Le nombre maximum de tokens à générer.
    :param temperature: Le paramètre de température pour la génération.
    :param top_p: Le paramètre de top-p (nucleus sampling).
    :return: La réponse générée par le modèle.
    """
    # Encoder l'entrée
    inputs = tokenizer(prompt, return_tensors="pt")

    # Générer la réponse
    outputs = model.generate(
        inputs.input_ids,
        max_length=max_tokens,
        temperature=temperature,
        top_p=top_p,
        do_sample=True,
    )

    # Décoder la réponse
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

# Créer l'interface Gradio
with gr.Blocks() as demo:
    gr.Markdown("# Interface pour interagir avec votre modèle")
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(label="Entrez votre message", placeholder="Posez une question ou écrivez quelque chose...")
            max_tokens = gr.Slider(minimum=10, maximum=500, value=100, step=10, label="Nombre maximum de tokens")
            temperature = gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Température")
            top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)")
            submit_button = gr.Button("Générer")
        
        with gr.Column():
            output_text = gr.Textbox(label="Réponse du modèle", interactive=False)

    # Lier la fonction à l'interface
    submit_button.click(
        fn=generate_text,
        inputs=[input_text, max_tokens, temperature, top_p],
        outputs=output_text,
    )

# Lancer l'application
demo.launch()