Add (key, data) to shuffler.
(self, key: type_utils.Key, data: bytes)
| 276 | self._in_memory = False |
| 277 | |
| 278 | def add(self, key: type_utils.Key, data: bytes) -> bool: |
| 279 | """Add (key, data) to shuffler.""" |
| 280 | if self._read_only: |
| 281 | raise AssertionError('add() cannot be called after __iter__.') |
| 282 | if not isinstance(data, bytes): |
| 283 | raise AssertionError( |
| 284 | f'Only bytes (not {type(data)}) can be stored in Shuffler! This' |
| 285 | ' likely indicates that non-integer keys were used when generating' |
| 286 | ' the dataset.' |
| 287 | ) |
| 288 | hkey = self._hasher.hash_key(key) |
| 289 | if self._ignore_duplicates: |
| 290 | if hkey in self._seen_keys: |
| 291 | return # pytype: disable=bad-return-type |
| 292 | self._seen_keys.add(hkey) |
| 293 | if self._disable_shuffling: |
| 294 | # Use the original key and not the hashed key to maintain the order. |
| 295 | hkey = typing.cast(int, key) |
| 296 | self._total_bytes += len(data) |
| 297 | if self._in_memory: |
| 298 | self._add_to_mem_buffer(hkey, data) |
| 299 | else: |
| 300 | self._add_to_bucket(hkey, data) |
| 301 | self._num_examples += 1 # pytype: disable=bad-return-type |
| 302 | |
| 303 | def __iter__(self) -> Iterator[type_utils.KeySerializedExample]: |
| 304 | self._read_only = True |