Spaces:
Sleeping
Sleeping
| 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 | |
| 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) | |