first commit

atroposlib/api/utils.py

@@ -0,0 +1,61 @@
+from typing import Dict, List, Optional, Tuple
+
+
+def grab_exact_from_heterogeneous_queue(
+    queue: List[Dict[str, List]], batch_size: int
+) -> Tuple[Optional[List], List]:
+    """
+    Grabs a batch of size batchsize from a queue of different sized items
+
+    e.g. queue = [{"tokens": [[1, 2, 3],[4, 5, 6, 7, 8]]}, {"tokens": [[9, 10]]}]
+
+    without going over the batchsize. This function will return a batch of size batchsize, and the new queue.
+
+    Because all groups are a common denominator of the batchsize, and all groups are a power of 2,
+    we can simplify a bit by assuming we can grab groups of groups to be equal to the maximum group size.
+    Note that we cannot drop items from groups, so we must grab the entire group if we grab it.
+
+    There may be a more efficient clearing mechanism by grouping these smaller groups heterogeneously, but
+    forcing them all into powers of two groups is a simple way to ensure we can grab a batch of the correct size.
+
+    :param queue:
+    :param batch_size:
+    :return: batch, new_queue
+    """
+    # check if we can even potentially grab a batch
+    if sum(len(item["tokens"]) for item in queue) < batch_size:
+        return None, queue
+    # Get max batch size
+    max_group_size = max(len(group["tokens"]) for group in queue)
+    group_sizes = set(len(group["tokens"]) for group in queue)
+    group_batching_storage = {i: [] for i in group_sizes}
+    # pack the groups into [max_group_size // group_size] packs
+    potential_batch = []
+    for i, item in enumerate(queue):
+        key = len(item["tokens"])
+        group_batching_storage[key].append({"group": item, "indx": i})
+        if len(group_batching_storage[key]) * key == max_group_size:
+            potential_batch.extend(group_batching_storage[key])
+            group_batching_storage[key] = []
+    if (
+        sum(len(grouped_items["group"]["tokens"]) for grouped_items in potential_batch)
+        < batch_size
+    ):
+        return None, queue
+    # we have a batch
+    batch = []
+    indxes_to_remove_from_queue = []
+    for item in potential_batch:
+        group = item["group"]
+        indx = item["indx"]
+        batch.append(group)
+        indxes_to_remove_from_queue.append(indx)
+        if sum(len(item["tokens"]) for item in batch) == batch_size:
+            break
+    if sum(len(item["tokens"]) for item in batch) != batch_size:
+        return None, queue
+    # remove the items from the queue
+    new_queue = [
+        item for i, item in enumerate(queue) if i not in indxes_to_remove_from_queue
+    ]
+    return batch, new_queue