Upload indexador.py

indexador.py

@@ -0,0 +1,90 @@
+import os
+import fitz  # PyMuPDF
+import faiss
+import pickle
+import numpy as np
+from transformers import AutoTokenizer, AutoModel
+import torch
+import re
+
+# Modelo de embeddings
+modelo_id = "jinaai/jina-embeddings-v2-base-es"
+tokenizer = AutoTokenizer.from_pretrained(modelo_id)
+model = AutoModel.from_pretrained(modelo_id)
+
+def cargar_pdfs(ruta="."):
+    textos = []
+    for archivo in os.listdir(ruta):
+        if archivo.endswith(".pdf"):
+            ruta_pdf = os.path.join(ruta, archivo)
+            print(f"Procesando: {archivo}")
+            doc = fitz.open(ruta_pdf)
+            texto = ""
+            for pagina in doc:
+                texto += pagina.get_text()
+            texto = texto.replace("\n", " ").replace("  ", " ").strip()
+            doc.close()
+            if texto:
+                textos.append(texto)
+    return textos
+
+def chunk_por_articulo(texto):
+    patron = r"(art[íi]culo[s]?[\s\n]+\d+[A-Z]?(?:\s*\.\s*)?)"
+    secciones = re.split(patron, texto, flags=re.IGNORECASE)
+
+    chunks = []
+    for i in range(1, len(secciones), 2):
+        titulo = secciones[i].strip()
+        contenido = secciones[i + 1].strip()
+        chunk = f"{titulo} {contenido}"
+        chunks.append(chunk)
+    return chunks
+
+def generar_embedding(textos, tokenizer, model, batch_size=32):
+    all_embeddings = []
+
+    for i in range(0, len(textos), batch_size):
+        batch = textos[i:i + batch_size]
+        inputs = tokenizer(batch, return_tensors="pt", padding=True, truncation=True)
+        with torch.no_grad():
+            outputs = model(**inputs)
+            last_hidden = outputs.last_hidden_state
+            mask = inputs["attention_mask"].unsqueeze(-1).expand(last_hidden.size()).float()
+            summed = torch.sum(last_hidden * mask, 1)
+            counted = torch.clamp(mask.sum(1), min=1e-9)
+            mean_pooled = summed / counted
+            all_embeddings.extend(mean_pooled.cpu().numpy())
+
+    return np.array(all_embeddings)
+
+def crear_index_y_guardar(ruta=".", archivo_salida="index.pkl"):
+    print("Cargando PDFs...")
+    textos = cargar_pdfs(ruta)
+
+    print("Dividiendo en chunks por artículos...")
+    chunks = []
+    for texto in textos:
+        trozos = chunk_por_articulo(texto)
+        chunks.extend(trozos)
+
+    if not chunks:
+        raise ValueError("No se generaron chunks. Revisa los archivos PDF.")
+
+    print(f"Total de chunks generados: {len(chunks)}")
+
+    print("Generando embeddings...")
+    embeddings = generar_embedding(chunks, tokenizer, model)
+
+    print("Creando índice FAISS...")
+    index = faiss.IndexFlatL2(embeddings.shape[1])
+    index.add(embeddings)
+
+    print(f"Guardando índice en: {archivo_salida}")
+    with open(archivo_salida, "wb") as f:
+        pickle.dump((index, chunks), f)
+
+    print("Indexación completada.")
+    return index, chunks
+
+if __name__ == "__main__":
+    crear_index_y_guardar()