modaic 0.1.0__py3-none-any.whl

modaic/databases/vector_database/benchmarks/baseline.py

@@ -0,0 +1,123 @@
+from typing import TYPE_CHECKING, Dict, Iterable, List, Literal, Optional, Tuple
+
+import immutables
+import numpy as np
+from PIL import Image
+from tqdm import tqdm
+
+from modaic.context.base import Context, Embeddable
+
+from .common import _add_item_embedme, _items_have_multiple_embedmes
+
+if TYPE_CHECKING:
+    from modaic.databases.vector_database import VectorDatabase
+
+
+def has_collection(vdb: "VectorDatabase", collection_name: str) -> bool:
+    """
+    Check if a collection exists in the vector database.
+
+    Args:
+        collection_name: The name of the collection to check
+
+    Returns:
+        True if the collection exists, False otherwise
+    """
+    return vdb.ext.backend.has_collection(collection_name)
+
+
+def _embed_and_add_records(
+    vdb: "VectorDatabase",
+    collection_name: str,
+    embedmes: Dict[str, List[str | Image.Image]] | Dict[None, List[str | Image.Image]],
+    contexts: List[Context],
+):
+    # TODO: could add functionality for multiple embedmes per context (e.g. you want to embed both an image and a text description of an image)
+    all_embeddings = {}
+    if collection_name not in vdb.collections:
+        raise ValueError(
+            f"Collection {collection_name} not found in VectorDatabase's indexes, Please use VectorDatabase.create_collection() to create a collection first. Alternatively, you can use VectorDatabase.load_collection() to add records to an existing collection."
+        )
+    try:
+        first_index = next(iter(vdb.collections[collection_name].indexes.keys()))
+        # NOTE: get embeddings for each index
+        for index_name, index_config in vdb.collections[collection_name].indexes.items():
+            embeddings = index_config.embedder(embedmes)
+            # NOTE: Ensure embeddings is a 2D array (DSPy returns 1D for single strings, 2D for lists)
+            if embeddings.ndim == 1:
+                embeddings = embeddings.reshape(1, -1)
+            # NOTE: If index_name is None use the only index for the collection
+            all_embeddings[index_name or first_index] = embeddings
+    except Exception as e:
+        raise ValueError(f"Failed to create embeddings for index: {index_name}") from e
+
+    data_to_insert: List[immutables.Map[str, np.ndarray]] = []
+    # FIXME Probably should add type checking to ensure context matches schema, not sure how to do this efficiently
+    for i, item in enumerate(contexts):
+        embedding_map: dict[str, np.ndarray] = {}
+        for index_name, embedding in all_embeddings.items():
+            embedding_map[index_name] = embedding[i]
+
+        # Create a record with embedding and validated metadata
+        record = vdb.ext.backend.create_record(embedding_map, item)
+        data_to_insert.append(record)
+
+    vdb.ext.backend.add_records(collection_name, data_to_insert)
+    del data_to_insert
+
+
+def add_records(
+    vdb: "VectorDatabase",
+    collection_name: str,
+    records: Iterable[Embeddable | Tuple[str | Image.Image, Context]],
+    batch_size: Optional[int] = None,
+    embedme_scope: Literal["auto", "context", "index"] = "auto",
+):
+    """
+    Add items to a collection in the vector database.
+    Uses the Context's get_embed_context() method and the embedder to create embeddings.
+
+    Args:
+        collection_name: The name of the collection to add records to
+        records: The records to add to the collection
+        batch_size: Optional batch size for processing records
+    """
+    if not records:
+        return
+
+    # NOTE: Make embedmes compatible with the ext's hybrid search function
+    if embedme_scope == "auto":
+        if _items_have_multiple_embedmes(records):
+            embedme_scope = "index"
+        else:
+            embedme_scope = "context"
+
+    if embedme_scope == "index":
+        embedmes: Dict[str, List[str | Image.Image]] = {k: [] for k in vdb.collections[collection_name].indexes.keys()}
+    else:
+        # CAVEAT: We make embedmes a dict with None as opposed to a list so we don't have to type check it
+        embedmes: Dict[None, List[str | Image.Image]] = {None: []}
+
+    serialized_contexts = []
+    # TODO: add multi-processing/multi-threading here, just ensure that the backend is thread-safe. Maybe we add a class level parameter to check if the vendor is thread-safe. Embedding will still need to happen on a single thread
+    with tqdm(total=len(records), desc="Adding records to vector database", position=0) as pbar:
+        for _, item in tqdm(
+            enumerate(records),
+            desc="Adding records to vector database",
+            position=0,
+            leave=False,
+        ):
+            _add_item_embedme(embedmes, item)
+            serialized_contexts.append(item)
+
+            if batch_size is not None and len(serialized_contexts) == batch_size:
+                _embed_and_add_records(vdb, collection_name, embedmes, serialized_contexts)
+                if embedme_scope == "index":
+                    embedmes = {k: [] for k in embedmes.keys()}
+                else:
+                    embedmes = {None: []}
+                serialized_contexts = []
+                pbar.update(batch_size)
+
+    if serialized_contexts:
+        _embed_and_add_records(vdb, collection_name, embedmes, serialized_contexts)

modaic/databases/vector_database/benchmarks/common.py

@@ -0,0 +1,48 @@
+from typing import Dict, Iterable, List, Tuple
+
+from more_itertools import peekable
+from PIL import Image
+
+from modaic.context.base import Context, Embeddable
+
+
+def _has_multiple_embedmes(
+    item: Embeddable,
+):
+    """
+    Check if the item has multiple embedmes.
+    """
+    return item.embedme.__code__.co_argcount == 2
+
+
+def _items_have_multiple_embedmes(
+    records: Iterable[Embeddable | Tuple[str | Image.Image, Context]],
+):
+    """
+    Check if the first record has multiple embedmes.
+    """
+    p = peekable(records)
+    first_item = p.peek()
+    if isinstance(first_item, Embeddable) and _has_multiple_embedmes(first_item):
+        return True
+    return False
+
+
+def _add_item_embedme(
+    embedmes: Dict[str | None, List[str | Image.Image]],
+    item: Embeddable | Tuple[str | Image.Image, Context],
+):
+    """
+    Adds an item to the embedmes dictionary.
+    """
+    # Fast type check for tuple
+    if type(item) is tuple:
+        embedme = item[0]
+        for index in embedmes.keys():
+            embedmes[index].append(embedme)
+    elif _has_multiple_embedmes(item):
+        # CAVEAT: Context objects that implement Embeddable protocol and take in an index name as a parameter also accept None as the default index.
+        for index in embedmes.keys():
+            embedmes[index].append(item.embedme(index))
+    else:
+        embedmes[None].append(item.embedme())

modaic/databases/vector_database/benchmarks/fork.py

@@ -0,0 +1,132 @@
+import multiprocessing as mp
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from typing import TYPE_CHECKING, Dict, Iterable, List, Literal, Optional, Tuple
+
+import numpy as np
+from PIL import Image
+
+from modaic.context.base import Context, Embeddable
+
+from .common import _add_item_embedme, _items_have_multiple_embedmes
+
+if TYPE_CHECKING:
+    from modaic.databases.vector_database.vector_database import VectorDatabase, VectorDBBackend
+MAX_IN_FLIGHT = 8
+
+
+def add_records(
+    self: "VectorDatabase",
+    collection_name: str,
+    records: Iterable[Embeddable | Tuple[str | Image.Image, Context]],
+    batch_size: Optional[int] = None,
+    embedme_scope: Literal["auto", "context", "index"] = "auto",
+):
+    """
+    Add items to a collection in the vector database.
+    Uses the Context's get_embed_context() method and the embedder to create embeddings.
+
+    Args:
+        collection_name: The name of the collection to add records to
+        records: The records to add to the collection
+        batch_size: Optional batch size for processing records
+    """
+    if not records:
+        return
+
+    # NOTE: Make embedmes compatible with the ext's hybrid search function
+    if embedme_scope == "auto":
+        if _items_have_multiple_embedmes(records):
+            embedme_scope = "index"
+        else:
+            embedme_scope = "context"
+
+    # TODO: add multi-processing/multi-threading here, just ensure that the backend is thread-safe. Maybe we add a class level parameter to check if the vendor is thread-safe. Embedding will still need to happen on a single thread
+    def gen_embeded_records():
+        if embedme_scope == "index":
+            embedmes: Dict[str, List[str | Image.Image]] = {
+                k: [] for k in self.collections[collection_name].indexes.keys()
+            }
+        else:
+            # CAVEAT: We make embedmes a dict with None as opposed to a list so we don't have to type check it
+            embedmes: Dict[None, List[str | Image.Image]] = {None: []}
+
+        serialized_contexts = []
+
+        for item in records:
+            _add_item_embedme(embedmes, item)
+            serialized_contexts.append(item)
+
+            if batch_size is not None and len(serialized_contexts) == batch_size:
+                yield _embed_records(self, collection_name, embedmes), serialized_contexts
+                if embedme_scope == "index":
+                    embedmes = {k: [] for k in embedmes.keys()}
+                else:
+                    embedmes = {None: []}
+                serialized_contexts = []
+
+        if serialized_contexts:
+            yield _embed_records(self, collection_name, embedmes), serialized_contexts
+
+    mp.set_start_method("fork")
+
+    with ProcessPoolExecutor(max_workers=4) as pool:
+        pending = set()
+        for records in gen_embeded_records():
+            if len(pending) >= MAX_IN_FLIGHT:
+                done = next(as_completed(pending))
+                pending.remove(done)
+                done.result()  # raise any backend error now
+
+            pending.add(pool.submit(backend_add_records, self.ext.backend, collection_name, records))
+        for fut in as_completed(pending):
+            fut.result()
+
+
+def _embed_records(
+    self: "VectorDatabase",
+    collection_name: str,
+    embedmes: Dict[str, List[str | Image.Image]] | Dict[None, List[str | Image.Image]],
+):
+    # TODO: could add functionality for multiple embedmes per context (e.g. you want to embed both an image and a text description of an image)
+    all_embeddings = {}
+    if collection_name not in self.collections:
+        raise ValueError(
+            f"Collection {collection_name} not found in VectorDatabase's indexes, Please use VectorDatabase.create_collection() to create a collection first. Alternatively, you can use VectorDatabase.load_collection() to add records to an existing collection."
+        )
+    try:
+        first_index = next(iter(self.collections[collection_name].indexes.keys()))
+        # NOTE: get embeddings for each index
+        for index_name, index_config in self.collections[collection_name].indexes.items():
+            embeddings = index_config.embedder(embedmes)
+            # NOTE: Ensure embeddings is a 2D array (DSPy returns 1D for single strings, 2D for lists)
+            if embeddings.ndim == 1:
+                embeddings = embeddings.reshape(1, -1)
+            # NOTE: If index_name is None use the only index for the collection
+            all_embeddings[index_name or first_index] = embeddings
+    except Exception as e:
+        raise ValueError(f"Failed to create embeddings for index: {index_name}") from e
+    return all_embeddings
+
+
+def _add_records(
+    self: "VectorDatabase",
+    collection_name: str,
+    embeddings: Dict[str, np.ndarray],
+    contexts: List[Context],
+):
+    data_to_insert: List[Dict[str, np.ndarray]] = []
+    # FIXME Probably should add type checking to ensure context matches schema, not sure how to do this efficiently
+    for i, item in enumerate(contexts):
+        embedding_map: dict[str, np.ndarray] = {}
+        for index_name, embedding in embeddings.items():
+            embedding_map[index_name] = embedding[i]
+
+        # Create a record with embedding and validated metadata
+        record = self.ext.backend.create_record(embedding_map, item)
+        data_to_insert.append(record)
+
+    self.ext.backend.add_records(collection_name, data_to_insert)
+
+
+def backend_add_records(backend: "VectorDBBackend", collection_name: str, records: List[Dict[str, np.ndarray]]):
+    backend.add_records(collection_name, records)

modaic/databases/vector_database/benchmarks/threaded.py

@@ -0,0 +1,119 @@
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import TYPE_CHECKING, Dict, Iterable, List, Literal, Optional, Tuple
+
+import immutables
+import numpy as np
+from PIL import Image
+
+from modaic.context.base import Context, Embeddable
+
+from .common import _add_item_embedme, _items_have_multiple_embedmes
+
+if TYPE_CHECKING:
+    from modaic.databases.vector_database.vector_database import VectorDatabase
+MAX_IN_FLIGHT = 8
+
+
+def add_records(
+    self: "VectorDatabase",
+    collection_name: str,
+    records: Iterable[Embeddable | Tuple[str | Image.Image, Context]],
+    batch_size: Optional[int] = None,
+    embedme_scope: Literal["auto", "context", "index"] = "auto",
+):
+    """
+    Add items to a collection in the vector database.
+    Uses the Context's get_embed_context() method and the embedder to create embeddings.
+
+    Args:
+        collection_name: The name of the collection to add records to
+        records: The records to add to the collection
+        batch_size: Optional batch size for processing records
+    """
+    if not records:
+        return
+
+    # NOTE: Make embedmes compatible with the ext's hybrid search function
+    if embedme_scope == "auto":
+        if _items_have_multiple_embedmes(records):
+            embedme_scope = "index"
+        else:
+            embedme_scope = "context"
+
+    # TODO: add multi-processing/multi-threading here, just ensure that the backend is thread-safe. Maybe we add a class level parameter to check if the vendor is thread-safe. Embedding will still need to happen on a single thread
+    def gen_embeded_records():
+        if embedme_scope == "index":
+            embedmes: Dict[str, List[str | Image.Image]] = {
+                k: [] for k in self.collections[collection_name].indexes.keys()
+            }
+        else:
+            # CAVEAT: We make embedmes a dict with None as opposed to a list so we don't have to type check it
+            embedmes: Dict[None, List[str | Image.Image]] = {None: []}
+
+        serialized_contexts = []
+
+        for item in records:
+            _add_item_embedme(embedmes, item)
+            serialized_contexts.append(item)
+
+            if batch_size is not None and len(serialized_contexts) == batch_size:
+                yield _embed_and_create_records(self, collection_name, embedmes, serialized_contexts)
+                if embedme_scope == "index":
+                    embedmes = {k: [] for k in embedmes.keys()}
+                else:
+                    embedmes = {None: []}
+                serialized_contexts = []
+
+        if serialized_contexts:
+            yield _embed_and_create_records(self, collection_name, embedmes, serialized_contexts)
+
+    with ThreadPoolExecutor(max_workers=8) as pool:
+        pending = set()
+        for records in gen_embeded_records():
+            if len(pending) >= MAX_IN_FLIGHT:
+                done = next(as_completed(pending))
+                pending.remove(done)
+                done.result()  # raise any backend error now
+
+            pending.add(pool.submit(self.ext.backend.add_records, collection_name, records))
+        for fut in as_completed(pending):
+            fut.result()
+
+
+def _embed_and_create_records(
+    self: "VectorDatabase",
+    collection_name: str,
+    embedmes: Dict[str, List[str | Image.Image]] | Dict[None, List[str | Image.Image]],
+    contexts: List[Context],
+):
+    # TODO: could add functionality for multiple embedmes per context (e.g. you want to embed both an image and a text description of an image)
+    all_embeddings = {}
+    if collection_name not in self.collections:
+        raise ValueError(
+            f"Collection {collection_name} not found in VectorDatabase's indexes, Please use VectorDatabase.create_collection() to create a collection first. Alternatively, you can use VectorDatabase.load_collection() to add records to an existing collection."
+        )
+    try:
+        first_index = next(iter(self.collections[collection_name].indexes.keys()))
+        # NOTE: get embeddings for each index
+        for index_name, index_config in self.collections[collection_name].indexes.items():
+            embeddings = index_config.embedder(embedmes)
+            # NOTE: Ensure embeddings is a 2D array (DSPy returns 1D for single strings, 2D for lists)
+            if embeddings.ndim == 1:
+                embeddings = embeddings.reshape(1, -1)
+            # NOTE: If index_name is None use the only index for the collection
+            all_embeddings[index_name or first_index] = embeddings
+    except Exception as e:
+        raise ValueError(f"Failed to create embeddings for index: {index_name}") from e
+
+    data_to_insert: List[immutables.Map[str, np.ndarray]] = []
+    # FIXME Probably should add type checking to ensure context matches schema, not sure how to do this efficiently
+    for i, item in enumerate(contexts):
+        embedding_map: dict[str, np.ndarray] = {}
+        for index_name, embedding in all_embeddings.items():
+            embedding_map[index_name] = embedding[i]
+
+        # Create a record with embedding and validated metadata
+        record = self.ext.backend.create_record(embedding_map, item)
+        data_to_insert.append(record)
+
+    return data_to_insert