import os
import torch
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, BitsAndBytesConfig
from threading import Thread

# --- CONFIGURATION ---
MODEL_ID = os.getenv("MODEL_ID", "Qwen/Qwen2.5-Coder-7B-Instruct")
MAX_NEW_TOKENS = int(os.getenv("MAX_NEW_TOKENS", 2048))
TEMPERATURE = float(os.getenv("TEMPERATURE", 0.15))

# --- CHARGEMENT DU MODÈLE ---
print("Chargement du tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

# Nouvelle méthode pour activer le 4-bit proprement
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

print("Chargement du modèle (compression 4-bit NF4)...")
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    device_map="auto",
    quantization_config=bnb_config, # On utilise la config ici au lieu de l'argument direct
    trust_remote_code=True
)

# --- FONCTION DE RÉPONSE ---
def respond(message, history):
    messages = [{"role": "system", "content": "Tu es CodeMind, un expert en code et en mathématiques."}]
    
    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})
            
    messages.append({"role": "user", "content": message})

    input_ids = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt"
    ).to(model.device)

    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)

    generate_kwargs = dict(
        input_ids=input_ids,
        streamer=streamer,
        max_new_tokens=MAX_NEW_TOKENS,
        temperature=TEMPERATURE,
        do_sample=True,
        top_p=0.9,
    )

    t = Thread(target=model.generate, kwargs=generate_kwargs)
    t.start()

    partial_message = ""
    for new_token in streamer:
        partial_message += new_token
        yield partial_message

# --- INTERFACE GRADIO ---
demo = gr.ChatInterface(
    respond,
    title="CodeMind AI",
    description="Assistant Qwen 7B corrigé pour GPU classique.",
)

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