Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,595 Bytes
2af6f6f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 |
import os
import fitz # PyMuPDF
import faiss
import pickle
import numpy as np
from transformers import AutoTokenizer, AutoModel
import torch
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_texto(texto, longitud=800):
return [texto[i:i+longitud] for i in range(0, len(texto), longitud)]
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.append(mean_pooled.numpy())
return np.vstack(all_embeddings)
def crear_index_y_guardar(ruta=".", modelo_id="jinaai/jina-embeddings-v2-base-es", archivo_salida="index.pkl"):
print("Cargando PDFs...")
textos = cargar_pdfs(ruta)
print("Dividiendo en chunks...")
chunks = []
for texto in textos:
chunks.extend(chunk_texto(texto))
if not chunks:
raise ValueError("No se generaron chunks. Revisa tus PDFs.")
print(f"Total de chunks generados: {len(chunks)}")
print("Generando embeddings...")
tokenizer = AutoTokenizer.from_pretrained(modelo_id)
model = AutoModel.from_pretrained(modelo_id)
embeddings = generar_embedding(chunks, tokenizer, model, batch_size=32)
print(f"Dimensión de embeddings: {embeddings.shape[1]}")
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()
|