from fastapi import FastAPI
from sentence_transformers import SentenceTransformer
from gliner import GLiNER
from pydantic import BaseModel
from typing import List
import uvicorn

print("Démarrage de l'application...")

app = FastAPI()

class TextRequest(BaseModel):
    text: str

class TextRequestAvecEtiquettes(BaseModel):
    text: str
    etiquettes: List[str]

# Chargement de SBERT
print("Chargement de SBERT...")
modele_sbert = SentenceTransformer('paraphrase-multilingual-mpnet-base-v2')
#modele_sbert = SentenceTransformer('alex246879/sbert-professionnel-2026')
print("SBERT chargé avec succès")

print("Chargement de GLiNER...")
modele_gliner = GLiNER.from_pretrained("urchade/gliner_multi_pii-v1")
print("GLiNER chargé avec succès")

@app.post("/vecteur")
def creat_vector(data: TextRequest):
    print(f"Requête /vecteur reçue : {data.text}")
    vecteur = modele_sbert.encode(data.text).tolist()
    print(f"Taille vecteur envoyé : {len(vecteur)}")
    return {"vecteur": vecteur}

@app.post("/vecteurs")
def creat_vectors(data: TextRequestAvecEtiquettes):
    print(f"Requête /vecteurs reçue : {data.text}")
    print(f"Étiquettes demandées : {data.etiquettes}")

    # Extraction des entités par étiquette via GLiNER
    print("Extraction des entités via GLiNER...")
    entites = modele_gliner.predict_entities(data.text, data.etiquettes)
    print(f"Entités extraites : {entites}")

    # Regroupement des textes extraits par étiquette
    entites_par_etiquette: dict[str, List[str]] = {}
    for entite in entites:
        label = entite["label"]
        texte = entite["text"]
        entites_par_etiquette.setdefault(label, []).append(texte)

    print(f"Regroupement par étiquette : {entites_par_etiquette}")

    # Génération d'un vecteur SBERT par étiquette
    resultat = {}
    for label, textes in entites_par_etiquette.items():
        print(f"Encodage SBERT pour '{label}' : {textes}")
        vecteur = modele_sbert.encode(" ".join(textes)).tolist()
        resultat[label] = vecteur
        print(f"Étiquette '{label}' : {len(textes)} entité(s), taille vecteur : {len(vecteur)}")

    print("Résultat prêt, envoi de la réponse")
    return resultat

if __name__ == "__main__":
    # Hugging Face utilise le port 7860 par défaut
    uvicorn.run(app, host="0.0.0.0", port=7860)