from fastapi import FastAPI
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModel
import torch
import torch.nn.functional as F
import os

# 📁 Rediriger les caches HF vers un dossier autorisé
CACHE_DIR = "/data"
os.environ["HF_HOME"] = CACHE_DIR
os.environ["TRANSFORMERS_CACHE"] = CACHE_DIR
os.environ["HF_MODULES_CACHE"] = CACHE_DIR  # ✅ <--- ajoute cette ligne
os.environ["HF_HUB_CACHE"] = CACHE_DIR


app = FastAPI()

# ✅ Modèle avec custom code
MODEL_NAME = "nomic-ai/nomic-embed-text-v1"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, cache_dir=CACHE_DIR, trust_remote_code=True)
model = AutoModel.from_pretrained(MODEL_NAME, cache_dir=CACHE_DIR, trust_remote_code=True)

class EmbedInput(BaseModel):
    text: str

@app.post("/embed")
async def embed_text(payload: EmbedInput):
    inputs = tokenizer(payload.text, return_tensors="pt", padding=True, truncation=True)
    with torch.no_grad():
        outputs = model(**inputs)
        embeddings = outputs.last_hidden_state[:, 0]
        normalized = F.normalize(embeddings, p=2, dim=1)
    return {"embedding": normalized[0].tolist()}