Spaces:
Running
Running
File size: 3,233 Bytes
e581cb0 | 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 | import json
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from sentence_transformers import SentenceTransformer, util
# --- 1. CARGA DE MODELOS ---
# Usamos un modelo para buscar informacion en los documentos
print("Cargando el modelo de busqueda...")
embed_model = SentenceTransformer("MongoDB/mdbr-leaf-ir")
# Usamos un modelo para que nos ayude a redactar la respuesta
print("Cargando el modelo de lenguaje...")
model_id = "microsoft/Phi-2"
tokenizer = AutoTokenizer.from_pretrained(model_id)
llm_model = AutoModelForCausalLM.from_pretrained(model_id)
# --- 2. BASE DE CONOCIMIENTO ---
# Leemos los documentos que tenemos guardados en el archivo JSON
with open("documents.json", "r", encoding="utf-8") as f:
docs_dict = json.load(f)
documents = list(docs_dict.values())
# Preparamos los datos para que el ordenador pueda buscar en ellos rapidamente
print("Preparando los documentos para la busqueda...")
docs_embeddings = embed_model.encode(documents, convert_to_tensor=True)
# --- 3. FUNCIONES DEL SISTEMA ---
def recuperar_documentos(consulta, top_k=2, umbral=0.4):
""" Busca en nuestra base de datos los textos que mas se parecen a la pregunta """
# Pasamos la pregunta a un formato que el modelo entienda
query_embedding = embed_model.encode(consulta, convert_to_tensor=True)
# Comparamos la pregunta con todos los documentos para ver cuales coinciden mejor
cos_scores = util.cos_sim(query_embedding, docs_embeddings)[0]
# Nos quedamos con los mejores resultados segun lo que hayamos configurado
top_results = torch.topk(cos_scores, k=min(top_k, len(documents)))
final_docs = []
for score, idx in zip(top_results[0], top_results[1]):
# Solo guardamos el documento si se parece lo suficiente a la pregunta
if score >= umbral:
final_docs.append(documents[idx])
return final_docs
def generar_respuesta(consulta, documentos_recuperados):
""" Redacta una respuesta usando la pregunta y la informacion encontrada """
# Si no hemos encontrado nada, usamos un texto por defecto
contexto = " ".join(documentos_recuperados) if documentos_recuperados else "No se ha encontrado informacion relevante."
# Preparamos las instrucciones para el modelo
prompt = f"Context: {contexto}\nQuestion: {consulta}\nAnswer:"
# Convertimos el texto a un formato que el modelo pueda procesar
inputs = tokenizer(prompt, return_tensors="pt")
# El modelo genera el texto de la respuesta
with torch.no_grad():
outputs = llm_model.generate(**inputs, max_new_tokens=100)
# Limpiamos el texto generado para quedarnos solo con lo importante
resultado_total = tokenizer.decode(outputs[0], skip_special_tokens=True)
respuesta_limpia = resultado_total.split("Answer:")[-1].strip()
return respuesta_limpia
def preguntar(consulta, top_k=2, umbral=0.4):
""" Funcion principal que busca la informacion y redacta la respuesta """
docs = recuperar_documentos(consulta, top_k, umbral)
respuesta = generar_respuesta(consulta, docs)
return respuesta, docs
|