Skip to content

PII Masking

The PIIMasking task identifies and masks Personally Identifiable Information (PII) in documents.

Usage

from sieves import tasks

task = tasks.PIIMasking(
    model=model,
)

Results

The PIIMasking task returns a unified Result object containing the masked_text and a list of pii_entities.

class Result(pydantic.BaseModel):
    """Result of a PII masking task. Contains the masked text and the identified PII entities.

    PII entities should be masked with [MASKED].

    Attributes:
        masked_text: Masked version of text.
        pii_entities: List of PII entities.
    """

    masked_text: str = pydantic.Field(description="The original text with PII entities replaced by placeholders.")
    pii_entities: list[PIIEntity] = pydantic.Field(description="List of all PII entities identified in the text.")

Evaluation

Performance of the PII masking task can be measured using the .evaluate() method.

  • Metric: Corpus-wide Micro-F1 Score (F1). PII entities are matched based on their text span (start/end offsets) and type.
  • Requirement: Each document must have ground-truth PII entities stored in doc.gold[task_id].
report = task.evaluate(docs)
print(f"PII F1-Score: {report.metrics['F1']}")

Ground Truth Formats

Ground truth has to be specified in doc.meta using Result instances.


Allows masking of PII (Personally Identifiable Information) in text documents.

PIIMasking

Bases: PredictiveTask[TaskPromptSignature, TaskResult, _TaskBridge]

Task for masking PII (Personally Identifiable Information) in text documents.

Source code in sieves/tasks/predictive/pii_masking/core.py
 30
 31
 32
 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
 94
 95
 96
 97
 98
 99
100
101
102
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
160
161
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
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
class PIIMasking(PredictiveTask[TaskPromptSignature, TaskResult, _TaskBridge]):
    """Task for masking PII (Personally Identifiable Information) in text documents."""

    def __init__(
        self,
        model: TaskModel,
        pii_types: Sequence[str] | dict[str, str] | None = None,
        task_id: str | None = None,
        include_meta: bool = True,
        batch_size: int = -1,
        overwrite: bool = False,
        prompt_instructions: str | None = None,
        fewshot_examples: Sequence[FewshotExample] = (),
        model_settings: ModelSettings = ModelSettings(),
        condition: Callable[[Doc], bool] | None = None,
    ) -> None:
        """Initialize PIIMasking task.

        :param model: Model to use.
        :param pii_types: Types of PII to mask. Supports three formats:
            - None (default): Uses all common PII types
            - List format: `["EMAIL", "PHONE", "SSN"]`
            - Dict format: `{"EMAIL": "Email addresses", "PHONE": "Phone numbers"}`
            The dict format allows you to provide descriptions that help the model better identify PII.
        :param task_id: Task ID.
        :param include_meta: Whether to include meta information generated by the task.
        :param batch_size: Batch size to use for inference. Use -1 to process all documents at once.
        :param overwrite: Whether to overwrite original document text with masked text.
        :param prompt_instructions: Custom prompt instructions. If None, default instructions are used.
        :param fewshot_examples: Few-shot examples.
        :param model_settings: Settings for structured generation.
        :param condition: Optional callable that determines whether to process each document.
        """
        if pii_types is not None:
            if isinstance(pii_types, dict):
                self._pii_types = list(pii_types.keys())
                self._pii_type_descriptions = pii_types
            else:
                self._pii_types = list(pii_types)
                self._pii_type_descriptions = {}
        else:
            self._pii_types = None
            self._pii_type_descriptions = {}

        self._pii_types_param = pii_types
        self._mask_placeholder = "[MASKED]"

        super().__init__(
            model=model,
            task_id=task_id,
            include_meta=include_meta,
            batch_size=batch_size,
            overwrite=overwrite,
            prompt_instructions=prompt_instructions,
            fewshot_examples=fewshot_examples,
            model_settings=model_settings,
            condition=condition,
        )

    @property
    @override
    def fewshot_example_type(self) -> type[FewshotExample]:
        """Return few-shot example type.

        :return: Few-shot example type.
        """
        return FewshotExample

    @property
    @override
    def prompt_signature(self) -> type[pydantic.BaseModel]:
        return TaskResult

    @property
    @override
    def metric(self) -> str:
        return "F1"

    @override
    def _compute_metrics(self, truths: list[Any], preds: list[Any], judge: dspy.LM | None = None) -> dict[str, float]:
        """Compute corpus-level metrics.

        :param truths: List of ground truths.
        :param preds: List of predictions.
        :param judge: Optional DSPy LM instance to use as judge for generative tasks.
        :return: Dictionary of metrics.
        """
        tp = 0
        fp = 0
        fn = 0

        for gold, pred in zip(truths, preds):
            # Extract entities.
            if gold is not None:
                assert isinstance(gold, TaskResult)
                true_entities = {(e.entity_type, e.text) for e in gold.pii_entities}
            else:
                true_entities = set()

            if pred is not None:
                assert isinstance(pred, TaskResult)
                pred_entities = {(e.entity_type, e.text) for e in pred.pii_entities}
            else:
                pred_entities = set()

            tp += len(true_entities & pred_entities)
            fp += len(pred_entities - true_entities)
            fn += len(true_entities - pred_entities)

        precision = tp / (tp + fp) if (tp + fp) > 0 else 0.0
        recall = tp / (tp + fn) if (tp + fn) > 0 else 0.0
        f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0.0

        return {self.metric: f1}

    @override
    def _init_bridge(self, model_type: ModelType) -> _TaskBridge:
        bridge_types: dict[ModelType, type[_TaskBridge]] = {
            ModelType.dspy: DSPyPIIMasking,
            ModelType.langchain: PydanticPIIMasking,
            ModelType.outlines: PydanticPIIMasking,
        }

        try:
            return bridge_types[model_type](
                task_id=self._task_id,
                prompt_instructions=self._custom_prompt_instructions,
                mask_placeholder=self._mask_placeholder,
                pii_types=self._pii_types_param,
                overwrite=self._overwrite,
                model_settings=self._model_settings,
                prompt_signature=self.prompt_signature,
                model_type=model_type,
                fewshot_examples=self._fewshot_examples,
            )
        except KeyError as err:
            raise KeyError(f"Model type {model_type} is not supported by {self.__class__.__name__}.") from err

    @staticmethod
    @override
    def supports() -> set[ModelType]:
        return {
            ModelType.dspy,
            ModelType.langchain,
            ModelType.outlines,
        }

    @property
    @override
    def _state(self) -> dict[str, Any]:
        # Store pii_types as dict if descriptions exist, else as original value
        pii_types_state = self._pii_type_descriptions if self._pii_type_descriptions else self._pii_types
        return {
            **super()._state,
            "pii_types": pii_types_state,
        }

    @override
    def to_hf_dataset(self, docs: Iterable[Doc], threshold: float | None = None) -> datasets.Dataset:
        # Define metadata.
        features = datasets.Features(
            {
                "text": datasets.Value("string"),
                "masked_text": datasets.Value("string"),
            }
        )
        info = datasets.DatasetInfo(
            description=f"PII masking dataset. Generated with sieves v{Config.get_version()}.",
            features=features,
        )

        # Fetch data used for generating dataset.
        try:
            data = [(doc.text, doc.results[self._task_id].masked_text) for doc in docs]
        except KeyError as err:
            raise KeyError(f"Not all documents have results for this task with ID {self._task_id}") from err

        def generate_data() -> Iterable[dict[str, Any]]:
            """Yield results as dicts.

            :return: Results as dicts.
            """
            for text, masked_text in data:
                yield {"text": text, "masked_text": masked_text}

        # Create dataset.
        return datasets.Dataset.from_generator(generate_data, features=features, info=info)

    @override
    def distill(
        self,
        base_model_id: str,
        framework: DistillationFramework,
        data: datasets.Dataset | Sequence[Doc],
        output_path: Path | str,
        val_frac: float,
        init_kwargs: dict[str, Any] | None = None,
        train_kwargs: dict[str, Any] | None = None,
        seed: int | None = None,
    ) -> None:
        raise NotImplementedError

    @override
    def _evaluate_dspy_example(self, truth: dspy.Example, pred: dspy.Prediction, trace: Any, model: dspy.LM) -> float:
        # Compute entity detection F1 score based on (entity_type, text) pairs
        true_entities = {(e["entity_type"], e["text"]) for e in truth["pii_entities"]}
        pred_entities = {(e["entity_type"], e["text"]) for e in pred.get("pii_entities", [])}

        if not true_entities:
            return 1.0 if not pred_entities else 0.0

        precision = len(true_entities & pred_entities) / len(pred_entities) if pred_entities else 0
        recall = len(true_entities & pred_entities) / len(true_entities)
        return 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0

fewshot_example_type property

Return few-shot example type.

Returns:

Type Description
type[FewshotExample]

Few-shot example type.

fewshot_examples property

Return few-shot examples.

Returns:

Type Description
Sequence[FewshotExample]

Few-shot examples.

id property

Return task ID.

Used by pipeline for results and dependency management.

Returns:

Type Description
str

Task ID.

prompt_signature_description property

Return prompt signature description.

Returns:

Type Description
str | None

Prompt signature description.

prompt_template property

Return prompt template.

Returns:

Type Description
str

Prompt template.

__add__(other)

Chain this task with another task or pipeline using the + operator.

This returns a new Pipeline that executes this task first, followed by the task(s) in other. The original task(s)/pipeline are not mutated.

Cache semantics: - If other is a Pipeline, the resulting pipeline adopts other's use_cache setting (because the left-hand side is a single task). - If other is a Task, the resulting pipeline defaults to use_cache=True.

Parameters:

Name Type Description Default
other Task | Pipeline

A Task or Pipeline to execute after this task.

required

Returns:

Type Description
Pipeline

A new Pipeline representing the chained execution.

Raises:

Type Description
TypeError

If other is not a Task or Pipeline.

Source code in sieves/tasks/core.py
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
def __add__(self, other: Task | Pipeline) -> Pipeline:
    """Chain this task with another task or pipeline using the ``+`` operator.

    This returns a new ``Pipeline`` that executes this task first, followed by the
    task(s) in ``other``. The original task(s)/pipeline are not mutated.

    Cache semantics:
    - If ``other`` is a ``Pipeline``, the resulting pipeline adopts ``other``'s
      ``use_cache`` setting (because the left-hand side is a single task).
    - If ``other`` is a ``Task``, the resulting pipeline defaults to ``use_cache=True``.

    :param other: A ``Task`` or ``Pipeline`` to execute after this task.
    :return: A new ``Pipeline`` representing the chained execution.
    :raises TypeError: If ``other`` is not a ``Task`` or ``Pipeline``.
    """
    # Lazy import to avoid circular dependency at module import time.
    from sieves.pipeline import Pipeline

    if isinstance(other, Pipeline):
        return Pipeline(tasks=[self, *other.tasks], use_cache=other.use_cache)

    if isinstance(other, Task):
        return Pipeline(tasks=[self, other])

    raise TypeError(f"Cannot chain Task with {type(other).__name__}")

__call__(docs)

Execute task with conditional logic.

Checks the condition for each document without materializing all docs upfront. Passes all documents that pass the condition to _call() for proper batching. Documents that fail the condition have results[task_id] set to None.

Parameters:

Name Type Description Default
docs Iterable[Doc]

Docs to process.

required

Returns:

Type Description
Iterable[Doc]

Processed docs (in original order).

Source code in sieves/tasks/core.py
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
def __call__(self, docs: Iterable[Doc]) -> Iterable[Doc]:
    """Execute task with conditional logic.

    Checks the condition for each document without materializing all docs upfront.
    Passes all documents that pass the condition to _call() for proper batching.
    Documents that fail the condition have results[task_id] set to None.

    :param docs: Docs to process.
    :return: Processed docs (in original order).
    """
    docs = iter(docs) if not isinstance(docs, Iterator) else docs

    # Materialize docs in batches. This doesn't incur additional memory overhead, as docs are materialized in
    # batches downstream anyway.
    batch_size = self._batch_size if self._batch_size > 0 else sys.maxsize
    while docs_batch := [doc for doc in itertools.islice(docs, batch_size)]:
        # First pass: determine which docs pass the condition by index.
        passing_indices: set[int] = {
            idx for idx, doc in enumerate(docs_batch) if self._condition is None or self._condition(doc)
        }

        # Process all passing docs in one batch.
        processed = self._call(d for i, d in enumerate(docs_batch) if i in passing_indices)
        processed_iter = iter(processed) if not isinstance(processed, Iterator) else processed

        # Iterate through original docs in order and yield results.
        for idx, doc in enumerate(docs_batch):
            if idx in passing_indices:
                # Doc passed condition - use processed result.
                yield next(processed_iter)
            else:
                # Doc failed condition - set `None` result and yield original.
                doc.results[self.id] = None
                yield doc

__init__(model, pii_types=None, task_id=None, include_meta=True, batch_size=-1, overwrite=False, prompt_instructions=None, fewshot_examples=(), model_settings=ModelSettings(), condition=None)

Initialize PIIMasking task.

Parameters:

Name Type Description Default
model TaskModel

Model to use.

required
pii_types Sequence[str] | dict[str, str] | None

Types of PII to mask. Supports three formats: - None (default): Uses all common PII types - List format: ["EMAIL", "PHONE", "SSN"] - Dict format: {"EMAIL": "Email addresses", "PHONE": "Phone numbers"} The dict format allows you to provide descriptions that help the model better identify PII.

None
task_id str | None

Task ID.

None
include_meta bool

Whether to include meta information generated by the task.

True
batch_size int

Batch size to use for inference. Use -1 to process all documents at once.

-1
overwrite bool

Whether to overwrite original document text with masked text.

False
prompt_instructions str | None

Custom prompt instructions. If None, default instructions are used.

None
fewshot_examples Sequence[FewshotExample]

Few-shot examples.

()
model_settings ModelSettings

Settings for structured generation.

ModelSettings()
condition Callable[[Doc], bool] | None

Optional callable that determines whether to process each document.

None
Source code in sieves/tasks/predictive/pii_masking/core.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
def __init__(
    self,
    model: TaskModel,
    pii_types: Sequence[str] | dict[str, str] | None = None,
    task_id: str | None = None,
    include_meta: bool = True,
    batch_size: int = -1,
    overwrite: bool = False,
    prompt_instructions: str | None = None,
    fewshot_examples: Sequence[FewshotExample] = (),
    model_settings: ModelSettings = ModelSettings(),
    condition: Callable[[Doc], bool] | None = None,
) -> None:
    """Initialize PIIMasking task.

    :param model: Model to use.
    :param pii_types: Types of PII to mask. Supports three formats:
        - None (default): Uses all common PII types
        - List format: `["EMAIL", "PHONE", "SSN"]`
        - Dict format: `{"EMAIL": "Email addresses", "PHONE": "Phone numbers"}`
        The dict format allows you to provide descriptions that help the model better identify PII.
    :param task_id: Task ID.
    :param include_meta: Whether to include meta information generated by the task.
    :param batch_size: Batch size to use for inference. Use -1 to process all documents at once.
    :param overwrite: Whether to overwrite original document text with masked text.
    :param prompt_instructions: Custom prompt instructions. If None, default instructions are used.
    :param fewshot_examples: Few-shot examples.
    :param model_settings: Settings for structured generation.
    :param condition: Optional callable that determines whether to process each document.
    """
    if pii_types is not None:
        if isinstance(pii_types, dict):
            self._pii_types = list(pii_types.keys())
            self._pii_type_descriptions = pii_types
        else:
            self._pii_types = list(pii_types)
            self._pii_type_descriptions = {}
    else:
        self._pii_types = None
        self._pii_type_descriptions = {}

    self._pii_types_param = pii_types
    self._mask_placeholder = "[MASKED]"

    super().__init__(
        model=model,
        task_id=task_id,
        include_meta=include_meta,
        batch_size=batch_size,
        overwrite=overwrite,
        prompt_instructions=prompt_instructions,
        fewshot_examples=fewshot_examples,
        model_settings=model_settings,
        condition=condition,
    )

deserialize(config, **kwargs) classmethod

Generate PredictiveTask instance from config.

Parameters:

Name Type Description Default
config Config

Config to generate instance from.

required
kwargs dict[str, Any]

Values to inject into loaded config.

{}

Returns:

Type Description
PredictiveTask[TaskPromptSignature, TaskResult, TaskBridge]

Deserialized PredictiveTask instance.

Source code in sieves/tasks/predictive/core.py
256
257
258
259
260
261
262
263
264
265
266
267
268
269
@classmethod
def deserialize(
    cls, config: Config, **kwargs: dict[str, Any]
) -> PredictiveTask[TaskPromptSignature, TaskResult, TaskBridge]:
    """Generate PredictiveTask instance from config.

    :param config: Config to generate instance from.
    :param kwargs: Values to inject into loaded config.
    :return PredictiveTask[TaskPromptSignature, TaskResult, _TaskBridge]: Deserialized PredictiveTask instance.
    """
    init_dict = config.to_init_dict(cls, **kwargs)
    init_dict["model_settings"] = ModelSettings.model_validate(init_dict["model_settings"])

    return cls(**init_dict)

evaluate(docs, judge=None, failure_threshold=0.5)

Evaluate task performance using DSPy-based evaluation.

Parameters:

Name Type Description Default
docs Iterable[Doc]

Documents to evaluate.

required
judge LM | None

Optional DSPy LM instance to use as judge for generative tasks.

None
failure_threshold float

Decision threshold for whether to mark predicitions as failures.

0.5

Returns:

Type Description
TaskEvaluationReport

Evaluation report.

Source code in sieves/tasks/predictive/core.py
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
def evaluate(
    self, docs: Iterable[Doc], judge: dspy.LM | None = None, failure_threshold: float = 0.5
) -> TaskEvaluationReport:
    """Evaluate task performance using DSPy-based evaluation.

    :param docs: Documents to evaluate.
    :param judge: Optional DSPy LM instance to use as judge for generative tasks.
    :param failure_threshold: Decision threshold for whether to mark predicitions as failures.
    :return: Evaluation report.
    """
    truths: list[Any] = []
    preds: list[Any] = []
    failures: list[Doc] = []

    # Evaluate each doc individually to identify failed predictions.
    for doc in docs:
        if self.id not in doc.results:
            continue

        pred = doc.results[self.id]
        gold = doc.gold.get(self.id, None)

        # Accumulate for corpus-level metrics.
        truths.append(gold)
        preds.append(pred)

        # If gold or prediction is None: we cannot do proper evalution, so we just check whether they're both None
        # to compute score for failure analysis.
        if gold is None or pred is None:
            if gold is not None or pred is not None:
                failures.append(doc)
        else:
            # Convert result and gold to DSPy representation.
            truth = dspy.Example(**self._task_result_to_dspy_dict(gold))
            pred_dspy = dspy.Prediction(**self._task_result_to_dspy_dict(pred))

            # Call internal evaluation logic for per-doc failure analysis.
            score = self._evaluate_dspy_example(truth, pred_dspy, trace=None, model=judge)

            if score < failure_threshold:
                failures.append(doc)

    # Evaluate on corpus level to obtain representative metrics.
    metrics = self._compute_metrics(truths, preds, judge=judge)

    return TaskEvaluationReport(
        metrics=metrics,
        task_id=self.id,
        failures=failures,
    )

optimize(optimizer, verbose=True)

Optimize task prompt and few-shot examples with the available optimization config.

Updates task to use best prompt and few-shot examples found by the optimizer.

Parameters:

Name Type Description Default
optimizer Optimizer

Optimizer to run.

required
verbose bool

Whether to suppress output. DSPy produces a good amount of logs, so this can be useful to not pollute your terminal. Only warnings and errors will be printed.

True

Returns:

Type Description
tuple[str, Sequence[FewshotExample]]

Best found prompt and few-shot examples.

Source code in sieves/tasks/predictive/core.py
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
def optimize(self, optimizer: optimization.Optimizer, verbose: bool = True) -> tuple[str, Sequence[FewshotExample]]:
    """Optimize task prompt and few-shot examples with the available optimization config.

    Updates task to use best prompt and few-shot examples found by the optimizer.

    :param optimizer: Optimizer to run.
    :param verbose: Whether to suppress output. DSPy produces a good amount of logs, so this can be useful to
        not pollute your terminal. Only warnings and errors will be printed.

    :return tuple[str, Sequence[FewshotExample]]: Best found prompt and few-shot examples.
    """
    assert len(self._fewshot_examples) > 1, "At least two few-shot examples need to be provided to optimize."

    # Run optimizer to get best prompt and few-shot examples.
    signature = self._get_task_signature()
    dspy_examples = [ex.to_dspy() for ex in self._fewshot_examples]

    def _pred_eval(truth: dspy.Example, pred: dspy.Prediction, trace: Any | None = None) -> float:
        """Wrap optimization evaluation, inject model.

        :param truth: Ground truth.
        :param pred: Predicted value.
        :param trace: Optional trace information.
        :return: Metric value between 0.0 and 1.0.
        :raises KeyError: If target fields are missing from truth or prediction.
        :raises ValueError: If similarity score cannot be parsed from LLM response.
        """
        return self._evaluate_dspy_example(truth, pred, trace, model=optimizer.model)

    if verbose:
        best_prompt, best_examples = optimizer(signature, dspy_examples, _pred_eval, verbose=verbose)
    else:
        # Temporarily suppress DSPy logs.
        dspy_logger = logging.getLogger("dspy")
        optuna_logger = logging.getLogger("optuna")
        original_dspy_level = dspy_logger.level
        original_optuna_level = optuna_logger.level

        try:
            dspy_logger.setLevel(logging.ERROR)
            optuna_logger.setLevel(logging.ERROR)
            with warnings.catch_warnings():
                warnings.simplefilter("ignore")
                best_prompt, best_examples = optimizer(signature, dspy_examples, _pred_eval, verbose=verbose)
        finally:
            dspy_logger.setLevel(original_dspy_level)
            optuna_logger.setLevel(original_optuna_level)

    # Update few-shot examples and prompt instructions.
    fewshot_example_cls = self._fewshot_examples[0].__class__
    self._fewshot_examples = [fewshot_example_cls.from_dspy(ex) for ex in best_examples]
    self._validate_fewshot_examples()
    self._custom_prompt_instructions = best_prompt

    # Reinitialize bridge to use new prompt and few-shot examples.
    self._bridge = self._init_bridge(ModelType.get_model_type(self._model_wrapper))

    return best_prompt, self._fewshot_examples

serialize()

Serialize task.

Returns:

Type Description
Config

Config instance.

Source code in sieves/tasks/core.py
147
148
149
150
151
152
def serialize(self) -> Config:
    """Serialize task.

    :return: Config instance.
    """
    return Config.create(self.__class__, {k: Attribute(value=v) for k, v in self._state.items()})

Bridges for PII masking task.

DSPyPIIMasking

Bases: PIIMaskingBridge[PromptSignature, Result, InferenceMode]

DSPy bridge for PII masking.

Source code in sieves/tasks/predictive/pii_masking/bridges.py
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
160
161
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 DSPyPIIMasking(PIIMaskingBridge[dspy_.PromptSignature, dspy_.Result, dspy_.InferenceMode]):
    """DSPy bridge for PII masking."""

    @override
    def _validate(self) -> None:
        assert self._model_type == ModelType.dspy

    @override
    @property
    def _default_prompt_instructions(self) -> str:
        return ""

    @override
    @property
    def inference_mode(self) -> dspy_.InferenceMode:
        """Return inference mode for DSPy model wrapper."""
        return self._model_settings.inference_mode or dspy_.InferenceMode.predict

    @property
    @override
    def _chunk_extractor(self) -> Callable[[Any], Iterable[pydantic.BaseModel]]:
        return lambda res: res.pii_entities

    @override
    def integrate(self, results: Sequence[dspy_.Result], docs: list[Doc]) -> list[Doc]:
        """Integrate results into docs."""
        for doc, result in zip(docs, results):
            # Store masked text and PII entities in results
            res = Result(
                masked_text=result.masked_text,
                pii_entities=[
                    PIIEntity.model_validate(e.model_dump() if hasattr(e, "model_dump") else e)
                    for e in result.pii_entities
                ],
            )
            doc.results[self._task_id] = res

            if self._overwrite:
                doc.text = result.masked_text

        return docs

    @override
    def consolidate(
        self,
        results: Sequence[dspy_.Result],
        docs_offsets: list[tuple[int, int]],
    ) -> Sequence[dspy_.Result]:
        """Consolidate results from multiple chunks."""
        # Delegate consolidation of entities to strategy.
        consolidated_entities_all = self._consolidation_strategy.consolidate(results, docs_offsets)

        # Merge results for each document.
        consolidated_results: list[dspy_.Result] = []
        for i, (start, end) in enumerate(docs_offsets):
            doc_results = results[start:end]
            masked_texts: list[str] = []

            for res in doc_results:
                masked_texts.append(res.masked_text)

            consolidated_results.append(
                dspy.Prediction.from_completions(
                    {
                        "masked_text": [" ".join(masked_texts).strip()],
                        "pii_entities": [consolidated_entities_all[i]],
                    },
                    signature=self.prompt_signature,
                )
            )
        return consolidated_results

inference_mode property

Return inference mode for DSPy model wrapper.

model_settings property

Return model settings.

Returns:

Type Description
ModelSettings

Model settings.

model_type property

Return model type.

Returns:

Type Description
ModelType

Model type.

prompt_signature property

Create output signature.

E.g.: Signature in DSPy, Pydantic objects in outlines, JSON schema in jsonformers. This is model type-specific.

Returns:

Type Description
type[TaskPromptSignature] | TaskPromptSignature

Output signature object. This can be an instance (e.g. a regex string) or a class (e.g. a Pydantic class).

prompt_template property

Return prompt template.

Chains _prompt_instructions, _prompt_example_xml and _prompt_conclusion.

Note: different model have different expectations as to how a prompt should look like. E.g. outlines supports the Jinja 2 templating format for insertion of values and few-shot examples, whereas DSPy integrates these things in a different value in the workflow and hence expects the prompt not to include these things. Mind model-specific expectations when creating a prompt template.

Returns:

Type Description
str

Prompt template as string. None if not used by model wrapper.

__init__(task_id, prompt_instructions, mask_placeholder, pii_types, overwrite, model_settings, prompt_signature, model_type, fewshot_examples=())

Initialize PII masking bridge.

Parameters:

Name Type Description Default
task_id str

Task ID.

required
prompt_instructions str | None

Custom prompt instructions. If None, default instructions are used.

required
mask_placeholder str

Placeholder for masked PII.

required
pii_types Sequence[str] | dict[str, str] | None

PII types to mask.

required
overwrite bool

Whether to overwrite original text.

required
model_settings ModelSettings

Settings for structured generation.

required
prompt_signature type[BaseModel]

Unified Pydantic prompt signature.

required
model_type ModelType

Model type.

required
fewshot_examples Sequence[BaseModel]

Few-shot examples.

()
Source code in sieves/tasks/predictive/pii_masking/bridges.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
def __init__(
    self,
    task_id: str,
    prompt_instructions: str | None,
    mask_placeholder: str,
    pii_types: Sequence[str] | dict[str, str] | None,
    overwrite: bool,
    model_settings: ModelSettings,
    prompt_signature: type[pydantic.BaseModel],
    model_type: ModelType,
    fewshot_examples: Sequence[pydantic.BaseModel] = (),
):
    """Initialize PII masking bridge.

    :param task_id: Task ID.
    :param prompt_instructions: Custom prompt instructions. If None, default instructions are used.
    :param mask_placeholder: Placeholder for masked PII.
    :param pii_types: PII types to mask.
    :param overwrite: Whether to overwrite original text.
    :param model_settings: Settings for structured generation.
    :param prompt_signature: Unified Pydantic prompt signature.
    :param model_type: Model type.
    :param fewshot_examples: Few-shot examples.
    """
    super().__init__(
        task_id=task_id,
        prompt_instructions=prompt_instructions,
        overwrite=overwrite,
        model_settings=model_settings,
        prompt_signature=prompt_signature,
        model_type=model_type,
        fewshot_examples=fewshot_examples,
    )

    self._mask_placeholder = mask_placeholder
    self._pii_types: list[str] | None = None
    self._pii_type_descriptions: dict[str, str] = {}

    if isinstance(pii_types, dict):
        self._pii_types = list(pii_types.keys())
        self._pii_type_descriptions = pii_types
    elif pii_types is not None:
        self._pii_types = pii_types
        self._pii_type_descriptions = {}

    self._pii_entity_cls = self._create_pii_entity_cls()
    self._consolidation_strategy = MultiEntityConsolidation(extractor=self._chunk_extractor)

consolidate(results, docs_offsets)

Consolidate results from multiple chunks.

Source code in sieves/tasks/predictive/pii_masking/bridges.py
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
@override
def consolidate(
    self,
    results: Sequence[dspy_.Result],
    docs_offsets: list[tuple[int, int]],
) -> Sequence[dspy_.Result]:
    """Consolidate results from multiple chunks."""
    # Delegate consolidation of entities to strategy.
    consolidated_entities_all = self._consolidation_strategy.consolidate(results, docs_offsets)

    # Merge results for each document.
    consolidated_results: list[dspy_.Result] = []
    for i, (start, end) in enumerate(docs_offsets):
        doc_results = results[start:end]
        masked_texts: list[str] = []

        for res in doc_results:
            masked_texts.append(res.masked_text)

        consolidated_results.append(
            dspy.Prediction.from_completions(
                {
                    "masked_text": [" ".join(masked_texts).strip()],
                    "pii_entities": [consolidated_entities_all[i]],
                },
                signature=self.prompt_signature,
            )
        )
    return consolidated_results

extract(docs)

Extract all values from doc instances that are to be injected into the prompts.

Parameters:

Name Type Description Default
docs Sequence[Doc]

Docs to extract values from.

required

Returns:

Type Description
Sequence[dict[str, Any]]

All values from doc instances that are to be injected into the prompts as a sequence.

Source code in sieves/tasks/predictive/bridges.py
178
179
180
181
182
183
184
def extract(self, docs: Sequence[Doc]) -> Sequence[dict[str, Any]]:
    """Extract all values from doc instances that are to be injected into the prompts.

    :param docs: Docs to extract values from.
    :return: All values from doc instances that are to be injected into the prompts as a sequence.
    """
    return [{"text": doc.text if doc.text else None} for doc in docs]

integrate(results, docs)

Integrate results into docs.

Source code in sieves/tasks/predictive/pii_masking/bridges.py
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
@override
def integrate(self, results: Sequence[dspy_.Result], docs: list[Doc]) -> list[Doc]:
    """Integrate results into docs."""
    for doc, result in zip(docs, results):
        # Store masked text and PII entities in results
        res = Result(
            masked_text=result.masked_text,
            pii_entities=[
                PIIEntity.model_validate(e.model_dump() if hasattr(e, "model_dump") else e)
                for e in result.pii_entities
            ],
        )
        doc.results[self._task_id] = res

        if self._overwrite:
            doc.text = result.masked_text

    return docs

PIIMaskingBridge

Bases: Bridge[_BridgePromptSignature, _BridgeResult, ModelWrapperInferenceMode], ABC

Abstract base class for PII masking bridges.

Source code in sieves/tasks/predictive/pii_masking/bridges.py
 30
 31
 32
 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
 94
 95
 96
 97
 98
 99
100
101
102
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
class PIIMaskingBridge(Bridge[_BridgePromptSignature, _BridgeResult, ModelWrapperInferenceMode], abc.ABC):
    """Abstract base class for PII masking bridges."""

    def __init__(
        self,
        task_id: str,
        prompt_instructions: str | None,
        mask_placeholder: str,
        pii_types: Sequence[str] | dict[str, str] | None,
        overwrite: bool,
        model_settings: ModelSettings,
        prompt_signature: type[pydantic.BaseModel],
        model_type: ModelType,
        fewshot_examples: Sequence[pydantic.BaseModel] = (),
    ):
        """Initialize PII masking bridge.

        :param task_id: Task ID.
        :param prompt_instructions: Custom prompt instructions. If None, default instructions are used.
        :param mask_placeholder: Placeholder for masked PII.
        :param pii_types: PII types to mask.
        :param overwrite: Whether to overwrite original text.
        :param model_settings: Settings for structured generation.
        :param prompt_signature: Unified Pydantic prompt signature.
        :param model_type: Model type.
        :param fewshot_examples: Few-shot examples.
        """
        super().__init__(
            task_id=task_id,
            prompt_instructions=prompt_instructions,
            overwrite=overwrite,
            model_settings=model_settings,
            prompt_signature=prompt_signature,
            model_type=model_type,
            fewshot_examples=fewshot_examples,
        )

        self._mask_placeholder = mask_placeholder
        self._pii_types: list[str] | None = None
        self._pii_type_descriptions: dict[str, str] = {}

        if isinstance(pii_types, dict):
            self._pii_types = list(pii_types.keys())
            self._pii_type_descriptions = pii_types
        elif pii_types is not None:
            self._pii_types = pii_types
            self._pii_type_descriptions = {}

        self._pii_entity_cls = self._create_pii_entity_cls()
        self._consolidation_strategy = MultiEntityConsolidation(extractor=self._chunk_extractor)

    @property
    @abc.abstractmethod
    def _chunk_extractor(self) -> Callable[[Any], Iterable[pydantic.BaseModel]]:
        """Return a callable that extracts a list of entities from a raw chunk result.

        :return: Extractor callable.
        """

    def _get_pii_type_descriptions(self) -> str:
        """Return a string with the PII type descriptions.

        :return: A string with the PII type descriptions.
        """
        if not self._pii_types:
            return ""

        pii_types_with_descriptions: list[str] = []
        for pii_type in self._pii_types:
            if pii_type in self._pii_type_descriptions:
                pii_types_with_descriptions.append(
                    f"  <pii_type_description>\n    <pii_type>{pii_type}</pii_type>\n    <description>"
                    f"{self._pii_type_descriptions[pii_type]}</description>\n  </pii_type_description>"
                )
            else:
                pii_types_with_descriptions.append(f"  <pii_type>{pii_type}</pii_type>")

        pii_type_desc_string = "\n".join(pii_types_with_descriptions)
        return f"<pii_type_descriptions>\n{pii_type_desc_string}\n</pii_type_descriptions>"

    def _create_pii_entity_cls(self) -> type[pydantic.BaseModel]:
        """Create PII entity class.

        :returns: PII entity class.
        """
        pii_types_list = []
        if self._pii_types:
            for pt in self._pii_types:
                pii_types_list.append(pt)
                if pt.lower() not in pii_types_list:
                    pii_types_list.append(pt.lower())
                if pt.upper() not in pii_types_list:
                    pii_types_list.append(pt.upper())

        PIIType = Literal[*pii_types_list] if pii_types_list else str  # type: ignore[invalid-type-form]

        class PIIEntityRuntime(PIIEntity, frozen=True):
            """PII entity."""

            entity_type: PIIType  # type: ignore[valid-type]

        return PIIEntityRuntime

inference_mode abstractmethod property

Return inference mode.

Returns:

Type Description
ModelWrapperInferenceMode

Inference mode.

model_settings property

Return model settings.

Returns:

Type Description
ModelSettings

Model settings.

model_type property

Return model type.

Returns:

Type Description
ModelType

Model type.

prompt_signature property

Create output signature.

E.g.: Signature in DSPy, Pydantic objects in outlines, JSON schema in jsonformers. This is model type-specific.

Returns:

Type Description
type[TaskPromptSignature] | TaskPromptSignature

Output signature object. This can be an instance (e.g. a regex string) or a class (e.g. a Pydantic class).

prompt_template property

Return prompt template.

Chains _prompt_instructions, _prompt_example_xml and _prompt_conclusion.

Note: different model have different expectations as to how a prompt should look like. E.g. outlines supports the Jinja 2 templating format for insertion of values and few-shot examples, whereas DSPy integrates these things in a different value in the workflow and hence expects the prompt not to include these things. Mind model-specific expectations when creating a prompt template.

Returns:

Type Description
str

Prompt template as string. None if not used by model wrapper.

__init__(task_id, prompt_instructions, mask_placeholder, pii_types, overwrite, model_settings, prompt_signature, model_type, fewshot_examples=())

Initialize PII masking bridge.

Parameters:

Name Type Description Default
task_id str

Task ID.

required
prompt_instructions str | None

Custom prompt instructions. If None, default instructions are used.

required
mask_placeholder str

Placeholder for masked PII.

required
pii_types Sequence[str] | dict[str, str] | None

PII types to mask.

required
overwrite bool

Whether to overwrite original text.

required
model_settings ModelSettings

Settings for structured generation.

required
prompt_signature type[BaseModel]

Unified Pydantic prompt signature.

required
model_type ModelType

Model type.

required
fewshot_examples Sequence[BaseModel]

Few-shot examples.

()
Source code in sieves/tasks/predictive/pii_masking/bridges.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
def __init__(
    self,
    task_id: str,
    prompt_instructions: str | None,
    mask_placeholder: str,
    pii_types: Sequence[str] | dict[str, str] | None,
    overwrite: bool,
    model_settings: ModelSettings,
    prompt_signature: type[pydantic.BaseModel],
    model_type: ModelType,
    fewshot_examples: Sequence[pydantic.BaseModel] = (),
):
    """Initialize PII masking bridge.

    :param task_id: Task ID.
    :param prompt_instructions: Custom prompt instructions. If None, default instructions are used.
    :param mask_placeholder: Placeholder for masked PII.
    :param pii_types: PII types to mask.
    :param overwrite: Whether to overwrite original text.
    :param model_settings: Settings for structured generation.
    :param prompt_signature: Unified Pydantic prompt signature.
    :param model_type: Model type.
    :param fewshot_examples: Few-shot examples.
    """
    super().__init__(
        task_id=task_id,
        prompt_instructions=prompt_instructions,
        overwrite=overwrite,
        model_settings=model_settings,
        prompt_signature=prompt_signature,
        model_type=model_type,
        fewshot_examples=fewshot_examples,
    )

    self._mask_placeholder = mask_placeholder
    self._pii_types: list[str] | None = None
    self._pii_type_descriptions: dict[str, str] = {}

    if isinstance(pii_types, dict):
        self._pii_types = list(pii_types.keys())
        self._pii_type_descriptions = pii_types
    elif pii_types is not None:
        self._pii_types = pii_types
        self._pii_type_descriptions = {}

    self._pii_entity_cls = self._create_pii_entity_cls()
    self._consolidation_strategy = MultiEntityConsolidation(extractor=self._chunk_extractor)

consolidate(results, docs_offsets) abstractmethod

Consolidate results for document chunks into document results.

Parameters:

Name Type Description Default
results Sequence[TaskResult]

Results per document chunk.

required
docs_offsets list[tuple[int, int]]

Chunk offsets per document. Chunks per document can be obtained with results[docs_chunk_offsets[i][0]:docs_chunk_offsets[i][1]].

required

Returns:

Type Description
Sequence[TaskResult]

Results per document as a sequence.

Source code in sieves/tasks/predictive/bridges.py
195
196
197
198
199
200
201
202
203
@abc.abstractmethod
def consolidate(self, results: Sequence[TaskResult], docs_offsets: list[tuple[int, int]]) -> Sequence[TaskResult]:
    """Consolidate results for document chunks into document results.

    :param results: Results per document chunk.
    :param docs_offsets: Chunk offsets per document. Chunks per document can be obtained with
        `results[docs_chunk_offsets[i][0]:docs_chunk_offsets[i][1]]`.
    :return: Results per document as a sequence.
    """

extract(docs)

Extract all values from doc instances that are to be injected into the prompts.

Parameters:

Name Type Description Default
docs Sequence[Doc]

Docs to extract values from.

required

Returns:

Type Description
Sequence[dict[str, Any]]

All values from doc instances that are to be injected into the prompts as a sequence.

Source code in sieves/tasks/predictive/bridges.py
178
179
180
181
182
183
184
def extract(self, docs: Sequence[Doc]) -> Sequence[dict[str, Any]]:
    """Extract all values from doc instances that are to be injected into the prompts.

    :param docs: Docs to extract values from.
    :return: All values from doc instances that are to be injected into the prompts as a sequence.
    """
    return [{"text": doc.text if doc.text else None} for doc in docs]

integrate(results, docs) abstractmethod

Integrate results into Doc instances.

Parameters:

Name Type Description Default
results Sequence[TaskResult]

Results from prompt executable.

required
docs list[Doc]

Doc instances to update.

required

Returns:

Type Description
list[Doc]

Updated doc instances as a list.

Source code in sieves/tasks/predictive/bridges.py
186
187
188
189
190
191
192
193
@abc.abstractmethod
def integrate(self, results: Sequence[TaskResult], docs: list[Doc]) -> list[Doc]:
    """Integrate results into Doc instances.

    :param results: Results from prompt executable.
    :param docs: Doc instances to update.
    :return: Updated doc instances as a list.
    """

PydanticPIIMasking

Bases: PIIMaskingBridge[BaseModel, BaseModel, ModelWrapperInferenceMode], ABC

Base class for Pydantic-based PII masking bridges.

Source code in sieves/tasks/predictive/pii_masking/bridges.py
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
class PydanticPIIMasking(PIIMaskingBridge[pydantic.BaseModel, pydantic.BaseModel, ModelWrapperInferenceMode], abc.ABC):
    """Base class for Pydantic-based PII masking bridges."""

    @override
    def _validate(self) -> None:
        assert self._model_type in {ModelType.langchain, ModelType.outlines}

    @property
    @override
    def _chunk_extractor(self) -> Callable[[Any], Iterable[pydantic.BaseModel]]:
        return lambda res: res.pii_entities

    @property
    def _default_prompt_instructions(self) -> str:
        pii_type_info = self._get_pii_type_descriptions() if self._pii_type_descriptions else ""
        return (
            "Identify and mask Personally Identifiable Information (PII) in the given text.\n"
            "{%- if pii_types|length > 0 %}\n"
            "Focus on these specific PII types: {{ pii_types|join(', ') }}.\n"
            "{%- else %}\n"
            "Mask all common types of PII such as names, addresses, phone numbers, emails, SSNs, credit card numbers, "
            "etc.\n"
            "{%- endif %}\n"
            f"{pii_type_info}\n"
            f'Replace each instance of PII with "{self._mask_placeholder}".\n'
            "Provide a confidence score between 0.0 and 1.0 for each entity found."
        )

    @override
    @property
    def _prompt_conclusion(self) -> str | None:
        return "========\n\n<text>{{ text }}</text>"

    @override
    def integrate(self, results: Sequence[pydantic.BaseModel], docs: list[Doc]) -> list[Doc]:
        for doc, result in zip(docs, results):
            assert hasattr(result, "masked_text")
            assert hasattr(result, "pii_entities")
            # Store masked text and PII entities in results
            doc.results[self._task_id] = Result(
                masked_text=result.masked_text,
                pii_entities=[PIIEntity.model_validate(e) for e in result.pii_entities],
            )

            if self._overwrite:
                doc.text = result.masked_text

        return docs

    @override
    def consolidate(
        self,
        results: Sequence[pydantic.BaseModel],
        docs_offsets: list[tuple[int, int]],
    ) -> Sequence[pydantic.BaseModel]:
        assert issubclass(self.prompt_signature, pydantic.BaseModel)

        consolidated_entities_all = self._consolidation_strategy.consolidate(results, docs_offsets)
        consolidated_results: list[pydantic.BaseModel] = []

        for i, (start, end) in enumerate(docs_offsets):
            doc_results = results[start:end]
            masked_texts: list[str] = []

            for res in doc_results:
                if res is None:
                    continue  # type: ignore[unreachable]

                assert hasattr(res, "masked_text")
                masked_texts.append(res.masked_text)

            consolidated_results.append(
                self.prompt_signature(
                    masked_text=" ".join(masked_texts).strip(),
                    pii_entities=consolidated_entities_all[i],
                )
            )

        return consolidated_results

    @override
    @property
    def model_type(self) -> ModelType:
        return self._model_type

    @override
    @property
    def inference_mode(self) -> outlines_.InferenceMode | langchain_.InferenceMode:
        if self._model_type == ModelType.outlines:
            return self._model_settings.inference_mode or outlines_.InferenceMode.json
        elif self._model_type == ModelType.langchain:
            return self._model_settings.inference_mode or langchain_.InferenceMode.structured

        raise ValueError(f"Unsupported model type: {self._model_type}")

model_settings property

Return model settings.

Returns:

Type Description
ModelSettings

Model settings.

prompt_signature property

Create output signature.

E.g.: Signature in DSPy, Pydantic objects in outlines, JSON schema in jsonformers. This is model type-specific.

Returns:

Type Description
type[TaskPromptSignature] | TaskPromptSignature

Output signature object. This can be an instance (e.g. a regex string) or a class (e.g. a Pydantic class).

prompt_template property

Return prompt template.

Chains _prompt_instructions, _prompt_example_xml and _prompt_conclusion.

Note: different model have different expectations as to how a prompt should look like. E.g. outlines supports the Jinja 2 templating format for insertion of values and few-shot examples, whereas DSPy integrates these things in a different value in the workflow and hence expects the prompt not to include these things. Mind model-specific expectations when creating a prompt template.

Returns:

Type Description
str

Prompt template as string. None if not used by model wrapper.

__init__(task_id, prompt_instructions, mask_placeholder, pii_types, overwrite, model_settings, prompt_signature, model_type, fewshot_examples=())

Initialize PII masking bridge.

Parameters:

Name Type Description Default
task_id str

Task ID.

required
prompt_instructions str | None

Custom prompt instructions. If None, default instructions are used.

required
mask_placeholder str

Placeholder for masked PII.

required
pii_types Sequence[str] | dict[str, str] | None

PII types to mask.

required
overwrite bool

Whether to overwrite original text.

required
model_settings ModelSettings

Settings for structured generation.

required
prompt_signature type[BaseModel]

Unified Pydantic prompt signature.

required
model_type ModelType

Model type.

required
fewshot_examples Sequence[BaseModel]

Few-shot examples.

()
Source code in sieves/tasks/predictive/pii_masking/bridges.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
def __init__(
    self,
    task_id: str,
    prompt_instructions: str | None,
    mask_placeholder: str,
    pii_types: Sequence[str] | dict[str, str] | None,
    overwrite: bool,
    model_settings: ModelSettings,
    prompt_signature: type[pydantic.BaseModel],
    model_type: ModelType,
    fewshot_examples: Sequence[pydantic.BaseModel] = (),
):
    """Initialize PII masking bridge.

    :param task_id: Task ID.
    :param prompt_instructions: Custom prompt instructions. If None, default instructions are used.
    :param mask_placeholder: Placeholder for masked PII.
    :param pii_types: PII types to mask.
    :param overwrite: Whether to overwrite original text.
    :param model_settings: Settings for structured generation.
    :param prompt_signature: Unified Pydantic prompt signature.
    :param model_type: Model type.
    :param fewshot_examples: Few-shot examples.
    """
    super().__init__(
        task_id=task_id,
        prompt_instructions=prompt_instructions,
        overwrite=overwrite,
        model_settings=model_settings,
        prompt_signature=prompt_signature,
        model_type=model_type,
        fewshot_examples=fewshot_examples,
    )

    self._mask_placeholder = mask_placeholder
    self._pii_types: list[str] | None = None
    self._pii_type_descriptions: dict[str, str] = {}

    if isinstance(pii_types, dict):
        self._pii_types = list(pii_types.keys())
        self._pii_type_descriptions = pii_types
    elif pii_types is not None:
        self._pii_types = pii_types
        self._pii_type_descriptions = {}

    self._pii_entity_cls = self._create_pii_entity_cls()
    self._consolidation_strategy = MultiEntityConsolidation(extractor=self._chunk_extractor)

extract(docs)

Extract all values from doc instances that are to be injected into the prompts.

Parameters:

Name Type Description Default
docs Sequence[Doc]

Docs to extract values from.

required

Returns:

Type Description
Sequence[dict[str, Any]]

All values from doc instances that are to be injected into the prompts as a sequence.

Source code in sieves/tasks/predictive/bridges.py
178
179
180
181
182
183
184
def extract(self, docs: Sequence[Doc]) -> Sequence[dict[str, Any]]:
    """Extract all values from doc instances that are to be injected into the prompts.

    :param docs: Docs to extract values from.
    :return: All values from doc instances that are to be injected into the prompts as a sequence.
    """
    return [{"text": doc.text if doc.text else None} for doc in docs]

Schemas for PII masking task.

FewshotExample

Bases: FewshotExample

Example for PII masking few-shot prompting.

Attributes: text: Input text. masked_text: Masked version of text. pii_entities: List of PII entities.

Source code in sieves/tasks/predictive/schemas/pii_masking.py
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
class FewshotExample(BaseFewshotExample):
    """Example for PII masking few-shot prompting.

    Attributes:
        text: Input text.
        masked_text: Masked version of text.
        pii_entities: List of PII entities.
    """

    masked_text: str
    pii_entities: list[PIIEntity]

    @property
    def target_fields(self) -> tuple[str, ...]:
        """Return target fields.

        :return: Target fields.
        """
        return ("masked_text", "pii_entities")

input_fields property

Defines which fields are inputs.

Returns:

Type Description
Sequence[str]

Sequence of field names.

target_fields property

Return target fields.

Returns:

Type Description
tuple[str, ...]

Target fields.

from_dspy(example) classmethod

Convert from dspy.Example.

Parameters:

Name Type Description Default
example Example

Example as dspy.Example.

required

Returns:

Type Description
Self

Example as FewshotExample.

Source code in sieves/tasks/predictive/schemas/core.py
63
64
65
66
67
68
69
70
@classmethod
def from_dspy(cls, example: dspy.Example) -> Self:
    """Convert from `dspy.Example`.

    :param example: Example as `dspy.Example`.
    :returns: Example as `FewshotExample`.
    """
    return cls(**example)

to_dspy()

Convert to dspy.Example.

Returns:

Type Description
Example

Example as dspy.Example.

Source code in sieves/tasks/predictive/schemas/core.py
56
57
58
59
60
61
def to_dspy(self) -> dspy.Example:
    """Convert to `dspy.Example`.

    :returns: Example as `dspy.Example`.
    """
    return dspy.Example(**ModelWrapper.convert_fewshot_examples([self])[0]).with_inputs(*self.input_fields)

PIIEntity

Bases: BaseModel

Personally Identifiable Information (PII) entity.

Attributes: entity_type: Type of PII. text: Entity text. score: Confidence score.

Source code in sieves/tasks/predictive/schemas/pii_masking.py
12
13
14
15
16
17
18
19
20
21
22
23
24
25
class PIIEntity(pydantic.BaseModel, frozen=True):
    """Personally Identifiable Information (PII) entity.

    Attributes:
        entity_type: Type of PII.
        text: Entity text.
        score: Confidence score.
    """

    entity_type: str = pydantic.Field(description="The type of PII identified (e.g., EMAIL, PHONE, SSN).")
    text: str = pydantic.Field(description="The original text of the PII entity.")
    score: float | None = pydantic.Field(
        default=None, description="Provide a confidence score for the PII identification, between 0 and 1."
    )

Result

Bases: BaseModel

Result of a PII masking task. Contains the masked text and the identified PII entities.

PII entities should be masked with [MASKED].

Attributes: masked_text: Masked version of text. pii_entities: List of PII entities.

Source code in sieves/tasks/predictive/schemas/pii_masking.py
50
51
52
53
54
55
56
57
58
59
60
61
class Result(pydantic.BaseModel):
    """Result of a PII masking task. Contains the masked text and the identified PII entities.

    PII entities should be masked with [MASKED].

    Attributes:
        masked_text: Masked version of text.
        pii_entities: List of PII entities.
    """

    masked_text: str = pydantic.Field(description="The original text with PII entities replaced by placeholders.")
    pii_entities: list[PIIEntity] = pydantic.Field(description="List of all PII entities identified in the text.")