Skip to content

OutMask module

OutMask

Bases: Module

A module to remove specific fields of the given data models

Example:

import synalinks
import asyncio

language_model = synalinks.LanguageModel(
    model="ollama/mistral",
)

class Document(synalinks.DataModel):
    title: str = synalinks.Field(
        description="The title of the document",
    )
    text: str = synalinks.Field(
        description="The content of the document",
    )

class Summary(synalinks.DataModel):
    summary: str = synalinks.Field(
        description="the concise summary of the document",
    )

async def main():
    inputs = Input(data_model=Document)
    summary = synalinks.ChainOfThought(
        data_model=Summary,
        language_model=language_model,
    )(inputs)
    masked_summary = synalinks.OutMask(
        # remove the thinking field from the chain of thought
        mask=["thinking"],
    )(summary)

    program = Program(
        inputs=inputs,
        outputs=masked_summary,
        name="summary_generator",
        description="Generate a summary from a document",
    )

Parameters:

Name Type Description Default
mask list

The list of keys to remove.

 None
name str

Optional. The name of the module.

 None
description str

Optional. The description of the module.

 None
trainable bool

Whether the module's variables should be trainable.

 False
Source code in synalinks/src/modules/masking/out_mask.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
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
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
@synalinks_export(
    [
        "synalinks.OutMask",
        "synalinks.modules.OutMask",
    ]
)
class OutMask(Module):
    """A module to remove specific fields of the given data models

    Example:

    ```python
    import synalinks
    import asyncio

    language_model = synalinks.LanguageModel(
        model="ollama/mistral",
    )

    class Document(synalinks.DataModel):
        title: str = synalinks.Field(
            description="The title of the document",
        )
        text: str = synalinks.Field(
            description="The content of the document",
        )

    class Summary(synalinks.DataModel):
        summary: str = synalinks.Field(
            description="the concise summary of the document",
        )

    async def main():
        inputs = Input(data_model=Document)
        summary = synalinks.ChainOfThought(
            data_model=Summary,
            language_model=language_model,
        )(inputs)
        masked_summary = synalinks.OutMask(
            # remove the thinking field from the chain of thought
            mask=["thinking"],
        )(summary)

        program = Program(
            inputs=inputs,
            outputs=masked_summary,
            name="summary_generator",
            description="Generate a summary from a document",
        )
    ```

    Args:
        mask (list): The list of keys to remove.
        name (str): Optional. The name of the module.
        description (str): Optional. The description of the module.
        trainable (bool): Whether the module's variables should be trainable.
    """

    def __init__(
        self,
        mask=None,
        name=None,
        description=None,
        trainable=False,
    ):
        if not mask or not isinstance(mask, list):
            raise ValueError("`mask` parameter should be a list of fields to remove")
        super().__init__(
            name=name,
            description=description,
        )
        self.mask = mask

    async def call(self, inputs):
        outputs = tree.map_structure(
            lambda x: x.out_mask(mask=self.mask),
            inputs,
        )
        return outputs