# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/pdf/00_pdf_utils.ipynb.

# %% auto 0
__all__ = ['pdf_files', 'pdf_docs', 'pdf_to_docs', 'text_cat', 'docs_to_cats', 'add_cats_to_docs']

# %% ../../nbs/pdf/00_pdf_utils.ipynb 2
from dreamai.imports import *
from ..imports import *
from ..utils import *
from ..chains import *


# %% ../../nbs/pdf/00_pdf_utils.ipynb 4
def pdf_files(pdf_folder):
    return get_files(pdf_folder, extensions=[".pdf"])


def pdf_docs(pdf, splitter, proc=True):
    docs = PyPDFLoader(str(pdf)).load_and_split(text_splitter=splitter)
    if proc:
        docs = [proc_doc_text(doc) for doc in docs]
    return docs


def pdf_to_docs(
    path,
    splitter=None,
    chunk_size=200,
    chunk_overlap=60,
    separators=["\n\n", "\n", "(?<=\. )", " ", ""],
    add_start_index=True,
    proc=True,
):
    if splitter is None:
        splitter = RecursiveCharacterTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
            separators=separators,
            add_start_index=add_start_index,
        )
    docs = []
    for pdf in resolve_data_path(path):
        if Path(pdf).suffix == ".pdf":
            d = pdf_docs(pdf, splitter=splitter, proc=proc)
            if len(d) > 0:
                docs.append(d)
            # docs.append(pdf_docs(pdf, splitter=splitter, proc=proc))
    return docs


def text_cat(
    text,
    cats_model,
    thresh=0.6,
    classes=["Work Experience", "Education", "Certifications", "Other"],
    other_class="Other",
):
    if other_class not in classes:
        classes.append(other_class)
    ot_id = classes.index(other_class)
    probs, ids = cats_model.predict_proba(text).detach().cpu().max(1)
    ids[probs < thresh] = ot_id
    pred_classes = np.array(classes)[ids]
    return pred_classes.tolist()


def docs_to_cats(docs, cats_model):
    return text_cat([doc.page_content for doc in docs], cats_model)

def add_cats_to_docs(docs, cats_model, key="category"):
    fn = partial(docs_to_cats, cats_model=cats_model)
    return add_docs_metadata(docs, fn, key)
