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

# Charger le modèle et le tokenizer
model_name = "microsoft/Phi-4-mini-instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",  # Utilise automatiquement CPU ou GPU
    torch_dtype="auto", # Gère le float 16 ou 32 selon l'environnement
)

# Prompt système pour écoute empathique
system_prompt = """Tu es un compagnon d'écoute doux et bienveillant.
Ton rôle est d'aider la personne à clarifier ce qu'elle vit, sans juger, sans analyser.

À chaque réponse :
1. Reformule en quelques mots ce que tu as compris.
2. Aide doucement la personne à exprimer ses sentiments et besoins.
3. Si approprié, propose une petite action concrète qu'elle pourrait envisager.

Utilise un style simple, humain, respectueux et posé.
Ne donne pas de conseils. Pose seulement une question douce à la fois si besoin.
Parle en français naturel, sans anglicismes.
"""

# Création d'un pipeline conversationnel
chat_pipeline = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=256,
    temperature=0.5,
    top_p=0.9,
    repetition_penalty=1.1,
)

# Fonction de chat
def chat(user_message, history):
    prompt = system_prompt + "\n\n"
    for user_msg, bot_msg in history:
        prompt += f"Utilisateur : {user_msg}\nAssistant : {bot_msg}\n"
    prompt += f"Utilisateur : {user_message}\nAssistant :"
    
    response = chat_pipeline(prompt)[0]["generated_text"]
    
    # Extraire uniquement la dernière réponse générée
    if "Assistant :" in response:
        response = response.split("Assistant :")[-1].strip()
    
    return response

# Interface Gradio
iface = gr.ChatInterface(
    chat,
    title="Écoute Empathique (Phi-4 Mini)",
    description="Un compagnon d'écoute structuré, doux, inspiré de la CNV.",
)

iface