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

MODEL = "prithivMLmods/Llama-SmolTalk-3.2-1B-Instruct"

# Charger le tokenizer
tokenizer = AutoTokenizer.from_pretrained(MODEL)

# Charger le modèle en 8 bits
model = AutoModelForCausalLM.from_pretrained(
    MODEL,
    device_map="auto",
    load_in_8bit=True
)

# Pipeline
chatbot = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    device_map="auto"
)

# Prompt système
system_prompt = (
    "Tu es Aria, une IA gentille, claire et polie. "
    "Réponds toujours en phrases complètes. "
    "Ne te lances pas dans un jeu de rôle, ne répète pas les messages précédents, "
    "et donne uniquement ta réponse."
)

def chat(message, history):
    history = history or []
    context = "\n".join([f"{user}\n{bot}" for user, bot in history[-3:]])
    prompt = f"{system_prompt}\n{context}\n{message}\nRéponse:"

    resp = chatbot(
        prompt,
        max_new_tokens=250,   # plus long pour éviter les coupures
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
        repetition_penalty=1.1
    )[0]["generated_text"]

    # Couper dès qu'il repart sur un nouveau tour
    reply = resp.split("Réponse:")[-1].strip()
    reply = reply.split("Utilisateur:")[0].strip()

    history.append((message, reply))
    return history, history

with gr.Blocks() as demo:
    chatbot_ui = gr.Chatbot()
    state = gr.State([])  # sauvegarde de l'historique
    msg = gr.Textbox(placeholder="Écris un message...")

    msg.submit(chat, [msg, state], [chatbot_ui, state])
    msg.submit(lambda: "", None, msg)  # reset input après envoi

demo.launch()