# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/04_pipelines.extractor.ipynb.

# %% auto 0
__all__ = ['Extractor']

# %% ../../nbs/04_pipelines.extractor.ipynb 3
import os
from typing import Any, Dict, Generator, List, Optional, Tuple, Union, Callable
import pandas as pd
from ..utils import segment

from ..ingest import load_single_document


class Extractor:
    def __init__(
        self,
        llm,
        prompt_template: Optional[str] = None,              
        **kwargs,
    ):
        """
        `Extractor` applies a given prompt to each sentence or paragraph in a document and returns the results.

        **Args:**

        - *llm*: An `onprem.LLM` object
        - *prompt_template*: A model specific prompt_template with a single placeholder named "{prompt}".
                             All prompts (e.g., Map-Reduce prompts) are wrapped within this prompt.
                             If supplied, overrides the `prompt_template` supplied to the `LLM` constructor.

        """
        self.llm = llm
        self.prompt_template = prompt_template if prompt_template is not None else llm.prompt_template



    def apply(self,
              ex_prompt_template:str, 
              fpath: Optional[str] = None,
              content: Optional[str] = None,
              unit:str='paragraph',
              filter_fn: Optional[Callable] = None,
              pdf_pages:List[int]=[],
              maxchars = 2048,
              stop:list=[]
             ):
        """
        Apply the prompt to each `unit` (where a "unit" is either a paragraph or sentence) optionally filtered by `filter_fn`.
        Results are stored in a `pandas.Dataframe`.


        **Args:**

        - *ex_prompt_template*: A prompt to apply to each `unit` in document. Should have a single variable, `{text}`.
                               Example: `"Extract universities from the following text delimited by ###:\n\n###{text}###"`
        - *fpath*: A path to to a single file of interest (e.g., a PDF or MS Word document). Mutually-exclusive with `content`.
        - *content*: Text content of a document of interest.  Mutually-exclusive with `fpath`.
        - *unit*: One of {'sentence', 'paragraph'}. 
        - *filter_fn*: A function that accepts a sentence or paragraph and returns `True` if prompt should be applied to it.
                       If `filter_fn` returns False, the text is ignored and excluded from results.
        - *pdf_pages*: If `fpath` is a PDF document, only apply prompt to text on page numbers listed in `pdf_pages`.
                       Page numbers start with 1, not 0 (e.g., `pdf_pages=[1,2,3]` for first three pages).
                       If list is empty, prompt is applied to every page.
        - *maxchars*: units (i.e., paragraphs or sentences) larger than `maxhcars` split.
        - *stop*: list of characters to trigger the LLM to stop generating.




        **Returns:**

        - pd.Dataframe: a Dataframe with results
        """
        if not(bool(fpath) != bool(content)):
            raise ValueError('Either fpath argument or content argument must be supplied but not both.')
            
        # setup extraction prompt
        extraction_prompt = ex_prompt_template if self.prompt_template is None else self.prompt_template.format(**{'prompt': ex_prompt_template})   

        # extract text
        if not content:
            if not os.path.isfile(fpath):
                raise ValueError(f'{fpath} is not a file')
            docs = load_single_document(fpath)
            ext = "." + fpath.rsplit(".", 1)[-1].lower()
            if ext == '.pdf' and pdf_pages:
                docs = [doc for i,doc in enumerate(docs) if i+1 in pdf_pages]
            content = '\n\n'.join([doc.page_content for doc in docs])
        
        # segment
        chunks = segment(content)
        extractions = []
        texts = []
        for chunk in chunks:
            if filter_fn and not filter_fn(chunk): continue
            prompt = extraction_prompt.format(text=chunk)
            extractions.append(self.llm.prompt(prompt, stop=stop))
            texts.append(chunk)
        df = pd.DataFrame({'Extractions':extractions, 'Texts':texts})
        return df
            
        return results
