(iterable, batch_size, batch_type)
| 555 | |
| 556 | |
| 557 | def _batch_iterator(iterable, batch_size, batch_type): |
| 558 | streams = None |
| 559 | max_length = 0 |
| 560 | |
| 561 | for example in iterable: |
| 562 | if not isinstance(example, tuple): |
| 563 | example = (example,) |
| 564 | |
| 565 | if batch_type == "examples": |
| 566 | if streams and len(streams[0]) == batch_size: |
| 567 | yield streams |
| 568 | streams = None |
| 569 | |
| 570 | elif batch_type == "tokens": |
| 571 | max_length = max(max_length, len(example[0])) |
| 572 | |
| 573 | if streams and (len(streams[0]) + 1) * max_length > batch_size: |
| 574 | yield streams |
| 575 | streams = None |
| 576 | max_length = len(example[0]) |
| 577 | |
| 578 | else: |
| 579 | raise ValueError("Invalid batch type %s" % batch_type) |
| 580 | |
| 581 | if streams is None: |
| 582 | streams = tuple([] for _ in example) |
| 583 | for batch, element in zip(streams, example): |
| 584 | if element is None and len(streams) > 1: |
| 585 | raise ValueError("Input iterables do not have the same length") |
| 586 | batch.append(element) |
| 587 | |
| 588 | if streams is not None: |
| 589 | yield streams |
no test coverage detected