Create app.py

app.py

@@ -0,0 +1,136 @@
+from sentence_transformers import SentenceTransformer
+import chromadb
+import pandas as pd
+import gradio as gr
+from datetime import datetime
+
+# ===== Configuración =====
+CHROMA_DIR = "chroma_db"
+MODEL_NAME = "mrm8488/multilingual-e5-large-ft-sts-spanish-matryoshka-768-16-5e"
+COLLECTION_NAMES = ["consumidor_beta"]
+
+# ===== 1. Conectar a la base de datos y cargar colecciones =====
+print("Conectando a ChromaDB...")
+client = chromadb.PersistentClient(path=CHROMA_DIR)
+collections = [client.get_collection(name) for name in COLLECTION_NAMES]
+
+# ===== 2. Cargar el modelo =====
+print("Cargando modelo...")
+model = SentenceTransformer(MODEL_NAME, trust_remote_code=True)
+print("✓ Sistema listo")
+
+# ===== 3. Función de búsqueda =====
+def semantic_search(query: str, n_results: int = 5):
+    print(f"Buscando: {query}")
+    
+    if not query.strip():
+        return pd.DataFrame(), ""
+    
+    query_embedding = model.encode(query).tolist()
+
+    all_results = []
+    for collection in collections:
+        results = collection.query(
+            query_embeddings=[query_embedding],
+            n_results=n_results,
+            include=["documents", "metadatas", "distances"]
+        )
+        cosine_similarities = [1 - dist for dist in results['distances'][0]]
+        
+        for i in range(len(results['ids'][0])):
+            result_dict = {
+                '✓': False,  # Columna de selección al inicio
+                'ID': results['ids'][0][i],
+                'Similitud': round(cosine_similarities[i], 4),
+                'Texto': results['documents'][0][i],
+                'Colección': collection.name
+            }
+            
+            metadata = results['metadatas'][0][i]
+            if metadata:
+                for key, value in metadata.items():
+                    # Si hay URL, convertir a link
+                    if key.lower() in ['url', 'link', 'enlace'] and value:
+                        result_dict[key] = f'<a href="{value}" target="_blank">🔗 Abrir</a>'
+                    else:
+                        result_dict[key] = value
+            
+            all_results.append(result_dict)
+
+    df = pd.DataFrame(all_results).sort_values('Similitud', ascending=False).head(n_results)
+    print(f"Resultados: {len(df)}")
+    
+    return df, ""
+
+# ===== 4. Función para exportar seleccionados a Excel =====
+def export_to_excel(df_with_selection):
+    if df_with_selection is None or len(df_with_selection) == 0:
+        gr.Warning("No hay datos para exportar")
+        return None
+    
+    # Filtrar solo los marcados como True en la columna '✓'
+    df_selected = df_with_selection[df_with_selection['✓'] == True].copy()
+    
+    if len(df_selected) == 0:
+        gr.Warning("No has seleccionado ninguna decisión")
+        return None
+    
+    # Quitar la columna de checkbox del export
+    df_export = df_selected.drop(columns=['✓'])
+    
+    # Limpiar HTML de los links para Excel
+    for col in df_export.columns:
+        if df_export[col].dtype == 'object':
+            df_export[col] = df_export[col].apply(
+                lambda x: x.replace('<a href="', '').replace('" target="_blank">🔗 Abrir</a>', '') 
+                if isinstance(x, str) and '<a href=' in x else x
+            )
+    
+    # Guardar a Excel
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    filename = f"jurisprudencia_seleccionada_{timestamp}.xlsx"
+    df_export.to_excel(filename, index=False, engine='openpyxl')
+    
+    print(f"✓ Exportado: {filename} ({len(df_export)} registros)")
+    gr.Info(f"✓ Exportadas {len(df_export)} decisiones")
+    return filename
+
+# ===== 5. Interfaz Gradio =====
+with gr.Blocks(title="Buscador Jurisprudencia") as demo:
+    gr.Markdown("# 🔍 Buscador de Jurisprudencia")
+    gr.Markdown("Busca, marca las decisiones útiles en la columna ✓ y exporta a Excel")
+    
+    # Query box
+    query_box = gr.Textbox(
+        show_label=False, 
+        placeholder="Escribe tu consulta legal y presiona Enter...",
+        lines=1
+    )
+    
+    # Tabla de resultados (EDITABLE para marcar checkboxes)
+    results_table = gr.Dataframe(
+        label="Resultados",
+        wrap=True,
+        interactive=True,  # Permite editar la columna de checkbox
+        datatype=["bool", "str", "number", "markdown", "str"]  # Tipos de columnas
+    )
+    
+    # Botón de exportar
+    with gr.Row():
+        export_button = gr.Button("📥 Exportar seleccionadas a Excel", variant="primary", scale=2)
+        file_output = gr.File(label="Descargar", scale=1)
+    
+    # Eventos
+    query_box.submit(
+        semantic_search, 
+        [query_box], 
+        [results_table, query_box]
+    )
+    
+    export_button.click(
+        export_to_excel,
+        [results_table],
+        [file_output]
+    )
+
+demo.launch()