Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 import pickle
 import os
+import json
 from transformers import AutoTokenizer, AutoModel, pipeline
 import torch
 import faiss
@@ -17,6 +18,19 @@ if os.path.exists("index.pkl"):
 else:
     raise FileNotFoundError("No se encontró el archivo 'index.pkl'.")
 
+# Cargar diccionario de sinónimos
+with open("sinonimos.json", "r", encoding="utf-8") as f:
+    diccionario_sinonimos = json.load(f)
+
+# Función para expandir palabras clave con sinónimos
+def expandir_con_sinonimos(palabras, diccionario):
+    resultado = set(palabras)
+    for palabra in palabras:
+        for clave, sinonimos in diccionario.items():
+            if palabra == clave or palabra in sinonimos:
+                resultado.update([clave] + sinonimos)
+    return list(resultado)
+
 # Modelo de embeddings
 model_id = "jinaai/jina-embeddings-v2-base-es"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
@@ -51,11 +65,10 @@ def responder(pregunta):
     result_chunks = [chunks[i] for i in indices[0]]
 
     palabras_clave = pregunta.lower().split()
-    #filtrados = [c for c in result_chunks if any(p in c.lower() for p in palabras_clave)]
-    #contexto_final = "\n\n".join(filtrados[:3]) if filtrados else "\n\n".join(result_chunks[:3])
-    contexto_final = "\n\n".join(result_chunks[:10])
-    print("Primeros chunks:")
-    print("\n\n".join(result_chunks[:3]))
+    palabras_expandidas = expandir_con_sinonimos(palabras_clave, diccionario_sinonimos)
+    
+    filtrados = [c for c in result_chunks if any(p in c.lower() for p in palabras_expandidas)]
+    contexto_final = "\n\n".join(filtrados[:3]) if filtrados else "\n\n".join(result_chunks[:3])
 
     prompt = f"""
     Eres un abogado colombiano especializado en el Código de Tránsito, Código de Policía y Código Penal. 
@@ -104,4 +117,4 @@ demo = gr.Interface(
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()