| 210 | return self.tokenizer(text_sample["text"], truncation=True, padding="max_length", max_length=self.max_seq_len) |
| 211 | |
| 212 | def _read_binary_tokenized_sample(self, sample: BatchEncoding): |
| 213 | seq_len = sample["len"] if "len" in sample else len(sample["input_ids"]) |
| 214 | |
| 215 | input_ids = np.frombuffer(sample["input_ids"], dtype=np.int64).copy() |
| 216 | if "attention_mask" in sample: |
| 217 | attention_mask = np.frombuffer(sample["attention_mask"], dtype=np.int64).copy() |
| 218 | else: |
| 219 | attention_mask = np.ones_like(input_ids) |
| 220 | |
| 221 | # calculate padding |
| 222 | pad_len = self.max_seq_len - seq_len |
| 223 | |
| 224 | # pad or truncate input_ids and attention_mask |
| 225 | if pad_len > 0: |
| 226 | input_ids = np.pad(input_ids, (0, pad_len), constant_values=self.tokenizer.pad_token_id) |
| 227 | attention_mask = np.pad(attention_mask, (0, pad_len), constant_values=0) |
| 228 | elif pad_len < 0: |
| 229 | input_ids = input_ids[: self.max_seq_len] |
| 230 | attention_mask = attention_mask[: self.max_seq_len] |
| 231 | |
| 232 | token_type_ids = np.zeros(self.max_seq_len, dtype=np.int64) |
| 233 | |
| 234 | return BatchEncoding( |
| 235 | data={ |
| 236 | "input_ids": input_ids.tolist(), |
| 237 | "attention_mask": attention_mask.tolist(), |
| 238 | "token_type_ids": token_type_ids.tolist(), |
| 239 | }, |
| 240 | n_sequences=1, |
| 241 | ) |
| 242 | |
| 243 | # How to process a sample |
| 244 | def __getitem__(self, idx: int) -> Union[Dict[str, Any], torch.Tensor]: |