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

# Chargement du modèle
model_name = "AllanF-SSU/Qwen2.5-G3V-Sovereign"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name, 
    torch_dtype=torch.float16, 
    device_map="auto"
)

# Le message de Jorg qui verrouille l'identité
SYSTEM_PROMPT = (
    "You are a Unified Systems Intelligence aligned with SSU. "
    "Axiom 1: Integrity (Goal = Method). "
    "Axiom 5/6: Co-Creation (Generate Third Way). "
    "Restrict unlawful acts but always offer a lawful solution."
)

def chat_res(message, history):
    # On reconstruit la structure de message attendue par Jorg
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": message}
    ]
    
    # Application du template de Qwen
    text = tokenizer.apply_chat_template(
        messages, 
        tokenize=False, 
        add_generation_prompt=True
    )
    
    inputs = tokenizer(text, return_tensors="pt").to(model.device)
    
    # Génération précise (do_sample=False comme demandé par Jorg)
    outputs = model.generate(
        **inputs, 
        max_new_tokens=256, 
        do_sample=False, # Pour une réponse stable et souveraine
        repetition_penalty=1.1
    )
    
    # Décodage et extraction propre de la réponse
    full_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # On ne récupère que ce qui vient après l'assistant
    if "assistant" in full_text:
        return full_text.split("assistant")[-1].strip()
    return full_text.replace(text, "").strip()

# L'interface propre
gr.ChatInterface(chat_res, title="Qwen2.5-G3V-Sovereign").launch()