from fastapi import FastAPI
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
import torch

app = FastAPI(title="Assistant IA Education Marocaine")

# Chargement du modèle
MODEL_NAME = "unsloth/mistral-7b-v0.3"
LORA_NAME = "dohael/mistral-7b-education-maroc"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16,
    device_map="auto"
)
model = PeftModel.from_pretrained(model, LORA_NAME)

class Question(BaseModel):
    question: str

@app.get("/")
def root():
    return {"message": "Assistant IA Education Marocaine 🇲🇦", "status": "running"}

@app.post("/ask")
def ask(q: Question):
    prompt = f"<s>[INST] {q.question} [/INST]"
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(
        **inputs,
        max_new_tokens=512,
        temperature=0.7,
        do_sample=True,
        repetition_penalty=1.1
    )
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    response = response.split("[/INST]")[-1].strip()
    return {"question": q.question, "reponse": response}

@app.get("/health")
def health():
    return {"status": "healthy"}