indexador.py

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()