# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/pdf/01_pdf_chains.ipynb.

# %% auto 0
__all__ = ['pdf_docs_chain', 'pdf_cats_chain', 'pdf_ems_chain', 'docs_faiss_chain', 'pdf_faiss_chain', 'index_query_chain']

# %% ../../nbs/pdf/01_pdf_chains.ipynb 2
from ..imports import *
from ..chains import *
from .utils import *

# %% ../../nbs/pdf/01_pdf_chains.ipynb 3
def pdf_docs_chain(
    chunk_size=200,
    chunk_overlap=20,
    input_variables=["pdf_folder"],
    output_variables=["docs_df"],
    verbose=False,
):
    "Chain that takes a PDF folder and returns a DataFrame of Documents."
    pdf_chain = transform_chain(
        create_pdf_df,
        input_variables=input_variables,
        output_variables=["pdfs_df"],
    )
    docs_chain = transform_chain(
        df_pdf_docs,
        input_variables=["pdfs_df"],
        output_variables=output_variables,
        transform_kwargs={"chunk_size": chunk_size, "chunk_overlap": chunk_overlap},
    )
    return SequentialChain(
        chains=[pdf_chain, docs_chain],
        input_variables=input_variables,
        output_variables=output_variables,
        verbose=verbose,
    )


def pdf_cats_chain(
    cats_model,
    input_variables=["docs_df"],
    output_variables=["cats_df"],
):
    "Chain that takes a DataFrame of Documents and adds categories using a SetFit model."
    return transform_chain(
        df_docs_cat,
        input_variables=input_variables,
        output_variables=output_variables,
        transform_kwargs={"cats_model": cats_model},
    )


def pdf_ems_chain(
    ems_model,
    ems_folder,
    input_variables=["docs_df"],
    output_variables=["ems_df"],
):
    "Chain that takes a DataFrame of Documents and writes embeddings to `ems_folder` using `ems_model`."
    transform_chain(
        df_docs_ems,
        input_variables=input_variables,
        output_variables=output_variables,
        transform_kwargs={
            "ems_model": ems_model,
            "ems_folder": ems_folder,
        },
    )


def docs_faiss_chain(
    ems_model,
    index_folder,
    index_name,
    input_variables=["docs_df"],
    output_variables=["faiss_df"],
):
    "Chain that takes a DataFrame of Documents and adds them to a FAISS index in `index_folder`."
    return transform_chain(
        df_to_faiss,
        input_variables=input_variables,
        output_variables=output_variables,
        transform_kwargs={
            "ems_model": ems_model,
            "index_folder": index_folder,
            "index_name": index_name,
        },
    )


def pdf_faiss_chain(
    ems_model,  # The SentenceTransformer model to use for vectorestore embeddings.
    index_folder,  # The folder to store the FAISS index.
    index_name,  # The name of the FAISS index.
    input_variables=["pdf_folder"],  # The input key for the PDF folder.
    output_variables=["faiss_df"],  # The output key for the final DataFrame.
    chunk_size=200,  # The number of characters per Document.
    chunk_overlap=20,  # The number of characters to overlap between Documents.
    docs_block_size=1500,  # The number of Documents to process in a single Ray task.
    num_cpus=12,  # The number of CPUs to use for Ray.
    num_gpus=1,  # The number of GPUs to use for Ray.
    verbose=False,
):
    """
    Chain that takes a `pdf_folder` and adds them to FAISS indexes in `index_folder`.
    If there are more than `docs_block_size` Documents, it will be divided and distributed into multiple indexes using Ray.
    """
    docs_chain = pdf_docs_chain(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        input_variables=input_variables,
        output_variables=["docs_df"],
    )
    faiss_chain = ray_chain(
        docs_faiss_chain(
            ems_model,
            index_folder,
            index_name,
            input_variables=["docs_df"],
            output_variables=output_variables,
        ),
        block_size=docs_block_size,
        num_cpus=num_cpus,
        num_gpus=num_gpus,
    )
    return SequentialChain(
        chains=[docs_chain, faiss_chain],
        input_variables=input_variables,
        output_variables=output_variables,
        verbose=verbose,
    )


def index_query_chain(
    ems_model,  # The SentenceTransformer model to use for vectorestore embeddings.
    index_folder,  # The folder with the FAISS indexes.
    index_name,  # The name of the FAISS index.
    input_variables=["query"],  # The input key for the query.
    output_variables=["search_results"],  # The output key for the search results.
    k=2,  # The number of results to return.
    block_size=10,  # The number of indexes to process in a single Ray task.
    num_cpus=12,  # The number of CPUs to use for Ray.
    num_gpus=1,  # The number of GPUs to use for Ray.
    verbose=False,
):
    """
    Chain that takes a query and returns the top `k` results from the FAISS indexes in `index_folder`.
    If there are more than `block_size` indexes, search will be distributed using Ray.
    """
    q_df_chain = transform_chain(
        create_idx_q_df,
        input_variables=input_variables,
        output_variables=["q_df"],
        transform_kwargs={"index_folder": index_folder, "index_name": index_name},
    )

    def apply_search(df, ems_model, k):
        return df.apply(df_search_faiss, axis=1, ems_model=ems_model, k=k)

    search_chain = transform_chain(
        partial(apply_search, ems_model=ems_model, k=k),
        input_variables=["q_df"],
        output_variables=["search"],
    )

    def flatten_res(df, k):
        return [sorted(flatten_list(df.results), key=lambda x: x[1])[:k]]

    res_chain = transform_chain(
        flatten_res,
        transform_kwargs={"k": k},
        input_variables=["search"],
        output_variables=output_variables,
    )

    return ray_chain(
        SequentialChain(
            chains=[q_df_chain, search_chain, res_chain],
            input_variables=input_variables,
            output_variables=output_variables,
            verbose=verbose,
        ),
        block_size=block_size,
        num_cpus=num_cpus,
        num_gpus=num_gpus,
    )
