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

# 🔥 Forcer la mise à jour de transformers AVANT de l'importer
os.system("pip install --no-cache-dir -U transformers")

# 🔥 Installation/mise à jour des dépendances uniquement si nécessaire
print("🚀 Vérification et mise à jour des dépendances...")
os.system("pip install --no-cache-dir -U peft accelerate torch scipy")

# 🔥 Correction de `libstdc++6` pour éviter les erreurs
os.system("apt-get update && apt-get install -y --reinstall libstdc++6")

print("✅ Dépendances corrigées et mises à jour !")

# 📌 Définition des modèles
BASE_MODEL = "google/gemma-2-2b-it"
LORA_MODEL = "Serveurperso/gemma-2-2b-it-LoRA"

print("🚀 Chargement du modèle Gemma 2B avec LoRA Mémé Ginette...")

# 📌 Forcer le CPU
device = "cpu"

# 📌 Chargement du modèle principal
try:
    model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        device_map=device,
        torch_dtype=torch.float32,
        trust_remote_code=True,  # ✅ Permet à Hugging Face d'utiliser les nouvelles archis
        token=os.getenv("HF_TOKEN")
    )

    # 📌 Application du LoRA
    model = PeftModel.from_pretrained(
        model,
        LORA_MODEL,
        device_map=device,
        torch_dtype=torch.float32,
        token=os.getenv("HF_TOKEN")
    )

    tokenizer = AutoTokenizer.from_pretrained(LORA_MODEL, token=os.getenv("HF_TOKEN"))

    print("✅ Modèle chargé avec succès !")

except Exception as e:
    print(f"❌ Erreur lors du chargement du modèle: {e}")
    exit(1)

# 📌 Interface Gradio
def chat(message):
    inputs = tokenizer(message, return_tensors="pt").to(device)
    outputs = model.generate(**inputs, max_length=128)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

iface = gr.Interface(fn=chat, inputs="text", outputs="text", title="Mémé Ginette Chatbot")

print("🚀 Interface Gradio lancée sur port 7860")
iface.launch(share=True)