app.py

import os
from fastapi import FastAPI
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

app = FastAPI()

# Définir un dossier cache accessible
os.environ["TRANSFORMERS_CACHE"] = "/tmp"

# Charger le modèle et le tokenizer depuis Hugging Face avec cache local
MODEL_NAME = "fatmata/psybot"
local_dir = "/tmp/model"  # Changer le dossier vers /tmp/model
os.makedirs(local_dir, exist_ok=True)  # Crée le dossier si nécessaire

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, cache_dir=local_dir)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, cache_dir=local_dir, torch_dtype=torch.float32)

@app.get("/")
def home():
    return {"message": "Bienvenue sur l'API PsyBot !"}

@app.post("/generate")
def generate_text(prompt: str):
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(**inputs, max_length=100)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return {"response": response}