Spaces:
Sleeping
Sleeping
File size: 1,581 Bytes
fb8a717 3f74446 fb8a717 4709f03 5ed1078 fb8a717 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 |
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)
|