roboml 0.2.0__py3-none-any.whl

roboml/databases/__init__.py

@@ -0,0 +1,5 @@
+from .chroma import ChromaDB
+
+__all__ = [
+    "ChromaDB",
+]

roboml/databases/_base.py

@@ -0,0 +1,166 @@
+from abc import abstractmethod
+from logging import Logger
+import inspect
+
+from pydantic import validate_call, ValidationError
+
+
+from roboml.models._encoding import EncodingModel
+from roboml.utils import Status, logging
+
+
+class VectorDBTemplate:
+    """
+    This class describes a VectorDB template.
+    """
+
+    def __init__(self, *, name: str, **_):
+        """__init__.
+        :param name:
+        :type name: str
+        :param _:
+        """
+        self.name: str = name
+        self.encoding_model: EncodingModel
+        self.logger: Logger = logging.getLogger(self.name)
+        self.status: Status = Status.LOADED
+
+    def initialize(self, **kwargs) -> None:
+        """
+        Calls the db initialization function and sets status accordingly
+        """
+        if self.status == Status.READY:
+            return
+
+        self.status = Status.INITIALIZING
+        try:
+            validated_init = validate_call(self._initialize)
+            validated_init(**kwargs)
+        except Exception as e:
+            self.logger.error(f"Initialization Error: {e}")
+            self.status = Status.INITIALIZATION_ERROR
+            raise e
+        self.logger.info(f"{self.__class__.__name__} VectorDB initialized")
+        self.status = Status.READY
+
+    def get_status(self):
+        """Returns status of the model node"""
+        return self.status.name
+
+    def add(self, **kwargs) -> dict:
+        """VectorDB specific add function.
+        :param kwargs:
+        :rtype: None
+        """
+        if self.status is not Status.READY:
+            self.logger.error("Error: DB method called before initialization")
+            raise Exception("DB Method Called called before initialization")
+        try:
+            data_model = inspect.signature(self._add).parameters["data"].annotation
+            data = data_model(**kwargs)
+        except ValidationError:
+            self.logger.error("Validation Error occured for inference input")
+            raise
+        return self._add(data)
+
+    def conditional_add(self, **kwargs) -> dict:
+        """VectorDB specific conditional_add function.
+        :param kwargs:
+        :rtype: None
+        """
+        if self.status is not Status.READY:
+            self.logger.error("Error: DB method called before initialization")
+            raise Exception("DB Method Called called before initialization")
+        try:
+            data_model = (
+                inspect.signature(self._conditional_add).parameters["data"].annotation
+            )
+            data = data_model(**kwargs)
+        except ValidationError:
+            self.logger.error("Validation Error occured for inference input")
+            raise
+        return self._conditional_add(data)
+
+    def metadata_query(self, **kwargs) -> dict:
+        """VectorDB specific metadata_query function.
+        :param kwargs:
+        :rtype: None
+        """
+        if self.status is not Status.READY:
+            self.logger.error("Error: DB method called before initialization")
+            raise Exception("DB Method Called called before initialization")
+        try:
+            data_model = (
+                inspect.signature(self._metadata_query).parameters["data"].annotation
+            )
+            data = data_model(**kwargs)
+        except ValidationError:
+            self.logger.error("Validation Error occured for inference input")
+            raise
+        return self._metadata_query(data)
+
+    def query(self, **kwargs) -> dict:
+        """VectorDB specific query function.
+        :param kwargs:
+        :rtype: None
+        """
+        if self.status is not Status.READY:
+            self.logger.error("Error: DB method called before initialization")
+            raise Exception("DB Method Called called before initialization")
+        try:
+            data_model = inspect.signature(self._query).parameters["data"].annotation
+            data = data_model(**kwargs)
+        except ValidationError:
+            self.logger.error("Validation Error occured for inference input")
+            raise
+        return self._query(data)
+
+    @abstractmethod
+    def _initialize(self, *_, **__) -> None:
+        """VectorDB specific initialize function, to be implemented in derived classes.
+        :param kwargs:
+        :rtype: None
+        """
+        raise NotImplementedError(
+            "VectorDB specific initialize method needs to be implemented by derived classes"
+        )
+
+    @abstractmethod
+    def _add(self, *_, **__) -> dict:
+        """VectorDB specific add function, to be implemented in derived classes.
+        :param kwargs:
+        :rtype: None
+        """
+        raise NotImplementedError(
+            "VectorDB specific add method needs to be implemented by derived classes"
+        )
+
+    @abstractmethod
+    def _conditional_add(self, *_, **__) -> dict:
+        """VectorDB specific conditional_add function, to be implemented in derived classes.
+        :param kwargs:
+        :rtype: None
+        """
+        raise NotImplementedError(
+            "VectorDB specific conditional_add method needs to be implemented by derived classes"
+        )
+
+    @abstractmethod
+    def _metadata_query(self, *_, **__) -> dict:
+        """VectorDB specific metadata_query function, to be implemented in derived classes.
+        :param kwargs:
+        :rtype: None
+        """
+        raise NotImplementedError(
+            "VectorDB specific metadata_query method needs to be implemented by derived classes"
+        )
+
+    @abstractmethod
+    def _query(self, *_, **__) -> dict:
+        """VectorDB specific query function, to be implemented in derived classes.
+        :param kwargs:
+        :rtype: None
+        """
+        raise NotImplementedError(
+            "VectorDB specific query method needs to be implemented by derived classes"
+        )

roboml/databases/chroma.py

@@ -0,0 +1,189 @@
+from typing import Optional
+from chromadb import PersistentClient
+from chromadb.api import ClientAPI
+from chromadb.config import Settings
+
+from roboml.interfaces import DBAdd, DBMetadataQuery, DBQuery
+from roboml.models._encoding import EncodingModel
+from roboml.ray import app, ingress_decorator
+
+from ._base import VectorDBTemplate
+
+
+@ingress_decorator
+class ChromaDB(VectorDBTemplate):
+    """
+    ChromaDB Wrapper.
+    """
+
+    def __init__(self, **kwargs):
+        """__init__.
+        :param kwargs:
+        """
+        super().__init__(**kwargs)
+        self.vectordb: ClientAPI
+
+    @app.post("/initialize")
+    def _initialize(
+        self,
+        db_location: str = "./data",
+        username: Optional[str] = None,
+        password: Optional[str] = None,
+        encoder: Optional[dict] = None,
+    ) -> None:
+        """
+        Initializes the db.
+        """
+        if username and password:
+            self.logger.warning(
+                "Username/password authentication can only be used with local ChromaDB client. Cannot use username/password authentication with roboml ChromaDB instance."
+            )
+        # initialize the encoding model
+        self.encoding_model = EncodingModel(name="encoding_model")
+        encoder = encoder or {}
+        self.encoding_model._initialize(**encoder)
+
+        # create a vectordb client
+        self.vectordb = PersistentClient(
+            settings=Settings(anonymized_telemetry=False),
+            path=db_location,
+        )
+
+    @app.post("/add")
+    def _add(self, data: DBAdd) -> dict:
+        """Add data to the given collection.
+        :param data:
+        :param type: DBAdd
+        :rtype: dict
+        """
+        # If specified reset existing collection
+        if data.reset_collection:
+            try:
+                self.vectordb.delete_collection(name=data.collection_name)
+            except ValueError:
+                self.logger.warning(
+                    f"Cannot delete collection with name {data.collection_name} as it does not exist."
+                )
+        try:
+            # create a collection if one doesnt exist
+            collection = self.vectordb.get_or_create_collection(
+                name=data.collection_name, metadata={"hnsw:space": data.distance_func}
+            )
+            # create embeddings
+            embeddings = self.encoding_model.embed_documents(data.documents)
+            # add to collection
+            collection.add(
+                documents=data.documents,
+                embeddings=embeddings,
+                metadatas=data.metadatas,
+                ids=data.ids,
+            )
+        except Exception as e:
+            self.logger.error(f"Exception occured: {e}")
+            raise
+
+        return {"output": "Success"}
+
+    @app.post("/conditional_add")
+    def _conditional_add(self, data: DBAdd) -> dict:
+        """First check if id exists if not then add data to the collection provided
+        Update metadatas of the ids that exist
+        :param data:
+        :param type: DBAdd
+        :rtype: dict
+        """
+        try:
+            # create a collection if one doesnt exist
+            collection = self.vectordb.get_or_create_collection(
+                name=data.collection_name, metadata={"hnsw:space": data.distance_func}
+            )
+            # check for ids that already exist in DB
+            already_existing = collection.get(ids=data.ids)
+
+            # if ids found in database, remove them from main input lists
+            # update metadatas of already existing IDs
+            metadatas_to_update = []
+            to_be_deleted = []
+            if already_existing_ids := already_existing["ids"]:
+                for idx in range(len(data.ids)):
+                    if data.ids[idx] in already_existing_ids:
+                        to_be_deleted.append(idx)
+                        metadatas_to_update.append(data.metadatas[idx].copy())
+                # do the metadata update
+                collection.update(
+                    ids=already_existing_ids, metadatas=metadatas_to_update
+                )
+        except Exception as e:
+            self.logger.error(f"Exception occured: {e}")
+            raise
+
+        # delete from indices that were updated
+        for idx in sorted(to_be_deleted, reverse=True):
+            del data.ids[idx]
+            del data.metadatas[idx]
+            del data.documents[idx]
+        # add the remaining data
+        return self._add(data) if data.ids else {"output": "Success"}
+
+    @app.post("/metadata_query")
+    def _metadata_query(self, data: DBMetadataQuery) -> dict:
+        """Retreive data by metadata query.
+        :param data:
+        :param type: DBMetadataQuery
+        :rtype: dict
+        """
+        # create filters for all metadata values
+        all_filters = []
+        for metadata in data.metadatas:
+            # create filter for each metadata hashmap
+            if len(metadata) > 1:
+                filter = {"$and": [{i: {"$eq": metadata[i]}} for i in metadata]}
+            # if there is only one metadata entry, $and is not needed
+            elif len(metadata) == 1:
+                filter = {list(metadata.keys())[0]: {"$eq": list(metadata.values())[0]}}
+            else:
+                continue
+            all_filters.append(filter)
+
+        # if no filters, return output as None
+        if len(all_filters) == 0:
+            self.logger.warning(
+                "The metadata filters received were empty, please call query method for retreiving data without metadata filtering."
+            )
+            return {"output": []}
+
+        # if there are multiple filters, add an $or
+        filters = {"$or": all_filters} if len(all_filters) > 1 else all_filters[0]
+
+        try:
+            # get collection
+            collection = self.vectordb.get_collection(name=data.collection_name)
+            # get filtered data
+            output = collection.get(where=filters)
+        except Exception as e:
+            self.logger.error(f"Exception occured: {e}")
+            raise
+        return {"output": output}
+
+    @app.post("/query")
+    def _query(self, data: DBQuery) -> dict:
+        """
+        Retreives results for a given DB query
+        :param data:
+        :param type: DBQuery
+        :rtype: dict
+        """
+        query_vec = self.encoding_model.embed_query(data.query)
+        query_vec = query_vec.tolist()
+        try:
+            # create a collection for the map data
+            collection = self.vectordb.get_collection(name=data.collection_name)
+            output = (
+                collection.query(query_embeddings=query_vec, n_results=data.n_results)
+                or []
+            )
+        except Exception as e:
+            self.logger.error(f"Exception occured: {e}")
+            raise
+
+        return {"output": output}

roboml/interfaces.py

@@ -0,0 +1,137 @@
+from typing import Optional, Union
+
+import numpy as np
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class NodeInit(BaseModel):
+    """NodeInit."""
+
+    node_name: str
+    node_type: str
+
+
+class NodeDeinit(BaseModel):
+    """NodeDeinit."""
+
+    node_name: str
+
+
+# IO Interfaces
+class AudioInput(BaseModel):
+    """
+    Input values for audio inference
+    """
+
+    query: Union[str, bytes] = Field(title="Audio input raw bytes", min_length=1)
+
+
+class ImageInput(BaseModel):
+    """
+    Input values for image inference
+    """
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    images: Union[list[str], list[np.ndarray]] = Field(
+        title="List of images as base64 strings or numpy arrays", min_length=1
+    )
+
+
+class TextInput(BaseModel):
+    """
+    Input values for text inference
+    """
+
+    query: Union[str, list[dict]] = Field(title="Input to the model", min_length=1)
+
+
+class TextToSpeechInput(TextInput):
+    """
+    Input values for text to speech inference
+    """
+
+    voice: Optional[str] = Field(title="Voice to use", default=None)
+    get_bytes: bool = Field(title="Get raw audio bytes", default=False)
+
+
+class SpeechToTextInput(AudioInput):
+    """
+    Input values for speech to text inference
+    """
+
+    max_new_tokens: int = Field(
+        title="Maximum number of new tokens to be generated", default=128
+    )
+
+
+class LLMInput(TextInput):
+    """
+    Input values for LLM inference
+    """
+
+    max_new_tokens: int = Field(
+        title="Maximum number of new tokens to be generated", default=100
+    )
+    temperature: float = Field(
+        title="Temperature with which inference is to be generated", default=0.7
+    )
+
+
+class VLLMInput(ImageInput, LLMInput):
+    """
+    Input values for multi modal LLM inference
+    """
+
+    pass
+
+
+class DetectionInput(ImageInput):
+    """Input for Detection models."""
+
+    threshold: float = Field(title="Detection confidence threshold", default=0.5)
+    get_dataset_labels: bool = Field(
+        title="Get dataset label string names", default=True
+    )
+    labels_to_track: Optional[list[str]] = Field(
+        title="List of labels to track. Only used if tracking is enabled during initialization",
+        default=None,
+    )
+
+
+# DB Interfaces
+class DBAdd(BaseModel):
+    """
+    Documents to be added to DB
+    """
+
+    collection_name: str = Field(title="DB Collection Name")
+    ids: list[str] = Field(title="Document IDs", min_length=1)
+    metadatas: list[dict] = Field(title="Document metadatas", min_length=1)
+    documents: list[str] = Field(title="Documents", min_length=1)
+    distance_func: str = Field(
+        title="Distance function for the collection", default="l2"
+    )
+    reset_collection: bool = Field(
+        title="Delete existing collection with the name defined in collection_name and create it again with the data provided",
+        default=False,
+    )
+
+
+class DBMetadataQuery(BaseModel):
+    """
+    For retreiving documents based on metadata
+    """
+
+    collection_name: str = Field(title="DB Collection Name")
+    metadatas: list[dict] = Field(title="Document metadatas", min_length=1)
+
+
+class DBQuery(BaseModel):
+    """
+    For retreiving documents based on query
+    """
+
+    collection_name: str = Field(title="DB Collection Name")
+    query: str = Field(title="Query string")
+    n_results: int = Field(title="Number of results", default=1)

roboml/main.py

@@ -0,0 +1,62 @@
+import argparse
+
+from ray import serve
+
+from roboml.ray.app_factory import AppFactory
+from roboml.resp_server.server import Server
+from roboml.utils import logger
+
+
+def parse_args_for_ray() -> argparse.Namespace:
+    """Parse arguments."""
+    parser = argparse.ArgumentParser(description="Run models and event handlers")
+    parser.add_argument(
+        "--host", type=str, help="Specify server host address. Default '127.0.0.1'"
+    )
+    parser.add_argument("--port", type=int, help="Specify server port. Default 8000")
+    parser.add_argument(
+        "--nodes_per_cpu",
+        type=int,
+        help="Specify number of nodes to run per CPU. Default None",
+    )
+    parser.add_argument(
+        "--nodes_per_gpu",
+        type=int,
+        help="Specify number of nodes to run per GPU. Default None",
+    )
+    return parser.parse_args()
+
+
+def ray() -> None:
+    """Main entry function for ray"""
+
+    args = parse_args_for_ray()
+    host = args.host or "0.0.0.0"
+    port = args.port or 8000
+    nodes_per_cpu = args.nodes_per_cpu or None
+    nodes_per_gpu = args.nodes_per_gpu or None
+    app_factory = AppFactory.bind(
+        nodes_per_cpu=nodes_per_cpu, nodes_per_gpu=nodes_per_gpu
+    )
+    serve.start(http_options=serve.HTTPOptions(host=host, port=port))
+    serve.run(app_factory, name="app_factory", blocking=True)
+
+
+def parse_args_for_resp() -> argparse.Namespace:
+    """Parse arguments."""
+    parser = argparse.ArgumentParser(description="Run models and event handlers")
+    parser.add_argument(
+        "--host", type=str, help="Specify server host address. Default '0.0.0.0'"
+    )
+    parser.add_argument("--port", type=int, help="Specify server port. Default 6379")
+    return parser.parse_args()
+
+
+def resp() -> None:
+    """Main entry function for resp"""
+
+    args = parse_args_for_resp()
+    host = args.host or "0.0.0.0"
+    port = args.port or 6379
+    server = Server(logger)
+    server._start_server(host=host, port=port)

roboml/models/__init__.py

@@ -0,0 +1,13 @@
+from .mllm import Idefics, TransformersMLLM
+from .llm import TransformersLLM
+from .speech_to_text import Whisper
+from .text_to_speech import Bark, SpeechT5
+
+__all__ = [
+    "Whisper",
+    "TransformersLLM",
+    "TransformersMLLM",
+    "SpeechT5",
+    "Idefics",
+    "Bark",
+]

roboml/models/_base.py

@@ -0,0 +1,86 @@
+from abc import abstractmethod
+from logging import Logger
+from threading import Lock
+from typing import Any, Optional
+
+from pydantic import ValidationError, validate_call
+import torch
+import inspect
+
+from roboml.utils import Status, logging
+
+
+class ModelTemplate:
+    """
+    This class describes a model template.
+    """
+
+    def __init__(self, *, name: str, init_timeout: int = 600, **_):
+        self.name: str = name
+        self.init_timeout: Optional[int] = init_timeout  # 10 minutes
+        self.device: str = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model: Any = None
+        self.pre_processor: Any = None
+        self.status: Status = Status.LOADED
+        self.logger: Logger = logging.getLogger(self.name)
+        self.lock = Lock()
+
+    def initialize(self, **kwargs) -> None:
+        """
+        Calls the models initialization function and sets status accordingly
+        """
+        if self.status == Status.READY:
+            return
+
+        self.status = Status.INITIALIZING
+        try:
+            validated_init = validate_call(self._initialize)
+            validated_init(**kwargs)
+        except Exception as e:
+            self.logger.error(f"Initialization Error: {e}")
+            self.status = Status.INITIALIZATION_ERROR
+            raise e
+        self.logger.info(f"{self.__class__.__name__} Model initialized")
+        self.status = Status.READY
+
+    def inference(self, **kwargs) -> dict:
+        """
+        Calls the models inference function and sets status accordingly
+        """
+        if self.status is not Status.READY:
+            self.logger.error("Error: Inference called before initialization")
+            raise Exception("Inference called before initialization")
+        try:
+            data_model = (
+                inspect.signature(self._inference).parameters["data"].annotation
+            )
+            data = data_model(**kwargs)
+        except ValidationError:
+            self.logger.error("Validation Error occured for inference input")
+            raise
+        with self.lock:
+            result = self._inference(data)
+        return result
+
+    def get_status(self):
+        """Returns status of the model node"""
+        return self.status.name
+
+    @abstractmethod
+    def _initialize(self, *_, **__) -> None:
+        """Model specific initialization function, to be implemented in derived classes.
+        :param kwargs:
+        :rtype: None
+        """
+        raise NotImplementedError(
+            "Model specific initializaion needs to be implemented by derived model classes"
+        )
+
+    @abstractmethod
+    def _inference(self, *_, **__) -> dict:
+        """Model specific inference function, to be implemented in derived classes.
+        :rtype: dict
+        """
+        raise NotImplementedError(
+            "Model specific inference needs to be implemented by derived model classes"
+        )