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

# Remplacez par le modèle de base et l'adaptateur LoRA que vous utilisez
BASE_MODEL = "bigcode/starcoder2-3b"
ADAPTER_REPO = "simnJS/autotrain-fxp6j-p5s8i"

# 1. Charger le tokenizer
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)

# 2. Charger le modèle de base en FP16 sur le GPU (si vous avez la VRAM nécessaire)
base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    torch_dtype=torch.float16,
    device_map="auto"   # Permet de placer le modèle sur le GPU automatiquement
)

# 3. Charger l'adaptateur LoRA
model = PeftModel.from_pretrained(
    base_model,
    ADAPTER_REPO,
    torch_dtype=torch.float16
)

# 4. Fonction pour générer une réponse
def generate_answer(user_message, history):
    """
    user_message: le dernier message de l'utilisateur
    history: liste de tuples (message_utilisateur, réponse_modèle)
    """
    # Construire le prompt en tenant compte de l'historique si besoin
    # Ici, on fait simple et on utilise juste le dernier message
    prompt = user_message

    # Encoder le prompt
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    # Générer la réponse
    outputs = model.generate(
        **inputs,
        max_new_tokens=100,
        temperature=0.7,
        do_sample=True,
        top_p=0.9
    )
    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
    # Ajouter dans l'historique
    history.append((user_message, answer))
    return history, history

# 5. Interface Gradio de type chatbot
with gr.Blocks() as demo:
    gr.Markdown("# Chat avec mon modèle LoRA Verse")
    chatbot = gr.Chatbot()
    msg = gr.Textbox(label="Tapez votre message ici...")
    state = gr.State([])  # pour stocker l'historique des messages

    def submit_message(user_message, history):
        return generate_answer(user_message, history)

    msg.submit(submit_message, inputs=[msg, state], outputs=[chatbot, state])
    # ou un bouton si vous préférez
    # send_btn = gr.Button("Envoyer")
    # send_btn.click(fn=submit_message, inputs=[msg, state], outputs=[chatbot, state])

demo.launch()