orient_express 2.3.0__tar.gz → 2.3.1__tar.gz

{orient_express-2.3.0 → orient_express-2.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: orient_express
-Version: 2.3.0
+Version: 2.3.1
 Summary: A library to simplify model deployment to Vertex AI
 Author: Alexey Zankevich
 Author-email: alex.zankevich@shiftsmart.com
@@ -21,6 +21,7 @@ Requires-Dist: pyyaml (>=6.0,<7.0)
 Requires-Dist: scikit-learn (>=1.5.0,<2.0.0)
 Requires-Dist: torch (>=2.5.0)
 Requires-Dist: torchvision (>=0.20.0)
+Requires-Dist: tqdm (>=4.67.1)
 Requires-Dist: xgboost (>=2.0.0,<3.0.0)
 Description-Content-Type: text/markdown
 

{orient_express-2.3.0 → orient_express-2.3.1}/orient_express/predictors/vector_index.py

@@ -86,38 +86,48 @@ class VectorIndex(Predictor):
             )
         return all_results
 
-    def aggregate(self) -> "VectorIndex":
+    def aggregate(self, per_label: bool = False) -> "VectorIndex":
         """
-        Aggregate the vectors into a single centroid per label.
+        Aggregate the vectors into a single centroid per unique label group.
+
+        Args:
+            per_label: if True, create one centroid per individual label
+                (splitting label groups apart). If False (default), create
+                one centroid per unique label group.
         """
-        label_to_indices: dict = {}
+        key_to_indices: dict = {}
         for i, label_group in enumerate(self.labels):
-            for label in label_group:
-                if label not in label_to_indices:
-                    label_to_indices[label] = []
-                label_to_indices[label].append(i)
-
-        unique_labels = sorted(label_to_indices.keys(), key=str)
+            if per_label:
+                keys = [label for label in label_group]
+            else:
+                keys = [tuple(sorted(label_group))]
+            for key in keys:
+                if key not in key_to_indices:
+                    key_to_indices[key] = []
+                key_to_indices[key].append(i)
+
+        sorted_keys = sorted(key_to_indices.keys(), key=str)
         centroids = np.empty(
-            (len(unique_labels), self.vectors.shape[1]), dtype=self.vectors.dtype
+            (len(sorted_keys), self.vectors.shape[1]), dtype=self.vectors.dtype
         )
 
-        for i, label in enumerate(unique_labels):
-            indices = label_to_indices[label]
+        new_labels = []
+        for i, key in enumerate(sorted_keys):
+            indices = key_to_indices[key]
             centroid = self.vectors[indices].mean(axis=0)
             norm = np.linalg.norm(centroid)
             if norm > 0:
                 centroid = centroid / norm
             centroids[i] = centroid
+            new_labels.append([key] if per_label else list(key))
 
-        new_labels = [[label] for label in unique_labels]
         return VectorIndex(vectors=centroids, labels=new_labels)
 
     def get_serving_container_image_uri(self) -> str:
         warnings.warn(
             "VectorIndex does not support serving via a container. Returning incompatible image URI."
         )
-        return "us-west1-docker.pkg.dev/shiftsmartinc/orient-express/image-onnx:v2.1.2"
+        return "us-west1-docker.pkg.dev/shiftsmart-api/orient-express/image-onnx:v2.1.2"
 
     def get_serving_container_health_route(self, model_name) -> str:
         return f"/v1/models/{model_name}"

{orient_express-2.3.0 → orient_express-2.3.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "orient_express"
-version = "2.3.0"
+version = "2.3.1"
 description = "A library to simplify model deployment to Vertex AI"
 authors = ["Alexey Zankevich <alex.zankevich@shiftsmart.com>", "Ian Myers <ian.myers@shiftsmart.com>"]
 readme = "README.md"
@@ -18,6 +18,7 @@ Pillow = ">=10.0.1"
 gcsfs = "^2024.10.0"
 scikit-learn = "^1.5.0"
 xgboost = "^2.0.0"
+tqdm = ">=4.67.1"
 onnxruntime-gpu = {version = "^1.20", markers = "sys_platform == 'linux' and platform_machine == 'x86_64' or sys_platform == 'win32' and platform_machine == 'AMD64'"}
 onnxruntime = {version = "^1.20", markers = "sys_platform == 'darwin' or platform_machine == 'aarch64' or platform_machine == 'arm64' or sys_platform == 'win32' and platform_machine == 'ARM64'"}