import gradio as gr
from unsloth import FastLanguageModel
import torch

# =========================
# Charger le modèle BuccAI depuis Hugging Face
# =========================
print("⏳ Chargement du modèle BuccAI...")
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="jfand/BuccAI",  # ⚡ ton repo modèle
    max_seq_length=2048,
    load_in_4bit=True,              # car version quantifiée
)
print("✅ Modèle chargé avec succès !")

# =========================
# Fonction de génération
# =========================
def generate_response(prompt, max_tokens=400):
    try:
        inputs = tokenizer(prompt, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_tokens,
            temperature=0.7,
            top_p=0.9,
            do_sample=True,
            repetition_penalty=1.15,
        )
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        return response.strip()
    except Exception as e:
        return f"⚠️ Erreur: {str(e)}"

# =========================
# Interface Gradio
# =========================
with gr.Blocks() as demo:
    gr.Markdown("# 🦷 BuccAI - Assistant Dentaire (Makandal Technologies)")

    with gr.Row():
        with gr.Column(scale=3):
            user_input = gr.Textbox(
                label="💬 Posez votre question dentaire",
                placeholder="Ex: Quels sont les symptômes de la gingivite ?",
                lines=3
            )
            max_tokens = gr.Slider(100, 1000, value=400, step=50, label="Max tokens")
            submit = gr.Button("Générer la réponse")
        with gr.Column(scale=4):
            output = gr.Textbox(label="🤖 Réponse de BuccAI", lines=15)

    submit.click(fn=generate_response, inputs=[user_input, max_tokens], outputs=output)

demo.launch()