Skip to content

PII Masking

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

FewshotExample

Bases: FewshotExample

Example for PII masking few-shot prompting.

Source code in sieves/tasks/predictive/pii_masking/core.py 
37
38
39
40
41
42
43
44
45
46
47
class FewshotExample(BaseFewshotExample):
    """Example for PII masking few-shot prompting."""

    reasoning: str
    masked_text: str
    pii_entities: list[PIIEntity]

    @override
    @property
    def target_fields(self) -> Sequence[str]:
        return "masked_text", "pii_entities"

input_fields property

Defines which fields are inputs.

Returns:

Type Description
Sequence[str]

Sequence of field names.

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/core.py 
76
77
78
79
80
81
82
83
@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/core.py 
69
70
71
72
73
74
def to_dspy(self) -> dspy.Example:
    """Convert to `dspy.Example`.

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

PIIEntity

Bases: BaseModel

PII entity.

Source code in sieves/tasks/predictive/pii_masking/core.py 
30
31
32
33
34
class PIIEntity(pydantic.BaseModel, frozen=True):
    """PII entity."""

    entity_type: str
    text: str

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 
 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
class PIIMasking(PredictiveTask[_TaskPromptSignature, _TaskResult, _TaskBridge]):
    """Task for masking PII (Personally Identifiable Information) in text documents."""

    def __init__(
        self,
        model: _TaskModel,
        pii_types: list[str] | None = None,
        mask_placeholder: str = "[MASKED]",
        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] = (),
        generation_settings: GenerationSettings = GenerationSettings(),
    ) -> None:
        """
        Initialize PIIMasking task.

        :param model: Model to use.
        :param pii_types: Types of PII to mask. If None, all common PII types will be masked.
                         E.g., ["NAME", "EMAIL", "PHONE", "ADDRESS", "SSN", "CREDIT_CARD", "DATE_OF_BIRTH"]
        :param mask_placeholder: String to replace PII with.
        :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. If None,
            task's default template is used.
        :param fewshot_examples: Few-shot examples.
        :param generation_settings: Settings for structured generation.
        """
        self._pii_types = pii_types
        self._mask_placeholder = mask_placeholder

        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,
            generation_settings=generation_settings,
        )

    @override
    def _init_bridge(self, engine_type: EngineType) -> _TaskBridge:
        bridge_types: dict[EngineType, type[_TaskBridge]] = {
            EngineType.dspy: DSPyPIIMasking,
            EngineType.langchain: LangChainPIIMasking,
            EngineType.outlines: OutlinesPIIMasking,
        }

        try:
            return bridge_types[engine_type](
                task_id=self._task_id,
                prompt_instructions=self._custom_prompt_instructions,
                mask_placeholder=self._mask_placeholder,
                pii_types=self._pii_types,
                overwrite=self._overwrite,
            )
        except KeyError as err:
            raise KeyError(f"Engine type {engine_type} is not supported by {self.__class__.__name__}.") from err

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

    @property
    @override
    def _state(self) -> dict[str, Any]:
        return {
            **super()._state,
            "pii_types": self._pii_types,
            "mask_placeholder": self._mask_placeholder,
        }

    @override
    def to_hf_dataset(self, docs: Iterable[Doc], threshold: float = 0.5) -> 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_optimization_example(
        self, truth: dspy.Example, pred: dspy.Prediction, model: dspy.LM, trace: Any | None = None
    ) -> 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_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 
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
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 the task on a set of documents.

Parameters:

Name Type Description Default
docs Iterable[Doc]

Documents to process.

 required

Returns:

Type Description
Iterable[Doc]

Processed documents.
Source code in sieves/tasks/predictive/core.py 
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
def __call__(self, docs: Iterable[Doc]) -> Iterable[Doc]:
    """Execute the task on a set of documents.

    :param docs: Documents to process.
    :return Iterable[Doc]: Processed documents.
    """
    # 1. Compile expected prompt signatures.
    signature = self._bridge.prompt_signature

    # 2. Build executable.
    executable = self._engine.build_executable(
        inference_mode=self._bridge.inference_mode,
        prompt_template=self.prompt_template,
        prompt_signature=signature,
        fewshot_examples=self._fewshot_examples,
    )

    # Compute batch-wise results.
    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)]:
        if len(docs_batch) == 0:
            break

        # 3. Extract values from docs to inject/render those into prompt templates.
        docs_values = list(self._bridge.extract(docs_batch))
        assert len(docs_values) == len(docs_batch)

        # 4. Map extracted docs values onto chunks.
        docs_chunks_offsets: list[tuple[int, int]] = []
        docs_chunks: list[dict[str, Any]] = []
        for doc, doc_values in zip(docs_batch, docs_values):
            assert doc.text
            doc_chunks_values = [doc_values | {"text": chunk} for chunk in (doc.chunks or [doc.text])]
            docs_chunks_offsets.append((len(docs_chunks), len(docs_chunks) + len(doc_chunks_values)))
            docs_chunks.extend(doc_chunks_values)

        # 5. Execute prompts per chunk.
        results = list(executable(docs_chunks))
        assert len(results) == len(docs_chunks)

        # 6. Consolidate chunk results.
        results = list(self._bridge.consolidate(results, docs_chunks_offsets))
        assert len(results) == len(docs_batch)

        # 7. Integrate results into docs.
        docs_batch = self._bridge.integrate(results, docs_batch)

        yield from docs_batch

__init__(model, pii_types=None, mask_placeholder='[MASKED]', task_id=None, include_meta=True, batch_size=-1, overwrite=False, prompt_instructions=None, fewshot_examples=(), generation_settings=GenerationSettings())

Initialize PIIMasking task.

Parameters:

Name Type Description Default
model _TaskModel

Model to use.

 required
pii_types list[str] | None

Types of PII to mask. If None, all common PII types will be masked. E.g., ["NAME", "EMAIL", "PHONE", "ADDRESS", "SSN", "CREDIT_CARD", "DATE_OF_BIRTH"]

 None
mask_placeholder str

String to replace PII with.

 '[MASKED]'
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. If None, task's default template is used.

 None
fewshot_examples Sequence[FewshotExample]

Few-shot examples.

 ()
generation_settings GenerationSettings

Settings for structured generation.

 GenerationSettings()
Source code in sieves/tasks/predictive/pii_masking/core.py 
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
def __init__(
    self,
    model: _TaskModel,
    pii_types: list[str] | None = None,
    mask_placeholder: str = "[MASKED]",
    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] = (),
    generation_settings: GenerationSettings = GenerationSettings(),
) -> None:
    """
    Initialize PIIMasking task.

    :param model: Model to use.
    :param pii_types: Types of PII to mask. If None, all common PII types will be masked.
                     E.g., ["NAME", "EMAIL", "PHONE", "ADDRESS", "SSN", "CREDIT_CARD", "DATE_OF_BIRTH"]
    :param mask_placeholder: String to replace PII with.
    :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. If None,
        task's default template is used.
    :param fewshot_examples: Few-shot examples.
    :param generation_settings: Settings for structured generation.
    """
    self._pii_types = pii_types
    self._mask_placeholder = mask_placeholder

    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,
        generation_settings=generation_settings,
    )

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 
237
238
239
240
241
242
243
244
245
246
247
248
249
250
@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["generation_settings"] = GenerationSettings.model_validate(init_dict["generation_settings"])

    return cls(**init_dict)

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 
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
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]
    pred_eval = functools.partial(self._evaluate_optimization_example, 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(EngineType.get_engine_type(self._engine))

    return best_prompt, self._fewshot_examples

serialize()

Serialize task.

Returns:

Type Description
Config

Config instance.
Source code in sieves/tasks/core.py 
88
89
90
91
92
93
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: PIIBridge[PromptSignature, Result, InferenceMode]

DSPy bridge for PII masking.

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

    @override
    @property
    def _default_prompt_instructions(self) -> str:
        default_pii_types_desc = "all types of personally identifiable information"
        pii_types_desc = ", ".join(self._pii_types) if self._pii_types else default_pii_types_desc
        return (
            f"Identify and mask {pii_types_desc} in the given text. Replace each PII instance with "
            f"'{self._mask_placeholder}'."
        )

    @override
    @property
    def _prompt_example_template(self) -> str | None:
        return None

    @override
    @property
    def _prompt_conclusion(self) -> str | None:
        return None

    @override
    @cached_property
    def prompt_signature(self) -> type[dspy_.PromptSignature]:
        """Define prompt signature for DSPy."""
        PIIEntity = self._pii_entity_cls

        class PIIMasking(dspy.Signature):  # type: ignore[misc]
            text: str = dspy.InputField(description="Text to mask PII from.")
            reasoning: str = dspy.OutputField(description="Reasoning about what PII was found and masked.")
            masked_text: str = dspy.OutputField(description="Text with all PII masked.")
            pii_entities: list[PIIEntity] = dspy.OutputField(description="List of PII entities that were masked.")  # type: ignore[valid-type]

        PIIMasking.__doc__ = jinja2.Template(self._prompt_instructions).render()

        return PIIMasking

    @override
    @property
    def inference_mode(self) -> dspy_.InferenceMode:
        """Return inference mode for DSPy engine."""
        return dspy_.InferenceMode.chain_of_thought

    @override
    def integrate(self, results: Iterable[dspy_.Result], docs: Iterable[Doc]) -> Iterable[Doc]:
        """Integrate results into docs."""
        for doc, result in zip(docs, results):
            # Store masked text and PII entities in results
            doc.results[self._task_id] = {
                "masked_text": result.masked_text,
                "pii_entities": result.pii_entities,
            }

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

        return docs

    @override
    def consolidate(
        self, results: Iterable[dspy_.Result], docs_offsets: list[tuple[int, int]]
    ) -> Iterable[dspy_.Result]:
        """Consolidate results from multiple chunks."""
        results = list(results)
        PIIEntity = self._pii_entity_cls

        # Merge results for each document
        for doc_offset in docs_offsets:
            doc_results = results[doc_offset[0] : doc_offset[1]]
            seen_entities: set[PIIEntity] = set()  # type: ignore[valid-type]
            entities: list[PIIEntity] = []  # type: ignore[valid-type]
            masked_texts: list[str] = []
            reasonings: list[str] = []

            for res in doc_results:
                reasonings.append(res.reasoning)
                masked_texts.append(res.masked_text)
                for entity in res.pii_entities:
                    if entity not in seen_entities:
                        entities.extend(res.pii_entities)
                        seen_entities.add(entity)

            yield dspy.Prediction.from_completions(
                {"masked_text": [" ".join(masked_texts)], "pii_entities": [entities], "reasoning": [str(reasonings)]},
                signature=self.prompt_signature,
            )

inference_mode property

Return inference mode for DSPy engine.

prompt_signature cached property

Define prompt signature for DSPy.

prompt_template property

Return prompt template.

Chains _prompt_instructions, _prompt_example_template and _prompt_conclusion.

Note: different engines have different expectations as 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 engine-specific expectations when creating a prompt template.

Returns:

Type Description
str

Prompt template as string. None if not used by engine.

__init__(task_id, prompt_instructions, overwrite, mask_placeholder, pii_types)

Initialize PIIBridge.

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
overwrite bool

Whether to overwrite text with masked text.

 required
mask_placeholder str

String to replace PII with.

 required
pii_types list[str] | None

Types of PII to mask. If None, all common PII types will be masked.

 required
Source code in sieves/tasks/predictive/pii_masking/bridges.py 
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
def __init__(
    self,
    task_id: str,
    prompt_instructions: str | None,
    overwrite: bool,
    mask_placeholder: str,
    pii_types: list[str] | None,
):
    """
    Initialize PIIBridge.

    :param task_id: Task ID.
    :param prompt_instructions: Custom prompt instructions. If None, default instructions are used.
    :param overwrite: Whether to overwrite text with masked text.
    :param mask_placeholder: String to replace PII with.
    :param pii_types: Types of PII to mask. If None, all common PII types will be masked.
    """
    super().__init__(
        task_id=task_id,
        prompt_instructions=prompt_instructions,
        overwrite=overwrite,
    )
    self._mask_placeholder = mask_placeholder
    self._pii_types = pii_types
    self._pii_entity_cls = self._create_pii_entity_cls()

consolidate(results, docs_offsets)

Consolidate results from multiple chunks.

Source code in sieves/tasks/predictive/pii_masking/bridges.py 
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
@override
def consolidate(
    self, results: Iterable[dspy_.Result], docs_offsets: list[tuple[int, int]]
) -> Iterable[dspy_.Result]:
    """Consolidate results from multiple chunks."""
    results = list(results)
    PIIEntity = self._pii_entity_cls

    # Merge results for each document
    for doc_offset in docs_offsets:
        doc_results = results[doc_offset[0] : doc_offset[1]]
        seen_entities: set[PIIEntity] = set()  # type: ignore[valid-type]
        entities: list[PIIEntity] = []  # type: ignore[valid-type]
        masked_texts: list[str] = []
        reasonings: list[str] = []

        for res in doc_results:
            reasonings.append(res.reasoning)
            masked_texts.append(res.masked_text)
            for entity in res.pii_entities:
                if entity not in seen_entities:
                    entities.extend(res.pii_entities)
                    seen_entities.add(entity)

        yield dspy.Prediction.from_completions(
            {"masked_text": [" ".join(masked_texts)], "pii_entities": [entities], "reasoning": [str(reasonings)]},
            signature=self.prompt_signature,
        )

extract(docs)

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

Parameters:

Name Type Description Default
docs Iterable[Doc]

Docs to extract values from.

 required

Returns:

Type Description
Iterable[dict[str, Any]]

All values from doc instances that are to be injected into the prompts
Source code in sieves/tasks/predictive/bridges.py 
110
111
112
113
114
115
116
def extract(self, docs: Iterable[Doc]) -> Iterable[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 Iterable[dict[str, Any]]: All values from doc instances that are to be injected into the prompts
    """
    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 
111
112
113
114
115
116
117
118
119
120
121
122
123
124
@override
def integrate(self, results: Iterable[dspy_.Result], docs: Iterable[Doc]) -> Iterable[Doc]:
    """Integrate results into docs."""
    for doc, result in zip(docs, results):
        # Store masked text and PII entities in results
        doc.results[self._task_id] = {
            "masked_text": result.masked_text,
            "pii_entities": result.pii_entities,
        }

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

    return docs

LangChainPIIMasking

Bases: PydanticBasedPIIMasking[InferenceMode]

LangChain bridge for PII masking.

Source code in sieves/tasks/predictive/pii_masking/bridges.py 
274
275
276
277
278
279
280
class LangChainPIIMasking(PydanticBasedPIIMasking[langchain_.InferenceMode]):
    """LangChain bridge for PII masking."""

    @override
    @property
    def inference_mode(self) -> langchain_.InferenceMode:
        return langchain_.InferenceMode.structured

prompt_signature cached property

Define prompt signature for Pydantic-based engines.

prompt_template property

Return prompt template.

Chains _prompt_instructions, _prompt_example_template and _prompt_conclusion.

Note: different engines have different expectations as 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 engine-specific expectations when creating a prompt template.

Returns:

Type Description
str

Prompt template as string. None if not used by engine.

__init__(task_id, prompt_instructions, overwrite, mask_placeholder, pii_types)

Initialize PIIBridge.

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
overwrite bool

Whether to overwrite text with masked text.

 required
mask_placeholder str

String to replace PII with.

 required
pii_types list[str] | None

Types of PII to mask. If None, all common PII types will be masked.

 required
Source code in sieves/tasks/predictive/pii_masking/bridges.py 
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
def __init__(
    self,
    task_id: str,
    prompt_instructions: str | None,
    overwrite: bool,
    mask_placeholder: str,
    pii_types: list[str] | None,
):
    """
    Initialize PIIBridge.

    :param task_id: Task ID.
    :param prompt_instructions: Custom prompt instructions. If None, default instructions are used.
    :param overwrite: Whether to overwrite text with masked text.
    :param mask_placeholder: String to replace PII with.
    :param pii_types: Types of PII to mask. If None, all common PII types will be masked.
    """
    super().__init__(
        task_id=task_id,
        prompt_instructions=prompt_instructions,
        overwrite=overwrite,
    )
    self._mask_placeholder = mask_placeholder
    self._pii_types = pii_types
    self._pii_entity_cls = self._create_pii_entity_cls()

extract(docs)

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

Parameters:

Name Type Description Default
docs Iterable[Doc]

Docs to extract values from.

 required

Returns:

Type Description
Iterable[dict[str, Any]]

All values from doc instances that are to be injected into the prompts
Source code in sieves/tasks/predictive/bridges.py 
110
111
112
113
114
115
116
def extract(self, docs: Iterable[Doc]) -> Iterable[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 Iterable[dict[str, Any]]: All values from doc instances that are to be injected into the prompts
    """
    return ({"text": doc.text if doc.text else None} for doc in docs)

OutlinesPIIMasking

Bases: PydanticBasedPIIMasking[InferenceMode]

Outlines bridge for PII masking.

Source code in sieves/tasks/predictive/pii_masking/bridges.py 
265
266
267
268
269
270
271
class OutlinesPIIMasking(PydanticBasedPIIMasking[outlines_.InferenceMode]):
    """Outlines bridge for PII masking."""

    @override
    @property
    def inference_mode(self) -> outlines_.InferenceMode:
        return outlines_.InferenceMode.json

prompt_signature cached property

Define prompt signature for Pydantic-based engines.

prompt_template property

Return prompt template.

Chains _prompt_instructions, _prompt_example_template and _prompt_conclusion.

Note: different engines have different expectations as 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 engine-specific expectations when creating a prompt template.

Returns:

Type Description
str

Prompt template as string. None if not used by engine.

__init__(task_id, prompt_instructions, overwrite, mask_placeholder, pii_types)

Initialize PIIBridge.

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
overwrite bool

Whether to overwrite text with masked text.

 required
mask_placeholder str

String to replace PII with.

 required
pii_types list[str] | None

Types of PII to mask. If None, all common PII types will be masked.

 required
Source code in sieves/tasks/predictive/pii_masking/bridges.py 
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
def __init__(
    self,
    task_id: str,
    prompt_instructions: str | None,
    overwrite: bool,
    mask_placeholder: str,
    pii_types: list[str] | None,
):
    """
    Initialize PIIBridge.

    :param task_id: Task ID.
    :param prompt_instructions: Custom prompt instructions. If None, default instructions are used.
    :param overwrite: Whether to overwrite text with masked text.
    :param mask_placeholder: String to replace PII with.
    :param pii_types: Types of PII to mask. If None, all common PII types will be masked.
    """
    super().__init__(
        task_id=task_id,
        prompt_instructions=prompt_instructions,
        overwrite=overwrite,
    )
    self._mask_placeholder = mask_placeholder
    self._pii_types = pii_types
    self._pii_entity_cls = self._create_pii_entity_cls()

extract(docs)

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

Parameters:

Name Type Description Default
docs Iterable[Doc]

Docs to extract values from.

 required

Returns:

Type Description
Iterable[dict[str, Any]]

All values from doc instances that are to be injected into the prompts
Source code in sieves/tasks/predictive/bridges.py 
110
111
112
113
114
115
116
def extract(self, docs: Iterable[Doc]) -> Iterable[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 Iterable[dict[str, Any]]: All values from doc instances that are to be injected into the prompts
    """
    return ({"text": doc.text if doc.text else None} for doc in docs)

PIIBridge

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

Abstract base class for PII masking bridges.

Source code in sieves/tasks/predictive/pii_masking/bridges.py 
20
21
22
23
24
25
26
27
28
29
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
class PIIBridge(Bridge[_BridgePromptSignature, _BridgeResult, EngineInferenceMode], abc.ABC):
    """Abstract base class for PII masking bridges."""

    def __init__(
        self,
        task_id: str,
        prompt_instructions: str | None,
        overwrite: bool,
        mask_placeholder: str,
        pii_types: list[str] | None,
    ):
        """
        Initialize PIIBridge.

        :param task_id: Task ID.
        :param prompt_instructions: Custom prompt instructions. If None, default instructions are used.
        :param overwrite: Whether to overwrite text with masked text.
        :param mask_placeholder: String to replace PII with.
        :param pii_types: Types of PII to mask. If None, all common PII types will be masked.
        """
        super().__init__(
            task_id=task_id,
            prompt_instructions=prompt_instructions,
            overwrite=overwrite,
        )
        self._mask_placeholder = mask_placeholder
        self._pii_types = pii_types
        self._pii_entity_cls = self._create_pii_entity_cls()

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

        :returns: PII entity class.
        """
        pii_types = self._pii_types
        PIIType = Literal[*pii_types] if pii_types else str

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

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

        return PIIEntity

inference_mode abstractmethod property

Return inference mode.

Returns:

Type Description
EngineInferenceMode

Inference mode.

prompt_signature abstractmethod property

Create output signature.

E.g.: Signature in DSPy, Pydantic objects in outlines, JSON schema in jsonformers. This is engine-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_template and _prompt_conclusion.

Note: different engines have different expectations as 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 engine-specific expectations when creating a prompt template.

Returns:

Type Description
str

Prompt template as string. None if not used by engine.

__init__(task_id, prompt_instructions, overwrite, mask_placeholder, pii_types)

Initialize PIIBridge.

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
overwrite bool

Whether to overwrite text with masked text.

 required
mask_placeholder str

String to replace PII with.

 required
pii_types list[str] | None

Types of PII to mask. If None, all common PII types will be masked.

 required
Source code in sieves/tasks/predictive/pii_masking/bridges.py 
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
def __init__(
    self,
    task_id: str,
    prompt_instructions: str | None,
    overwrite: bool,
    mask_placeholder: str,
    pii_types: list[str] | None,
):
    """
    Initialize PIIBridge.

    :param task_id: Task ID.
    :param prompt_instructions: Custom prompt instructions. If None, default instructions are used.
    :param overwrite: Whether to overwrite text with masked text.
    :param mask_placeholder: String to replace PII with.
    :param pii_types: Types of PII to mask. If None, all common PII types will be masked.
    """
    super().__init__(
        task_id=task_id,
        prompt_instructions=prompt_instructions,
        overwrite=overwrite,
    )
    self._mask_placeholder = mask_placeholder
    self._pii_types = pii_types
    self._pii_entity_cls = self._create_pii_entity_cls()

consolidate(results, docs_offsets) abstractmethod

Consolidate results for document chunks into document results.

Parameters:

Name Type Description Default
results Iterable[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
Iterable[TaskResult]

Results per document.
Source code in sieves/tasks/predictive/bridges.py 
127
128
129
130
131
132
133
134
135
@abc.abstractmethod
def consolidate(self, results: Iterable[TaskResult], docs_offsets: list[tuple[int, int]]) -> Iterable[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 Iterable[_TaskResult]: Results per document.
    """

extract(docs)

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

Parameters:

Name Type Description Default
docs Iterable[Doc]

Docs to extract values from.

 required

Returns:

Type Description
Iterable[dict[str, Any]]

All values from doc instances that are to be injected into the prompts
Source code in sieves/tasks/predictive/bridges.py 
110
111
112
113
114
115
116
def extract(self, docs: Iterable[Doc]) -> Iterable[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 Iterable[dict[str, Any]]: All values from doc instances that are to be injected into the prompts
    """
    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 Iterable[TaskResult]

Results from prompt executable.

 required
docs Iterable[Doc]

Doc instances to update.

 required

Returns:

Type Description
Iterable[Doc]

Updated doc instances.
Source code in sieves/tasks/predictive/bridges.py 
118
119
120
121
122
123
124
125
@abc.abstractmethod
def integrate(self, results: Iterable[TaskResult], docs: Iterable[Doc]) -> Iterable[Doc]:
    """Integrate results into Doc instances.

    :param results: Results from prompt executable.
    :param docs: Doc instances to update.
    :return Iterable[Doc]: Updated doc instances.
    """

PydanticBasedPIIMasking

Bases: PIIBridge[BaseModel, BaseModel, EngineInferenceMode], ABC

Base class for Pydantic-based PII masking bridges.

Source code in sieves/tasks/predictive/pii_masking/bridges.py 
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
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
class PydanticBasedPIIMasking(PIIBridge[pydantic.BaseModel, pydantic.BaseModel, EngineInferenceMode], abc.ABC):
    """Base class for Pydantic-based PII masking bridges."""

    @property
    def _default_prompt_instructions(self) -> str:
        return """
        Identify and mask Personally Identifiable Information (PII) in the given text.
        {% if pii_types|length > 0 -%}
            Focus on these specific PII types: {{ pii_types|join(', ') }}.
        {% else -%}
            Mask all common types of PII such as names, addresses, phone numbers, emails, SSNs, credit
            card numbers, etc.
        {% endif -%}
        Replace each instance of PII with "{{ mask_placeholder }}".
        """

    @override
    @property
    def _prompt_example_template(self) -> str | None:
        return """
        {% if examples|length > 0 -%}
            <examples>
            {%- for example in examples %}
                <example>
                    <text>"{{ example.text }}"</text>
                    <output>
                        <reasoning>{{ example.reasoning }}</reasoning>
                        <masked_test>{{ example.masked_text }}</masked_test>
                        <pii_entities_found>{{ example.pii_entities }}</pii_entities_found>
                    </output>
                </example>
            {% endfor -%}
            </examples>
        {% endif -%}
        """

    @override
    @property
    def _prompt_conclusion(self) -> str | None:
        return """
        ========

        <text>{{ text }}</text>
        <output>
        """

    @override
    @cached_property
    def prompt_signature(self) -> type[pydantic.BaseModel]:
        """Define prompt signature for Pydantic-based engines."""
        PIIEntity = self._pii_entity_cls

        class PIIMasking(pydantic.BaseModel, frozen=True):
            """PII masking output."""

            reasoning: str
            masked_text: str
            pii_entities: list[PIIEntity]  # type: ignore[valid-type]

        return PIIMasking

    @override
    def integrate(self, results: Iterable[pydantic.BaseModel], docs: Iterable[Doc]) -> Iterable[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] = {"masked_text": result.masked_text, "pii_entities": result.pii_entities}

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

        return docs

    @override
    def consolidate(
        self, results: Iterable[pydantic.BaseModel], docs_offsets: list[tuple[int, int]]
    ) -> Iterable[pydantic.BaseModel]:
        results = list(results)
        PIIEntity = self._pii_entity_cls

        # Merge results for each document
        for doc_offset in docs_offsets:
            doc_results = results[doc_offset[0] : doc_offset[1]]
            seen_entities: set[PIIEntity] = set()  # type: ignore[valid-type]
            entities: list[PIIEntity] = []  # type: ignore[valid-type]
            masked_texts: list[str] = []
            reasonings: list[str] = []

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

                assert hasattr(res, "reasoning")
                assert hasattr(res, "masked_text")
                assert hasattr(res, "pii_entities")

                reasonings.append(res.reasoning)
                masked_texts.append(res.masked_text)
                for entity in res.pii_entities:
                    if entity not in seen_entities:
                        entities.extend(res.pii_entities)
                        seen_entities.add(entity)

            yield self.prompt_signature(
                reasoning=str(reasonings), masked_text=" ".join(masked_texts), pii_entities=entities
            )

inference_mode abstractmethod property

Return inference mode.

Returns:

Type Description
EngineInferenceMode

Inference mode.

prompt_signature cached property

Define prompt signature for Pydantic-based engines.

prompt_template property

Return prompt template.

Chains _prompt_instructions, _prompt_example_template and _prompt_conclusion.

Note: different engines have different expectations as 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 engine-specific expectations when creating a prompt template.

Returns:

Type Description
str

Prompt template as string. None if not used by engine.

__init__(task_id, prompt_instructions, overwrite, mask_placeholder, pii_types)

Initialize PIIBridge.

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
overwrite bool

Whether to overwrite text with masked text.

 required
mask_placeholder str

String to replace PII with.

 required
pii_types list[str] | None

Types of PII to mask. If None, all common PII types will be masked.

 required
Source code in sieves/tasks/predictive/pii_masking/bridges.py 
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
def __init__(
    self,
    task_id: str,
    prompt_instructions: str | None,
    overwrite: bool,
    mask_placeholder: str,
    pii_types: list[str] | None,
):
    """
    Initialize PIIBridge.

    :param task_id: Task ID.
    :param prompt_instructions: Custom prompt instructions. If None, default instructions are used.
    :param overwrite: Whether to overwrite text with masked text.
    :param mask_placeholder: String to replace PII with.
    :param pii_types: Types of PII to mask. If None, all common PII types will be masked.
    """
    super().__init__(
        task_id=task_id,
        prompt_instructions=prompt_instructions,
        overwrite=overwrite,
    )
    self._mask_placeholder = mask_placeholder
    self._pii_types = pii_types
    self._pii_entity_cls = self._create_pii_entity_cls()

extract(docs)

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

Parameters:

Name Type Description Default
docs Iterable[Doc]

Docs to extract values from.

 required

Returns:

Type Description
Iterable[dict[str, Any]]

All values from doc instances that are to be injected into the prompts
Source code in sieves/tasks/predictive/bridges.py 
110
111
112
113
114
115
116
def extract(self, docs: Iterable[Doc]) -> Iterable[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 Iterable[dict[str, Any]]: All values from doc instances that are to be injected into the prompts
    """
    return ({"text": doc.text if doc.text else None} for doc in docs)