pdf_processor.py

import fitz
from tqdm.auto import tqdm
from spacy.lang.en import English

def text_formatter(text: str) -> str:
    """
    Formatea el texto extraído de un PDF eliminando saltos de línea
    y espacios sobrantes.
    """
    return text.replace("\n", " ").strip()

def read_pdf(pdf_bytes: bytes) -> list[dict]:
    """
    Lee un archivo PDF a partir de sus bytes y extrae información de cada página:
    - Número de página
    - Conteo de caracteres
    - Conteo de palabras
    - Conteo aproximado de oraciones (separadas por '. ')
    - Conteo de tokens aproximado (asumiendo 1 token ~ 4 caracteres)
    - Texto completo de la página

    Parámetros:
    -----------
    pdf_bytes : bytes
        Contenido binario de un PDF.

    Retorna:
    --------
    list[dict]
        Lista de diccionarios con la información de cada página.
    """
    doc = fitz.open(stream=pdf_bytes, filetype="pdf")
    pages_and_text = []
    for page_number, page in tqdm(enumerate(doc), desc="Leyendo PDF"):
        text = page.get_text()
        text = text_formatter(text=text)
        pages_and_text.append(
            {
                "page_number": page_number + 1,
                "page_char_count": len(text),
                "page_word_count": len(text.split(" ")),
                "page_sentence_count_raw": len(text.split(". ")),
                "page_token_count": len(text) / 4,  # Estimación aproximada
                "text": text,
            }
        )
    return pages_and_text

def process_chunks(pages_and_text: list[dict]) -> list[dict]:
    """
    Procesa cada página para dividir el texto en oraciones y luego agrupar
    esas oraciones en 'chunks' de 10 oraciones cada uno.

    Parámetros:
    -----------
    pages_and_text : list[dict]
        Lista de diccionarios que contienen el texto y metadatos de cada página.

    Retorna:
    --------
    list[dict]
        Lista de diccionarios donde cada diccionario representa un chunk de oraciones.
    """
    nlp = English()
    nlp.add_pipe("sentencizer")

    # Dividir el texto en oraciones
    for item in tqdm(pages_and_text, desc="Dividiendo en oraciones"):
        item["sentences"] = [str(sent) for sent in nlp(item["text"]).sents]
        item["sentence_chunks"] = split_list(item["sentences"], 10)
        item["num_of_chunks"] = len(item["sentence_chunks"])

    # Crear una lista de todos los chunks
    pages_and_chunks = []
    for item in tqdm(pages_and_text, desc="Creando 'chunks'"):
        for sentence_chunk in item["sentence_chunks"]:
            chunk_dict = {
                "page_number": item["page_number"],
                "sentence_chunk": " ".join(sentence_chunk).strip(),
                "chunk_word_count": sum(len(sentence.split()) for sentence in sentence_chunk),
                "chunk_token_count": sum(len(sentence) for sentence in sentence_chunk) / 4,
            }
            pages_and_chunks.append(chunk_dict)

    return pages_and_chunks

def split_list(input_list: list[str], slice_size: int) -> list[list[str]]:
    """
    Divide una lista en sublistas de tamaño determinado.

    Parámetros:
    -----------
    input_list : list[str]
        Lista original que se desea dividir.
    slice_size : int
        Tamaño de cada sublista.

    Retorna:
    --------
    list[list[str]]
        Lista de sublistas, cada una con una longitud máxima de 'slice_size'.
    """
    return [
        input_list[i : i + slice_size] for i in range(0, len(input_list), slice_size)
    ]