api-semantic / api_semantica.py
Francogv's picture
Update api_semantica.py
4709f03 verified
from fastapi import FastAPI
from pydantic import BaseModel
from sentence_transformers import SentenceTransformer, util
import numpy as np
import json
import uvicorn
import os
os.environ["HF_HOME"] = "/tmp"
os.environ["TRANSFORMERS_CACHE"] = "/tmp"
os.environ["SENTENCE_TRANSFORMERS_HOME"] = "/tmp"
app = FastAPI()
# Cargar el model
modelo = SentenceTransformer("Francogv/Modelo-Embedding")
# Cargar la base de embeddings
with open("base_semantica.json", "r", encoding="utf-8") as f:
base = json.load(f)
# Pydantic schema para entrada
class PreguntaInput(BaseModel):
pregunta: str
top_k: int = 3
# Función de búsqueda
def buscar_semanticamente(pregunta, top_k=3):
emb_pregunta = modelo.encode(pregunta)
resultados = []
for item in base:
emb_item = np.array(item["embedding"], dtype=np.float32)
score = util.cos_sim(emb_pregunta, emb_item).item()
resultados.append((score, item))
resultados.sort(reverse=True, key=lambda x: x[0])
return resultados[:top_k]
# Endpoint principal
@app.post("/buscar")
async def buscar(input: PreguntaInput):
resultados = buscar_semanticamente(input.pregunta, input.top_k)
return {
"pregunta": input.pregunta,
"resultados": [
{
"score": round(score, 4),
"titulo": item["titulo"],
"url": item["url"],
"texto": item["texto"]
} for score, item in resultados
]
}
# Para correr en Hugging Face
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=7860)