orient_express 0.3.2__tar.gz → 0.3.4__tar.gz

{orient_express-0.3.2 → orient_express-0.3.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: orient_express
-Version: 0.3.2
+Version: 0.3.4
 Summary: A library to simplify model deployment to Vertex AI
 Author: Alexey Zankevich
 Author-email: alex.zankevich@shiftsmart.com
@@ -12,6 +12,7 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Requires-Dist: google-cloud-aiplatform
 Requires-Dist: google-cloud-storage
+Requires-Dist: jupyter (>=1.1.1,<2.0.0)
 Requires-Dist: pandas
 Description-Content-Type: text/markdown
 

{orient_express-0.3.2 → orient_express-0.3.4}/orient_express/model_wrapper.py

@@ -1,4 +1,6 @@
 import os
+from typing import Optional
+
 import joblib
 import logging
 import pandas as pd
@@ -9,20 +11,20 @@ from google.cloud import aiplatform
 class ModelExpress:
     def __init__(
         self,
-        model_name,
-        project_name,
-        bucket_name,
-        model_version=None,
-        model=None,
-        region="us-central1",
-        serialized_model_path="model.joblib",
-        serving_container_image_uri="us-west1-docker.pkg.dev/shiftsmart-api/orient-express/xgboost-scikit-learn:latest",
-        serving_container_predict_route="/v1/models/orient-express-model:predict",
-        serving_container_health_route="/v1/models/orient-express-model",
-        endpoint_name=None,
+        model_name: str,
+        project_name: str,
+        bucket_name: str,
+        model_version: Optional[int] = None,
+        model: object = None,
+        region: str = "us-central1",
+        serialized_model_path: str = "model.joblib",
+        serving_container_image_uri: str = "us-west1-docker.pkg.dev/shiftsmart-api/orient-express/xgboost-scikit-learn:latest",
+        serving_container_predict_route: str = "/v1/models/orient-express-model:predict",
+        serving_container_health_route: str = "/v1/models/orient-express-model",
+        endpoint_name: Optional[str] = None,
         machine_type="n1-standard-4",
-        min_replica_count=1,
-        max_replica_count=1,
+        min_replica_count: int = 1,
+        max_replica_count: int = 1,
     ):
         self.model = model
         self.model_name = model_name
@@ -44,15 +46,22 @@ class ModelExpress:
         else:
             self.endpoint_name = endpoint_name
 
+        self._vertex_initialized = False
+
     def colab_auth(self):
         from google.colab import auth
 
         auth.authenticate_user()
 
     def _vertex_init(self):
-        aiplatform.init(project=self.project_name, location=self.region)
-
-    def get_latest_vertex_model(self, model_name):
+        if not self._vertex_initialized:
+            aiplatform.init(project=self.project_name, location=self.region)
+            self._vertex_initialized = True
+
+    def get_latest_vertex_model(self, model_name: str):
+        """If there are a few models with the same name, load the most recent one.
+        It's highly recommended to keep only 1 model with the same name to avoid the confusion
+        """
         self._vertex_init()
 
         # Search for models with the specified display name
@@ -89,7 +98,7 @@ class ModelExpress:
 
         return self.create_model_version(new_version, last_model)
 
-    def get_artifacts_path(self, version, file_name=None):
+    def get_artifacts_path(self, version: int, file_name: str = None):
         dir_name = f"models/{self.model_name}/{version}"
         if file_name:
             return f"{dir_name}/{file_name}"
@@ -152,7 +161,9 @@ class ModelExpress:
             traffic_percentage=100,
         )
 
-    def remote_predict(self, input_df):
+    def remote_predict(self, input_df: pd.DataFrame):
+        self._vertex_init()
+
         if not self.endpoint:
             endpoint = self.get_endpoint()
             if not endpoint:
@@ -165,27 +176,31 @@ class ModelExpress:
         predictions = self.endpoint.predict(instances=instances)
         return predictions.predictions
 
-    def local_predict(self, input_df):
+    def local_predict(self, input_df: pd.DataFrame):
+        self._vertex_init()
+
         if not self.model:
             self.load_model_from_registry()
 
         return self.model.predict(input_df)
 
-    def local_predict_proba(self, input_df):
+    def local_predict_proba(self, input_df: pd.DataFrame):
         if not self.model:
             self.load_model_from_registry()
 
         return self.model.predict_proba(input_df)
 
     def load_model_from_registry(self):
+        self._vertex_init()
+
+        vertex_model = self.get_latest_vertex_model(self.model_name)
+
         if self.model_version:
+            # reload the model using a specific model version
             vertex_model = aiplatform.Model(
-                model_name=self.model_name, version=self.model_version
+                model_name=vertex_model.resource_name, version=str(self.model_version)
             )
 
-        else:
-            vertex_model = self.get_latest_vertex_model(self.model_name)
-
         if not vertex_model:
             raise Exception(f"Model '{self.model_name}' not found in the registry.")
 
@@ -194,7 +209,7 @@ class ModelExpress:
 
         self.model = joblib.load(self.serialized_model_path)
 
-    def download_artifacts(self, artifact_uri):
+    def download_artifacts(self, artifact_uri: str):
         storage_client = storage.Client()
         bucket_name, artifact_path = artifact_uri.replace("gs://", "").split("/", 1)
         bucket = storage_client.bucket(bucket_name)

orient_express-0.3.4/orient_express/sklearn_pipeline.py

@@ -0,0 +1,81 @@
+from sklearn.base import BaseEstimator, TransformerMixin
+from sklearn.preprocessing import LabelEncoder
+
+
+class LabelEncoderTransformer(BaseEstimator, TransformerMixin):
+    """
+    A wrapper class that integrates a model and a label encoder to handle
+    encoded predictions and their corresponding probabilities.
+
+    Attributes:
+        model (object): A machine learning model with `fit`, `predict`, and `predict_proba` methods.
+        label_encoder (object): A label encoder that maps encoded class labels to their original string labels.
+
+    Methods:
+        fit(X, y):
+            Trains the model on the provided features and labels.
+        predict(X):
+            Predicts the classes for the given features and returns the original string labels.
+        predict_proba(X):
+            Returns the probabilities for each class for the given features, along with their original class labels.
+    """
+
+    def __init__(self, model: BaseEstimator, label_encoder: LabelEncoder):
+        """
+        Initializes the LabelEncoderTransformer.
+
+        Args:
+            model (BaseEstimator): A machine learning model.
+            label_encoder (LabelEncoder): A label encoder with `fit` and `inverse_transform` methods.
+        """
+        self.model = model
+        self.label_encoder = label_encoder
+
+    def fit(self, X, y):
+        """
+        Fits the model to the training data.
+
+        Args:
+            X (array-like): Feature matrix.
+            y (array-like): Target labels.
+
+        Returns:
+            self: Fitted LabelEncoderTransformer instance.
+        """
+        self.model.fit(X, y)
+        return self
+
+    def predict(self, X):
+        """
+        Predicts the target labels and returns the original string labels.
+
+        Args:
+            X (array-like): Feature matrix for predictions.
+
+        Returns:
+            array-like: Predicted class labels in their original form.
+        """
+        encoded_predictions = self.model.predict(X)
+        return self.label_encoder.inverse_transform(encoded_predictions)
+
+    def predict_proba(self, X):
+        """
+        Returns class probabilities along with their original labels.
+
+        Args:
+            X (array-like): Feature matrix for predictions.
+
+        Returns:
+            list: A list of lists, where each inner list contains `[class_name, probability]` pairs.
+        """
+        # Get raw probabilities
+        probabilities = self.model.predict_proba(X)
+
+        # Get all class labels
+        class_names = self.label_encoder.classes_
+        # Combine class names with probabilities
+        combined_output = [
+            [[class_name, prob] for class_name, prob in zip(class_names, sample_probs)]
+            for sample_probs in probabilities
+        ]
+        return combined_output

{orient_express-0.3.2 → orient_express-0.3.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "orient_express"
-version = "0.3.2"
+version = "0.3.4"
 description = "A library to simplify model deployment to Vertex AI"
 authors = ["Alexey Zankevich <alex.zankevich@shiftsmart.com>"]
 readme = "README.md"
@@ -10,6 +10,7 @@ python = ">=3.9,<3.13"
 google-cloud-aiplatform = "*"
 google-cloud-storage = "*"
 pandas = "*"
+jupyter = "^1.1.1"
 
 [tool.poetry.dev-dependencies]
 black = "24.10.0"