Skip to content

Model training

Config File

The input config file for training a SAFE model is very similar to the GPT2 config file, with the addition of an optional num_labels attribute for training with descriptors regularization.

{
  "activation_function": "gelu_new",
  "attn_pdrop": 0.1,
  "bos_token_id": 10000,
  "embd_pdrop": 0.1,
  "eos_token_id": 1,
  "initializer_range": 0.02,
  "layer_norm_epsilon": 1e-05,
  "model_type": "gpt2",
  "n_embd": 768,
  "n_head": 12,
  "n_inner": null,
  "n_layer": 12,
  "n_positions": 1024,
  "reorder_and_upcast_attn": false,
  "resid_pdrop": 0.1,
  "scale_attn_by_inverse_layer_idx": false,
  "scale_attn_weights": true,
  "summary_activation": "tanh",
  "summary_first_dropout": 0.1,
  "summary_proj_to_labels": true,
  "summary_type": "cls_index",
  "summary_hidden_size": 128,
  "summary_use_proj": true,
  "transformers_version": "5.16.1",
  "use_cache": true,
  "vocab_size": 10000,
  "num_labels": 9
}

SAFE Model

PropertyHead(config: PretrainedConfig)

Bases: Module

Compute a single vector summary of a sequence hidden states.

Parameters:

  • config ([`PretrainedConfig`]) –

    The config used by the model. Relevant arguments in the config class of the model are (refer to the actual config class of your model for the default values it uses):

    • summary_type (str) -- The method to use to make this summary. Accepted values are:

      - "last" -- Take the last token hidden state (like XLNet) - "first" -- Take the first token hidden state (like Bert) - "mean" -- Take the mean of all tokens hidden states - "cls_index" -- Supply a Tensor of classification token position (GPT/GPT-2)

    • summary_activation (Optional[str]) -- Set to "tanh" to add a tanh activation to the output, another string, or None to add no activation.
Source code in safe/trainer/model.py 
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
def __init__(self, config: PretrainedConfig):
    super().__init__()

    self.summary_type = getattr(config, "summary_type", "cls_index")
    self.summary = torch.nn.Identity()
    last_hidden_size = config.hidden_size

    if getattr(config, "summary_hidden_size", None) and config.summary_hidden_size > 0:
        self.summary = nn.Linear(config.hidden_size, config.summary_hidden_size)
        last_hidden_size = config.summary_hidden_size

    activation_string = getattr(config, "summary_activation", None)
    self.activation: Callable = (
        get_activation(activation_string) if activation_string else nn.Identity()
    )

    self.out = torch.nn.Identity()
    if getattr(config, "num_labels", None) and config.num_labels > 0:
        num_labels = config.num_labels
        self.out = nn.Linear(last_hidden_size, num_labels)

activation: Callable = get_activation(activation_string) if activation_string else nn.Identity() instance-attribute

out = torch.nn.Identity() instance-attribute

summary = torch.nn.Identity() instance-attribute

summary_type = getattr(config, 'summary_type', 'cls_index') instance-attribute

forward(hidden_states: torch.FloatTensor, cls_index: Optional[torch.LongTensor] = None) -> torch.FloatTensor

Compute a single vector summary of a sequence hidden states.

Parameters:

  • hidden_states (FloatTensor) –

    torch.FloatTensor of shape [batch_size, seq_len, hidden_size]) The hidden states of the last layer.

  • cls_index (Optional[LongTensor], default: None ) –

    torch.LongTensor of shape [batch_size] or [batch_size, ...] where ... are optional leading dimensions of hidden_states, optional Used if summary_type == "cls_index" and takes the last token of the sequence as classification token.

Returns:

  • FloatTensor –

    torch.FloatTensor: The summary of the sequence hidden states.

Source code in safe/trainer/model.py 
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
def forward(
    self,
    hidden_states: torch.FloatTensor,
    cls_index: Optional[torch.LongTensor] = None,
) -> torch.FloatTensor:
    """
    Compute a single vector summary of a sequence hidden states.

    Args:
        hidden_states: `torch.FloatTensor` of shape `[batch_size, seq_len, hidden_size]`)
            The hidden states of the last layer.
        cls_index: `torch.LongTensor` of shape `[batch_size]` or `[batch_size, ...]`
            where ... are optional leading dimensions of `hidden_states`, *optional*
            Used if `summary_type == "cls_index"` and takes the last token of the sequence as classification token.

    Returns:
        `torch.FloatTensor`: The summary of the sequence hidden states.
    """
    if self.summary_type == "last":
        output = hidden_states[:, -1]
    elif self.summary_type == "first":
        output = hidden_states[:, 0]
    elif self.summary_type == "mean":
        output = hidden_states.mean(dim=1)
    elif self.summary_type == "cls_index":
        batch_size = hidden_states.shape[0]
        if cls_index is None:
            cls_index = torch.full(
                (batch_size,),
                hidden_states.shape[1] - 1,
                device=hidden_states.device,
                dtype=torch.long,
            )
        output = hidden_states[torch.arange(batch_size, device=hidden_states.device), cls_index]
    else:
        raise NotImplementedError

    output = self.summary(output)
    output = self.activation(output)
    return self.out(output)

SAFEDoubleHeadsModel(config)

Bases: GPT2DoubleHeadsModel

The safe model is a dual head GPT2 model with a language modeling head and an optional multi-task regression head

Source code in safe/trainer/model.py 
 99
100
101
102
103
104
def __init__(self, config):
    self.num_labels = getattr(config, "num_labels", None)
    super().__init__(config)
    self.config.num_labels = self.num_labels
    del self.multiple_choice_head
    self.multiple_choice_head = PropertyHead(config)

multiple_choice_head = PropertyHead(config) instance-attribute

num_labels = getattr(config, 'num_labels', None) instance-attribute

forward(input_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Any] = None, attention_mask: Optional[torch.FloatTensor] = None, token_type_ids: Optional[torch.LongTensor] = None, position_ids: Optional[torch.LongTensor] = None, head_mask: Optional[torch.FloatTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, mc_token_ids: Optional[torch.LongTensor] = None, labels: Optional[torch.LongTensor] = None, mc_labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, inputs: Optional[Any] = None, encoder_hidden_states: Optional[torch.Tensor] = None, encoder_attention_mask: Optional[torch.Tensor] = None, **kwargs) -> Union[Tuple, GPT2DoubleHeadsModelOutput]

Parameters:

  • mc_token_ids (`torch.LongTensor` of shape `(batch_size, num_choices)`, *optional*, default to index of the last token of the input, default: None ) –

    Index of the classification token in each input sequence. Selected in the range [0, input_ids.size(-1) - 1].

  • labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*, default: None ) –

    Labels for language modeling. Note that the labels are shifted inside the model, i.e. you can set labels = input_ids. Indices are selected in [-100, 0, ..., config.vocab_size - 1]. All labels set to -100 are ignored (masked), the loss is only computed for labels in [0, ..., config.vocab_size - 1]

  • mc_labels (`torch.LongTensor` of shape `(batch_size, n_tasks)`, *optional*, default: None ) –

    Labels for computing the supervized loss for regularization.

  • inputs (Optional[Any], default: None ) –

    List of inputs, put here because the trainer removes information not in signature

Source code in safe/trainer/model.py 
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
def forward(
    self,
    input_ids: Optional[torch.LongTensor] = None,
    past_key_values: Optional[Any] = None,
    attention_mask: Optional[torch.FloatTensor] = None,
    token_type_ids: Optional[torch.LongTensor] = None,
    position_ids: Optional[torch.LongTensor] = None,
    head_mask: Optional[torch.FloatTensor] = None,
    inputs_embeds: Optional[torch.FloatTensor] = None,
    mc_token_ids: Optional[torch.LongTensor] = None,
    labels: Optional[torch.LongTensor] = None,
    mc_labels: Optional[torch.LongTensor] = None,
    use_cache: Optional[bool] = None,
    output_attentions: Optional[bool] = None,
    output_hidden_states: Optional[bool] = None,
    return_dict: Optional[bool] = None,
    inputs: Optional[Any] = None,  # do not remove because of trainer
    encoder_hidden_states: Optional[torch.Tensor] = None,
    encoder_attention_mask: Optional[torch.Tensor] = None,
    **kwargs,
) -> Union[Tuple, GPT2DoubleHeadsModelOutput]:
    r"""

    Args:
        mc_token_ids (`torch.LongTensor` of shape `(batch_size, num_choices)`, *optional*, default to index of the last token of the input):
            Index of the classification token in each input sequence. Selected in the range `[0, input_ids.size(-1) -
            1]`.
        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
            `labels = input_ids`. Indices are selected in `[-100, 0, ..., config.vocab_size - 1]`. All labels set to
            `-100` are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size - 1]`
        mc_labels (`torch.LongTensor` of shape `(batch_size, n_tasks)`, *optional*):
            Labels for computing the supervized loss for regularization.
        inputs: List of inputs, put here because the trainer removes information not in signature
    """
    del inputs
    return_dict = return_dict if return_dict is not None else self.config.return_dict
    transformer_outputs = self.transformer(
        input_ids,
        past_key_values=past_key_values,
        attention_mask=attention_mask,
        token_type_ids=token_type_ids,
        position_ids=position_ids,
        head_mask=head_mask,
        inputs_embeds=inputs_embeds,
        use_cache=use_cache,
        output_attentions=output_attentions,
        output_hidden_states=output_hidden_states,
        return_dict=return_dict,
        encoder_hidden_states=encoder_hidden_states,
        encoder_attention_mask=encoder_attention_mask,
        **kwargs,
    )

    hidden_states = transformer_outputs[0]
    lm_logits = self.lm_head(hidden_states)

    if mc_token_ids is None and self.config.pad_token_id is not None and input_ids is not None:
        mc_token_ids = (torch.ne(input_ids, self.config.pad_token_id).sum(-1) - 1).to(
            lm_logits.device
        )

    # Set device for model parallelism
    if getattr(self, "model_parallel", False):
        torch.cuda.set_device(self.transformer.first_device)
        hidden_states = hidden_states.to(self.lm_head.weight.device)

    mc_loss = None
    mc_logits = None
    if mc_labels is not None and getattr(self.config, "num_labels", 0) > 0:
        mc_logits = self.multiple_choice_head(hidden_states, mc_token_ids).squeeze(-1)
        mc_labels = mc_labels.to(mc_logits.device)
        loss_fct = MSELoss()
        mc_loss = loss_fct(
            mc_logits.view(-1, mc_logits.size(-1)), mc_labels.view(-1, mc_logits.size(-1))
        )

    lm_loss = None
    if labels is not None:
        labels = labels.to(lm_logits.device)
        shift_logits = lm_logits[..., :-1, :].contiguous()
        shift_labels = labels[..., 1:].contiguous()
        loss_fct = CrossEntropyLoss()
        lm_loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))

    if not return_dict:
        output = (lm_logits, mc_logits) + transformer_outputs[1:]
        return (
            lm_loss,
            mc_loss,
        ) + output

    return GPT2DoubleHeadsModelOutput(
        loss=lm_loss,
        mc_loss=mc_loss,
        logits=lm_logits,
        mc_logits=mc_logits,
        past_key_values=transformer_outputs.past_key_values,
        hidden_states=transformer_outputs.hidden_states,
        attentions=transformer_outputs.attentions,
    )

Trainer

SAFETrainer(*args, prop_loss_coeff: float = 0.001, **kwargs)

Bases: Trainer

Custom trainer for training SAFE model.

This custom trainer changes the loss function to support the property head

Source code in safe/trainer/trainer_utils.py 
12
13
14
def __init__(self, *args, prop_loss_coeff: float = 1e-3, **kwargs):
    super().__init__(*args, **kwargs)
    self.prop_loss_coeff = prop_loss_coeff

prop_loss_coeff = prop_loss_coeff instance-attribute

compute_loss(model, inputs, return_outputs=False, num_items_in_batch=None)

How the loss is computed by Trainer. By default, all models return the loss in the first element.

Source code in safe/trainer/trainer_utils.py 
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
def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
    """
    How the loss is computed by Trainer. By default, all models return the loss in the first element.
    """
    del num_items_in_batch
    labels = (
        inputs.pop("labels") if self.label_smoother is not None and "labels" in inputs else None
    )
    outputs = model(**inputs)
    # Preserve the past state when the Trainer is configured to retain it.
    past_index = getattr(self.args, "past_index", -1)
    if past_index >= 0:
        self._past = outputs[past_index]

    if labels is not None:
        loss = self.label_smoother(outputs, labels, shift_labels=True)
    else:
        if isinstance(outputs, dict) and "loss" not in outputs:
            raise ValueError(
                "The model did not return a loss from the inputs, only the following keys: "
                f"{','.join(outputs.keys())}. For reference, the inputs it received are {','.join(inputs.keys())}."
            )
        # We don't use .loss here since the model may return tuples instead of ModelOutput.
        loss = outputs["loss"] if isinstance(outputs, dict) else outputs[0]

    mc_loss = outputs.get("mc_loss", None) if isinstance(outputs, dict) else outputs[1]
    if mc_loss is not None:
        loss = loss + self.prop_loss_coeff * mc_loss
    return (loss, outputs) if return_outputs else loss

Data Collator

SAFECollator(tokenizer: Tokenizer, pad_to_multiple_of: Optional[int] = None, input_key: str = 'inputs', label_key: str = 'labels', property_key: str = 'descriptors', include_descriptors: bool = False, max_length: Optional[int] = None)

Collate function for language modelling tasks

Note

The collate function is based on the default DataCollatorForLanguageModeling in huggingface see: https://github.com/huggingface/transformers/blob/v4.19.2/src/transformers/data/data_collator.py

Default collator for huggingface transformers in izanagi.

Parameters:

  • tokenizer (Tokenizer) –

    Huggingface tokenizer

  • input_key (str, default: 'inputs' ) –

    key to use for input ids

  • label_key (str, default: 'labels' ) –

    key to use for labels

  • property_key (str, default: 'descriptors' ) –

    key to use for properties

  • include_descriptors (bool, default: False ) –

    whether to include training on descriptors or not

  • pad_to_multiple_of (Optional[int], default: None ) –

    pad to multiple of this value

Source code in safe/trainer/collator.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
def __init__(
    self,
    tokenizer: Tokenizer,
    pad_to_multiple_of: Optional[int] = None,
    input_key: str = "inputs",
    label_key: str = "labels",
    property_key: str = "descriptors",
    include_descriptors: bool = False,
    max_length: Optional[int] = None,
):
    """
    Default collator for huggingface transformers in izanagi.

    Args:
        tokenizer: Huggingface tokenizer
        input_key: key to use for input ids
        label_key: key to use for labels
        property_key: key to use for properties
        include_descriptors: whether to include training on descriptors or not
        pad_to_multiple_of: pad to multiple of this value
    """

    self.tokenizer = tokenizer
    self.pad_to_multiple_of = pad_to_multiple_of
    self.input_key = input_key
    self.label_key = label_key
    self.property_key = property_key
    self.include_descriptors = include_descriptors
    self.max_length = max_length

include_descriptors = include_descriptors instance-attribute

input_key = input_key instance-attribute

label_key = label_key instance-attribute

max_length = max_length instance-attribute

pad_to_multiple_of = pad_to_multiple_of instance-attribute

property_key = property_key instance-attribute

tokenizer = tokenizer instance-attribute

__call__(samples: List[Union[List[int], Any, Dict[str, Any]]])

Call collate function

Parameters:

  • samples (List[Union[List[int], Any, Dict[str, Any]]]) –

    list of examples

Source code in safe/trainer/collator.py 
 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
def __call__(self, samples: List[Union[List[int], Any, Dict[str, Any]]]):
    """
    Call collate function

    Args:
        samples: list of examples
    """
    # Handle dict or lists with proper padding and conversion to tensor.
    tokenizer = self.get_tokenizer()

    # examples = samples
    examples = copy.deepcopy(samples)
    inputs = [example.pop(self.input_key, None) for example in examples]
    mc_labels = (
        torch.tensor([example.pop(self.property_key, None) for example in examples]).float()
        if self.property_key in examples[0]
        else None
    )

    if "input_ids" not in examples[0] and inputs is not None:
        batch = tokenizer(
            inputs,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=self.max_length,
            pad_to_multiple_of=self.pad_to_multiple_of,
        )
    else:
        batch = tokenizer.pad(
            examples,
            return_tensors="pt",
            padding=True,
            pad_to_multiple_of=self.pad_to_multiple_of,
            max_length=self.max_length,
        )

    # If special token mask has been preprocessed, pop it from the dict.
    batch.pop("special_tokens_mask", None)
    labels = batch.get(self.label_key, batch["input_ids"].clone())
    if tokenizer.pad_token_id is not None:
        labels[labels == tokenizer.pad_token_id] = -100
    batch[self.label_key] = labels

    if mc_labels is not None and self.include_descriptors:
        batch.update(
            {
                "mc_labels": mc_labels,
                # "input_text": inputs,
            }
        )
    return batch

get_tokenizer() cached

Get underlying tokenizer

Source code in safe/trainer/collator.py 
50
51
52
53
54
55
@functools.lru_cache()
def get_tokenizer(self):
    """Get underlying tokenizer"""
    if isinstance(self.tokenizer, SAFETokenizer):
        return self.tokenizer.get_pretrained()
    return self.tokenizer

Data Utils

batch_iterator(datasets, batch_size=100, n_examples=None, column='inputs')

Source code in safe/trainer/data_utils.py 
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
def batch_iterator(datasets, batch_size=100, n_examples=None, column="inputs"):
    if isinstance(datasets, Mapping):
        datasets = list(datasets.values())

    if not isinstance(datasets, (list, tuple)):
        datasets = [datasets]

    for dataset in datasets:
        iter_dataset = iter(dataset)
        if n_examples is not None and n_examples > 0:
            for _ in tqdm(range(0, n_examples, batch_size)):
                out = [next(iter_dataset)[column] for _ in range(batch_size)]
                yield out
        else:
            for out in tqdm(iter(partial(take, batch_size, iter_dataset), [])):
                yield [x[column] for x in out]

get_dataset(data_path, name: Optional[str] = None, tokenizer: Optional[Callable] = None, cache_dir: Optional[str] = None, streaming: bool = True, use_auth_token: bool = False, tokenize_column: Optional[str] = 'inputs', property_column: Optional[str] = 'descriptors', max_length: Optional[int] = None, num_shards=1024)

Get the datasets from the config file

Source code in safe/trainer/data_utils.py 
 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
def get_dataset(
    data_path,
    name: Optional[str] = None,
    tokenizer: Optional[Callable] = None,
    cache_dir: Optional[str] = None,
    streaming: bool = True,
    use_auth_token: bool = False,
    tokenize_column: Optional[str] = "inputs",
    property_column: Optional[str] = "descriptors",
    max_length: Optional[int] = None,
    num_shards=1024,
):
    """Get the datasets from the config file"""
    raw_datasets = {}
    if data_path is not None:
        data_path = upath.UPath(str(data_path))

        if data_path.exists():
            # then we need to load from disk
            data_path = str(data_path)
            # for some reason, the datasets package is not able to load the dataset
            # because the split where not originally proposed
            raw_datasets = datasets.load_from_disk(data_path)

            if streaming:
                if isinstance(raw_datasets, datasets.DatasetDict):
                    previous_num_examples = {k: len(dt) for k, dt in raw_datasets.items()}
                    raw_datasets = datasets.IterableDatasetDict(
                        {
                            k: dt.to_iterable_dataset(num_shards=num_shards)
                            for k, dt in raw_datasets.items()
                        }
                    )
                    for k, dt in raw_datasets.items():
                        if previous_num_examples[k] is not None:
                            setattr(dt, "num_examples", previous_num_examples[k])
                else:
                    num_examples = len(raw_datasets)
                    raw_datasets = raw_datasets.to_iterable_dataset(num_shards=num_shards)
                    setattr(raw_datasets, "num_examples", num_examples)

        else:
            data_path = str(data_path)
            raw_datasets = datasets.load_dataset(
                data_path,
                name=name,
                cache_dir=cache_dir,
                token=True if use_auth_token else None,
                streaming=streaming,
            )
    # that means we need to return a tokenized version of the dataset

    if property_column not in ["mc_labels", None]:
        raw_datasets = raw_datasets.rename_column(property_column, "mc_labels")

    columns_to_remove = None
    if tokenize_column is not None:
        columns_to_remove = [
            x
            for x in (get_dataset_column_names(raw_datasets) or [])
            if x not in [tokenize_column, "mc_labels"] and "label" not in x
        ] or None

    if tokenizer is None:
        if columns_to_remove is not None:
            raw_datasets = raw_datasets.remove_columns(columns_to_remove)
        return raw_datasets

    return raw_datasets.map(
        partial(
            tokenize_fn,
            tokenizer=tokenizer,
            tokenize_column=tokenize_column,
            max_length=max_length,
        ),
        batched=True,
        remove_columns=columns_to_remove,
    )

get_dataset_column_names(dataset: Union[datasets.Dataset, datasets.IterableDataset, Mapping])

Get the column names in a dataset

Parameters:

  • dataset (Union[Dataset, IterableDataset, Mapping]) –

    dataset to get the column names from

Source code in safe/trainer/data_utils.py 
18
19
20
21
22
23
24
25
26
27
28
29
30
31
def get_dataset_column_names(dataset: Union[datasets.Dataset, datasets.IterableDataset, Mapping]):
    """Get the column names in a dataset

    Args:
        dataset: dataset to get the column names from

    """
    if isinstance(dataset, (datasets.IterableDatasetDict, Mapping)):
        column_names = {split: dataset[split].column_names for split in dataset}
    else:
        column_names = dataset.column_names
    if isinstance(column_names, dict):
        column_names = list(column_names.values())[0]
    return column_names

take(n, iterable)

Return first n items of the iterable as a list

Source code in safe/trainer/data_utils.py 
13
14
15
def take(n, iterable):
    "Return first n items of the iterable as a list"
    return list(itertools.islice(iterable, n))

tokenize_fn(row: Dict[str, Any], tokenizer: Callable, tokenize_column: str = 'inputs', max_length: Optional[int] = None, padding: bool = False)

Perform the tokenization of a row Args: row: row to tokenize tokenizer: tokenizer to use tokenize_column: column to tokenize max_length: maximum size of the tokenized sequence padding: whether to pad the sequence

Source code in safe/trainer/data_utils.py 
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
def tokenize_fn(
    row: Dict[str, Any],
    tokenizer: Callable,
    tokenize_column: str = "inputs",
    max_length: Optional[int] = None,
    padding: bool = False,
):
    """Perform the tokenization of a row
    Args:
        row: row to tokenize
        tokenizer: tokenizer to use
        tokenize_column: column to tokenize
        max_length: maximum size of the tokenized sequence
        padding: whether to pad the sequence
    """
    # there's probably a way to do this with the tokenizer settings
    # but again, gotta move fast

    fast_tokenizer = (
        tokenizer.get_pretrained() if isinstance(tokenizer, SAFETokenizer) else tokenizer
    )

    return fast_tokenizer(
        row[tokenize_column],
        truncation=(max_length is not None),
        max_length=max_length,
        padding=padding,
        return_tensors=None,
    )