app.py

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
import torch

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

# Autoriser toutes les origines (CORS)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

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,
    dtype=torch.float16,
    device_map="cpu"
)
model = PeftModel.from_pretrained(
    model, 
    LORA_NAME,
    is_trainable=False
)
model.eval()

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")
    with torch.no_grad():
        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"}