# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/pdf/04_mappers.ipynb.

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

# %% ../../nbs/pdf/04_mappers.ipynb 2
from ..imports import *
from ..utils import *
from ..mapper import *
from .extract import *
from .df import *

# %% ../../nbs/pdf/04_mappers.ipynb 4
class SegsMapper(Mapper):
    """
    A custom mapper for PDF segmentation. It uses the `df_segs` function to extract the segments from the PDF.
    """

    def __init__(
        self,
        segs_model="HamzaFarhan/PDFSegs",  # The name of the model to use for segmentation.
        udf=df_segs,  # The function to use for extracting the segments.
        udf_kwargs=dict(  # The keyword arguments to pass to the `udf`.
            thresh=0.6,
            classes=["Work Experience", "Education", "Certifications", "Other"],
            other_class="Other",
        ),
        **kwargs,
    ):
        segs_model = load_segs_model(segs_model, device=default_device())
        udf_kwargs["segs_model"] = segs_model
        super().__init__(**locals_to_params(locals()))
