codekingpro/portable-devtools

1from typing import Optional, Tuple, List
2from chromadb.api import BaseAPI
3from chromadb.api.types import (
4    Documents,
5    Embeddings,
6    IDs,
7    Metadatas,
8)
9
10
11def create_batches(
12    api: BaseAPI,
13    ids: IDs,
14    embeddings: Optional[Embeddings] = None,
15    metadatas: Optional[Metadatas] = None,
16    documents: Optional[Documents] = None,
17) -> List[Tuple[IDs, Optional[Embeddings], Optional[Metadatas], Optional[Documents]]]:
18    _batches: List[
19        Tuple[IDs, Optional[Embeddings], Optional[Metadatas], Optional[Documents]]
20    ] = []
21    if len(ids) > api.get_max_batch_size():
22        # create split batches
23        for i in range(0, len(ids), api.get_max_batch_size()):
24            _batches.append(
25                (
26                    ids[i : i + api.get_max_batch_size()],
27                    embeddings[i : i + api.get_max_batch_size()]
28                    if embeddings is not None
29                    else None,
30                    metadatas[i : i + api.get_max_batch_size()] if metadatas else None,
31                    documents[i : i + api.get_max_batch_size()] if documents else None,
32                )
33            )
34    else:
35        _batches.append((ids, embeddings, metadatas, documents))
36    return _batches
37