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

MODEL_DIR = "./mistral-7b-brvm-finetuned"

# Fonction d’entraînement (appelle ton script de fine-tuning)
def train_model():
    os.system("python finetune.py")  # tu mets ton code d'entraînement dans finetune.py
    return "✅ Entraînement terminé ! Le modèle est sauvegardé dans " + MODEL_DIR

# Chargement du modèle (fine-tuné si dispo, sinon base)
def load_model():
    model_name = MODEL_DIR if os.path.exists(MODEL_DIR) else "mistralai/Mistral-7B-Instruct-v0.3"
    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        device_map="auto",
        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
        trust_remote_code=True,
    )
    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1)
    return pipe

# On charge le pipeline une fois au démarrage
pipe = load_model()

# Fonction de test du modèle
def chat(prompt):
    outputs = pipe(prompt, max_new_tokens=200, do_sample=True, temperature=0.7, top_p=0.9)
    return outputs[0]["generated_text"]

# Interface Gradio
with gr.Blocks() as demo:
    gr.Markdown("# 🐟 BRVM Finetuner (Mistral-7B)")

    with gr.Tab("🚀 Entraînement"):
        train_btn = gr.Button("Lancer l’entraînement")
        train_output = gr.Textbox(label="Logs")
        train_btn.click(fn=train_model, outputs=train_output)

    with gr.Tab("💬 Tester le modèle"):
        input_text = gr.Textbox(label="Votre question :", placeholder="Posez une question...")
        output_text = gr.Textbox(label="Réponse du modèle")
        submit_btn = gr.Button("Envoyer")
        submit_btn.click(fn=chat, inputs=input_text, outputs=output_text)

demo.launch()