File size: 2,558 Bytes
39d6011 | 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 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | import json
import torch
import numpy as np
from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer, AutoModelForCausalLM
# ================================
# CARGA DE MODELOS
# ================================
embedding_model = SentenceTransformer("MongoDB/mdbr-leaf-ir")
tokenizer = AutoTokenizer.from_pretrained("PleIAs/Pleias-RAG-350M")
llm_model = AutoModelForCausalLM.from_pretrained("PleIAs/Pleias-RAG-350M")
# ================================
# CARGA DE DOCUMENTOS (FIX AQUÍ)
# ================================
with open("documents.json", "r", encoding="utf-8") as f:
documents_dict = json.load(f)
# 🔥 IMPORTANTE: convertir dict → lista de textos
documents = list(documents_dict.values())
# ================================
# PRECOMPUTO DE EMBEDDINGS
# ================================
doc_embeddings = embedding_model.encode(documents, convert_to_tensor=True)
# ================================
# RECUPERAR DOCUMENTOS
# ================================
def recuperar_documentos(consulta, top_k=2, umbral=0.4):
query_embedding = embedding_model.encode(consulta, convert_to_tensor=True)
similitudes = torch.nn.functional.cosine_similarity(
query_embedding.unsqueeze(0),
doc_embeddings,
dim=1
)
indices_ordenados = torch.argsort(similitudes, descending=True)
resultados = []
for idx in indices_ordenados:
if similitudes[idx].item() >= umbral:
resultados.append(documents[idx])
if len(resultados) >= top_k:
break
return resultados
# ================================
# GENERAR RESPUESTA
# ================================
def generar_respuesta(consulta, documentos_recuperados):
contexto = " ".join(documentos_recuperados)
prompt = f"""Answer the question based only on the context provided
Context: {contexto}
Question: {consulta}
Answer:"""
inputs = tokenizer(prompt, return_tensors="pt")
outputs = llm_model.generate(
**inputs,
max_new_tokens=150
)
respuesta = tokenizer.decode(outputs[0], skip_special_tokens=True)
# eliminar el prompt del output
respuesta = respuesta.replace(prompt, "").strip()
return respuesta
# ================================
# FUNCIÓN FINAL
# ================================
def preguntar(consulta, top_k=2, umbral=0.4):
docs = recuperar_documentos(consulta, top_k, umbral)
return generar_respuesta(consulta, docs) |