Spaces:
Running
Running
| 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 | |