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

# %% auto 0
__all__ = ['text_to_segments', 'print_segments', 'load_segs_model', 'load_ems_model', 'write_segments']

# %% ../nbs/01_segment.ipynb 2
from .imports import *
from .core import *

# %% ../nbs/01_segment.ipynb 3
def text_to_segments(text: Union[str, List[str]], segs_model: SetFitModel,
                     thresh: float=0.6, classes: List=['Work Experience', 'Education', 'Certifications', 'Other']):
    preds = segs_model.predict(text).detach().cpu().numpy()
    probs = segs_model.predict_proba(text).detach().cpu().numpy().max(1)
    preds[probs < thresh] = 3
    pred_classes = [classes[p] for p in preds]
    segments = {pc:[] for pc in classes}
    for i,pc in enumerate(pred_classes):
        txt = text[i]
        if txt not in segments[pc]:
            segments[pc].append(text[i])
    return segments, pred_classes, probs

def print_segments(segments: dict, limit: int=10):
    for k,v in segments.items():
        print(f'{k}: {len(v)}')
        for s in v[:limit]:
            print('\t', end='')
            pprint(s, width=100)
            print()
        print('-'*100)
        
def load_segs_model(model_name: str='HamzaFarhan/PDFSegs', device: Union[str, torch.device]='cpu'):
    return SetFitModel.from_pretrained(model_name).to(device)

def load_ems_model(model_name: str='HamzaFarhan/PDFSegs', device: Union[str, torch.device]='cpu'):
    return SentenceTransformer(model_name, device=device)

def write_segments(segs_model: SetFitModel, data_path: Union[str, Path, List[Union[str, Path]]], output_path: Union[str, Path]='pdf_segments'):
    """
    Extracts text from PDFs and writes segments to JSON files.
    
    Parameters
    ----------
    segs_model : SetFitModel
        SetFit model for segment classification.
    data_path : Union[str, Path, List[Union[str, Path]]]
        Path to PDFs. Can be a single file, a directory, or a list of files/directories.
    output_path : Union[str, Path]
        Folder to write JSON files. Defaults to 'pdf_segments'.
    
    Returns
    -------
    None
    """
    output_path = Path(output_path)
    os.makedirs(output_path, exist_ok=True)
    pdfs = resolve_data_path(data_path)
    for file in pdfs:
        try:
            pdf_text = extract_text(file)[str(file)]
            segments, _, _ = text_to_segments(pdf_text, segs_model, thresh=0.6)
            fn = (output_path/Path(file).stem).with_suffix('.json')
            with open(fn, 'w') as f:
                json.dump(segments, f, indent=4)
        except Exception as e:
            print(f'\nCould not write segments for file: {str(file)}\n{e}')
