Skip to content

base

BaseClassificationProvider

Bases: AbstractPageTaskProvider[bytes, ClassificationConfig, ClassificationOutput]

The base classification provider.

Source code in docprompt/tasks/classification/base.py 
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
class BaseClassificationProvider(
    AbstractPageTaskProvider[bytes, ClassificationConfig, ClassificationOutput]
):
    """
    The base classification provider.
    """

    capabilities = [PageLevelCapabilities.PAGE_CLASSIFICATION]

    class Meta:
        abstract = True

    def process_document_node(
        self,
        document_node: "DocumentNode",
        task_config: ClassificationConfig = None,
        start: Optional[int] = None,
        stop: Optional[int] = None,
        contribute_to_document: bool = True,
        **kwargs,
    ):
        assert (
            task_config is not None
        ), "task_config must be provided for classification tasks"

        raster_bytes = []
        for page_number in range(start or 1, (stop or len(document_node)) + 1):
            image_bytes = document_node.page_nodes[
                page_number - 1
            ].rasterizer.rasterize("default")
            raster_bytes.append(image_bytes)

        # TODO: This is a somewhat dangerous way of requiring these kwargs to be drilled
        # through, potentially a decorator solution to be had here
        kwargs = {**self._default_invoke_kwargs, **kwargs}
        results = self._invoke(raster_bytes, config=task_config, **kwargs)

        return {
            i: res
            for i, res in zip(
                range(start or 1, (stop or len(document_node)) + 1), results
            )
        }

BasePageClassificationOutputParser

Bases: ABC, BaseOutputParser[ClassificationConfig, ClassificationOutput]

The output parser for the page classification system.

Source code in docprompt/tasks/classification/base.py 
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
class BasePageClassificationOutputParser(
    ABC, BaseOutputParser[ClassificationConfig, ClassificationOutput]
):
    """The output parser for the page classification system."""

    name: str = Field(...)
    type: ClassificationTypes = Field(...)
    labels: LabelType = Field(...)
    confidence: bool = Field(False)

    @classmethod
    def from_task_input(cls, task_input: ClassificationConfig, provider_name: str):
        return cls(
            type=task_input.type,
            name=provider_name,
            labels=task_input.labels,
            confidence=task_input.confidence,
        )

    def resolve_match(self, _match: Union[re.Match, None]) -> LabelType:
        """Get the regex pattern for the output parser."""

        if not _match:
            raise ValueError("Could not find the answer in the text.")

        val = _match.group(1)
        if self.type == ClassificationTypes.BINARY:
            if val not in self.labels:
                raise ValueError(f"Invalid label: {val}")
            return val

        elif self.type == ClassificationTypes.SINGLE_LABEL:
            if val not in self.labels:
                raise ValueError(f"Invalid label: {val}")
            return val

        elif self.type == ClassificationTypes.MULTI_LABEL:
            labels = val.split(", ")
            for label in labels:
                if label not in self.labels:
                    raise ValueError(f"Invalid label: {label}")
            return labels
        else:
            raise ValueError(f"Invalid classification type: {self.type}")

    def resolve_confidence(self, _match: Union[re.Match, None]) -> ConfidenceLevel:
        """Get the confidence level from the text."""

        if not _match:
            return None

        val = _match.group(1).lower()

        return ConfidenceLevel(val)

    @abstractmethod
    def parse(self, text: str) -> ClassificationOutput: ...

resolve_confidence(_match)

Get the confidence level from the text.

Source code in docprompt/tasks/classification/base.py 
148
149
150
151
152
153
154
155
156
def resolve_confidence(self, _match: Union[re.Match, None]) -> ConfidenceLevel:
    """Get the confidence level from the text."""

    if not _match:
        return None

    val = _match.group(1).lower()

    return ConfidenceLevel(val)

resolve_match(_match)

Get the regex pattern for the output parser.

Source code in docprompt/tasks/classification/base.py 
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
def resolve_match(self, _match: Union[re.Match, None]) -> LabelType:
    """Get the regex pattern for the output parser."""

    if not _match:
        raise ValueError("Could not find the answer in the text.")

    val = _match.group(1)
    if self.type == ClassificationTypes.BINARY:
        if val not in self.labels:
            raise ValueError(f"Invalid label: {val}")
        return val

    elif self.type == ClassificationTypes.SINGLE_LABEL:
        if val not in self.labels:
            raise ValueError(f"Invalid label: {val}")
        return val

    elif self.type == ClassificationTypes.MULTI_LABEL:
        labels = val.split(", ")
        for label in labels:
            if label not in self.labels:
                raise ValueError(f"Invalid label: {label}")
        return labels
    else:
        raise ValueError(f"Invalid classification type: {self.type}")

ClassificationConfig

Bases: BaseModel

Source code in docprompt/tasks/classification/base.py 
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
class ClassificationConfig(BaseModel):
    type: ClassificationTypes
    labels: LabelType
    descriptions: Optional[List[str]] = Field(
        None, description="The descriptions for each label (if any)."
    )

    instructions: Optional[str] = Field(
        None,
        description="Additional instructions to pass to the LLM for the task. Required for Binary Classification.",
    )

    confidence: bool = Field(False)

    @model_validator(mode="before")
    def validate_label_bindings(cls, data: Any) -> Any:
        """Validate the the label/description bindings based on the type."""

        classification_type = data.get("type", None)
        if classification_type == ClassificationTypes.SINGLE_LABEL:
            labels = data.get("labels", None)
            if not labels:
                raise ValueError(
                    "labels must be provided for single_label classification"
                )
            return data

        elif classification_type == ClassificationTypes.BINARY:
            instructions = data.get("instructions", None)
            if not instructions:
                raise ValueError(
                    "instructions must be provided for binary classification"
                )
            data["labels"] = ["YES", "NO"]
            return data

        elif classification_type == ClassificationTypes.MULTI_LABEL:
            labels = data.get("labels", None)
            if not labels:
                raise ValueError(
                    "labels must be provided for multi_label classification"
                )
            return data

    @model_validator(mode="after")
    def validate_descriptions_length(self):
        if self.descriptions is not None:
            labels = self.labels
            if labels is not None and len(self.descriptions) != len(labels):
                raise ValueError("descriptions must have the same length as labels")
        return self

    @property
    def formatted_labels(self):
        """Produce the formatted labels for the prompt template."""
        raw_labels = self.labels
        if self.descriptions:
            for label, description in zip(raw_labels, self.descriptions):
                yield f"{label}: {description}"
        else:
            yield from raw_labels

formatted_labels property

Produce the formatted labels for the prompt template.

validate_label_bindings(data)

Validate the the label/description bindings based on the type.

Source code in docprompt/tasks/classification/base.py 
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
@model_validator(mode="before")
def validate_label_bindings(cls, data: Any) -> Any:
    """Validate the the label/description bindings based on the type."""

    classification_type = data.get("type", None)
    if classification_type == ClassificationTypes.SINGLE_LABEL:
        labels = data.get("labels", None)
        if not labels:
            raise ValueError(
                "labels must be provided for single_label classification"
            )
        return data

    elif classification_type == ClassificationTypes.BINARY:
        instructions = data.get("instructions", None)
        if not instructions:
            raise ValueError(
                "instructions must be provided for binary classification"
            )
        data["labels"] = ["YES", "NO"]
        return data

    elif classification_type == ClassificationTypes.MULTI_LABEL:
        labels = data.get("labels", None)
        if not labels:
            raise ValueError(
                "labels must be provided for multi_label classification"
            )
        return data

ConfidenceLevel

Bases: str, Enum

The confidence level of the classification.

Source code in docprompt/tasks/classification/base.py 
19
20
21
22
23
24
class ConfidenceLevel(str, Enum):
    """The confidence level of the classification."""

    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"