formatting

Author: tylerhutcherson

README.md

@@ -1,5 +1,11 @@
 # Redis Model Store
 
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+![Language](https://img.shields.io/github/languages/top/redis-applied-ai/redis-model-store)
+[![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
+![GitHub last commit](https://img.shields.io/github/last-commit/redis-applied-ai/redis-model-store)
+[![pypi](https://badge.fury.io/py/redisvl.svg)](https://pypi.org/project/redis-model-store/)
+
 `redis-model-store` is a simple Python library designed to handle versioning and serialization of AI/ML models into Redis. It provides a streamlined way to manage your machine learning models in Redis.
 
 ## Features
@@ -8,7 +14,6 @@
 - **Sharding for Large Models**: Splits large serialized payloads into manageable chunks to optimize Redis storage.
 - **Version Management**: Automatically manages model versions in Redis, allowing you to store and retrieve specific versions.
 
----
 
 ## Installation
 ```bash
@@ -49,12 +54,16 @@ model = RandomForestClassifier()
 model.fit(X_train, y_train)
 
 # Save the model to Redis
-version = model_store.save_model(model, "random_forest")
+version = model_store.save_model(model, name="random_forest", description="Random forest classifier model")
 ```
 
 ### Load models
 ```python
 # Grab the latest model
-model = model_store.load_model("random_forest")
+model = model_store.load_model(name="random_forest")
+
+# Grab a specific model version
+model = model_store.load_model(name="random_forest", version=version)
 ```
 
+## 

model_store/serialize.py

@@ -4,6 +4,7 @@
 
 class SerializationError(Exception):
     """Raised when model serialization or deserialization fails."""
+
     pass
 
 

model_store/shard_manager.py

@@ -36,7 +36,8 @@ def __init__(self, shard_size: int, serializer: Serializer = PickleSerializer())
         self.shard_size = shard_size
         self.serializer = serializer
 
-    def _shardify(self, data: bytes) -> Iterator[bytes]:
+    @staticmethod
+    def _shardify(data: bytes, shard_size: int) -> Iterator[bytes]:
         """
         Split serialized data into fixed-size shards.
 
@@ -47,8 +48,8 @@ def _shardify(self, data: bytes) -> Iterator[bytes]:
             bytes: Successive shards of the data, each up to shard_size in length.
         """
         total_size = len(data)
-        for start in range(0, total_size, self.shard_size):
-            yield data[start : start + self.shard_size]
+        for start in range(0, total_size, shard_size):
+            yield data[start : start + shard_size]
 
     @staticmethod
     def shard_key(model_name: str, model_version: str, idx: int) -> str:
@@ -65,7 +66,7 @@ def shard_key(model_name: str, model_version: str, idx: int) -> str:
         """
         return f"shard:{model_name}:{model_version}:{idx}"
 
-    def to_shards(self, model: Any) -> List[bytes]:
+    def to_shards(self, model: Any) -> Iterator[bytes]:
         """
         Convert model into smaller chunks (shards) ready for storage.
 
@@ -74,14 +75,13 @@ def to_shards(self, model: Any) -> List[bytes]:
 
         Returns:
             List[bytes]: List of binary shards derived from the model.
-            TODO -- returns a generator here
 
-        Raises:
-            SerializationError: If model serialization fails.
+        Yields:
+            bytes: Successive shards of the data, each up to shard_size in length.
         """
         try:
             serialized_data = self.serializer.dumps(model)
-            return self._shardify(serialized_data)
+            return self._shardify(serialized_data, self.shard_size)
         except Exception as e:
             raise SerializationError(f"Failed to serialize model: {str(e)}") from e
 

model_store/store.py

@@ -11,6 +11,7 @@
 
 class ModelStoreError(Exception):
     """Raised when model store I/O operations fail."""
+
     pass
 
 
@@ -214,7 +215,9 @@ def load_model(self, name: str, version: Optional[str] = None) -> Any:
                 logger.debug(f"Retrieving specified model version: {version}")
                 model_version = self.model_registry.get_version(name, version)
 
-            logger.info(f"Retrieved model version metadata ({current_timestamp()-st:.4f}s)")
+            logger.info(
+                f"Retrieved model version metadata ({current_timestamp()-st:.4f}s)"
+            )
             model = self._from_redis(model_version.shard_keys)
 
             total_duration = current_timestamp() - total_start