| 490 | |
| 491 | |
| 492 | class BufferedIterator(Generic[T]): |
| 493 | def __init__(self, iterable: Iterable[T], buffer_size: int): |
| 494 | self.iterator = iter(iterable) |
| 495 | self.buffer = deque(maxlen=buffer_size) |
| 496 | self.buffer_size = buffer_size |
| 497 | self.lock = threading.Lock() |
| 498 | self.exhausted = False |
| 499 | self.filler_thread = threading.Thread(target=self._background_fill, daemon=True) |
| 500 | self.filler_thread.start() |
| 501 | |
| 502 | def _background_fill(self): |
| 503 | # Fill up the buffer, whenever possible, in the background |
| 504 | while not self.exhausted: |
| 505 | if len(self.buffer) < self.buffer_size: |
| 506 | try: |
| 507 | item = next(self.iterator) |
| 508 | with self.lock: |
| 509 | self.buffer.append(item) |
| 510 | except StopIteration: |
| 511 | self.exhausted = True |
| 512 | break |
| 513 | else: |
| 514 | time.sleep(0.01) # Sleep for a bit to avoid busy waiting |
| 515 | |
| 516 | def __iter__(self): |
| 517 | return self |
| 518 | |
| 519 | def __next__(self) -> T: |
| 520 | while True: |
| 521 | if not self.buffer: |
| 522 | if self.exhausted: |
| 523 | # We've exhausted the iterator and the buffer so we're done |
| 524 | raise StopIteration |
| 525 | else: |
| 526 | # The buffer is empty but the iterator is not exhausted yet. |
| 527 | # Let's give the filler thread a chance to add items to the buffer |
| 528 | time.sleep(0.01) |
| 529 | else: |
| 530 | with self.lock: |
| 531 | return self.buffer.popleft() |
| 532 | |
| 533 | |
| 534 | def split_packed_batch(batch: Any, microbatch_size: Union[int, float], padding_tolerance=1.0) -> Sequence: |