Method _fill_buffer

src/sequence_packer.py:196–233

Refills the internal buffer. - max_items_to_add: an amount less than or equal to the number of items to add Returns: the number of items actually added. The default implementation of this simply extends to src.buffer, which is initialized as a list in __in

(self, max_items_to_add=float("inf"))

        return None

    def _fill_buffer(self, max_items_to_add=float("inf")) -> int:
        """
        Refills the internal buffer.

        - max_items_to_add: an amount less than or equal to the number of items to add

        Returns: the number of items actually added.

        The default implementation of this simply extends to src.buffer, which is
        initialized as a list in __init__. Subclasses which want to use a different data
        structure for internal buffering should override this method and also add
        code in __init__ to initialize src.buffer appropriately.

        Any implementation of this MUST never place more than self.buffer_size items
        in the internal buffer.
        """
        items_added = 0
        # NOTE: this should be >=, kept as is to match model training code
        # TODO: change if training a new model
        while (self.buffer_size - len(self.buffer)) > self.src_batch_size:
            try:
                # if pulling another batch would fetch more than the requested max, stop
                if max_items_to_add < float("inf"):
                    if (items_added + self.src_batch_size) > max_items_to_add:
                        # print("Not adding, because of max_items_to_fetch")
                        break
                incoming_batch = next(self.src_iterator)
                assert (
                    len(incoming_batch) <= self.src_batch_size
                ), f"expected {len(incoming_batch)=} <= {self.src_batch_size=}"
                for item in incoming_batch:
                    if len(item["input_ids"]) > 0:  # ignore empty sequences
                        self.buffer.append(item["input_ids"])
                        items_added += 1
                        self._seqs_consumed += 1
            except StopIteration:
                break
        return items_added

    def _generate_batches(self):
        """