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

# 1. Chargement du modèle et du tokenizer
model_name = "Conlanger-LLM-CLEM/Gheya-dialogue-v1"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

def predict(message, history):
    # Encodage du message
    input_ids = tokenizer(message, return_tensors="pt").input_ids
    
    # Génération
    outputs = model.generate(
        input_ids,
        do_sample=True,
        temperature=0.7,
        max_new_tokens=200,
        top_k=50,
        top_p=0.95,
        early_stopping=True,
        num_beams=1
    )
    
    # Décodage
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# 2. Construction de l'interface avec gr.Blocks
# On applique le thème ici, au niveau global
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# Gheya Dialogue Bot")
    gr.Markdown("Discutez avec Gheya-dialogue-v1, le Premier modèle Gheya de conversation abstraite.")
    
    gr.ChatInterface(
        fn=predict,
        examples=["Bonjour, comment vas-tu ?", "Qui es-tu ?"],
    )

if __name__ == "__main__":
    demo.launch()