Updated model to include trained_slice_id (#421)

Author: schumry

CHANGELOG.md

@@ -6,6 +6,11 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 
+## [0.16.18](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.16.18) - 2024-02-06
+
+### Added
+- Add the ability to add and remove `trained_slice_id` to a model
+
 ## [0.16.17](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.16.17) - 2024-01-29
 
 ### Fixes

nucleus/__init__.py

@@ -111,6 +111,7 @@
     MESSAGE_KEY,
     MODEL_RUN_ID_KEY,
     MODEL_TAGS_KEY,
+    MODEL_TRAINED_SLICE_IDS_KEY,
     NAME_KEY,
     NUCLEUS_ENDPOINT,
     POINTS_KEY,
@@ -247,6 +248,7 @@ def models(self) -> List[Model]:
                 metadata=model["metadata"] or None,
                 client=self,
                 tags=model.get(MODEL_TAGS_KEY, []),
+                trained_slice_ids=model.get(MODEL_TRAINED_SLICE_IDS_KEY, None),
             )
             for model in model_objects["models"]
         ]
@@ -560,6 +562,7 @@ def create_model(
         metadata: Optional[Dict] = None,
         bundle_name: Optional[str] = None,
         tags: Optional[List[str]] = None,
+        trained_slice_ids: Optional[List[str]] = None,
     ) -> Model:
         """Adds a :class:`Model` to Nucleus.
 
@@ -579,7 +582,12 @@ def create_model(
         """
         response = self.make_request(
             construct_model_creation_payload(
-                name, reference_id, metadata, bundle_name, tags
+                name,
+                reference_id,
+                metadata,
+                bundle_name,
+                tags,
+                trained_slice_ids,
             ),
             "models/add",
         )
@@ -595,6 +603,7 @@ def create_model(
             bundle_name=bundle_name,
             client=self,
             tags=tags,
+            trained_slice_ids=trained_slice_ids,
         )
 
     def create_launch_model(
@@ -603,6 +612,7 @@ def create_launch_model(
         reference_id: str,
         bundle_args: Dict[str, Any],
         metadata: Optional[Dict] = None,
+        trained_slice_ids: Optional[List[str]] = None,
     ) -> Model:
         """
         Adds a :class:`Model` to Nucleus, as well as a Launch bundle from a given function.
@@ -694,6 +704,7 @@ def create_launch_model(
             reference_id,
             metadata,
             bundle_name,
+            trained_slice_ids=trained_slice_ids,
         )
 
     def create_launch_model_from_dir(
@@ -702,6 +713,7 @@ def create_launch_model_from_dir(
         reference_id: str,
         bundle_from_dir_args: Dict[str, Any],
         metadata: Optional[Dict] = None,
+        trained_slice_ids: Optional[List[str]] = None,
     ) -> Model:
         """
         Adds a :class:`Model` to Nucleus, as well as a Launch bundle from a directory.
@@ -816,6 +828,7 @@ def create_launch_model_from_dir(
             reference_id,
             metadata,
             bundle_name,
+            trained_slice_ids=trained_slice_ids,
         )
 
     @deprecated(

nucleus/constants.py

@@ -106,6 +106,7 @@
 METADATA_KEY = "metadata"
 MODEL_BUNDLE_NAME_KEY = "bundle_name"
 MODEL_TAGS_KEY = "tags"
+MODEL_TRAINED_SLICE_IDS_KEY = "trained_slice_ids"
 MODEL_ID_KEY = "model_id"
 MODEL_RUN_ID_KEY = "model_run_id"
 MODEL_PREDICTION_ID_KEY = "model_prediction_id"

nucleus/model.py

@@ -3,7 +3,13 @@
 import requests
 
 from .async_job import AsyncJob
-from .constants import METADATA_KEY, MODEL_TAGS_KEY, NAME_KEY, REFERENCE_ID_KEY
+from .constants import (
+    METADATA_KEY,
+    MODEL_TAGS_KEY,
+    MODEL_TRAINED_SLICE_IDS_KEY,
+    NAME_KEY,
+    REFERENCE_ID_KEY,
+)
 from .dataset import Dataset
 from .model_run import ModelRun
 from .prediction import (
@@ -101,6 +107,7 @@ def __init__(
         client,
         bundle_name=None,
         tags=None,
+        trained_slice_ids=None,
     ):
         self.id = model_id
         self.name = name
@@ -109,9 +116,10 @@ def __init__(
         self.bundle_name = bundle_name
         self.tags = tags if tags else []
         self._client = client
+        self.trained_slice_ids = trained_slice_ids if trained_slice_ids else []
 
     def __repr__(self):
-        return f"Model(model_id='{self.id}', name='{self.name}', reference_id='{self.reference_id}', metadata={self.metadata}, bundle_name={self.bundle_name}, tags={self.tags}, client={self._client})"
+        return f"Model(model_id='{self.id}', name='{self.name}', reference_id='{self.reference_id}', metadata={self.metadata}, bundle_name={self.bundle_name}, tags={self.tags}, client={self._client}, trained_slice_ids={self.trained_slice_ids})"
 
     def __eq__(self, other):
         return (
@@ -120,6 +128,7 @@ def __eq__(self, other):
             and (self.metadata == other.metadata)
             and (self._client == other._client)
             and (self.bundle_name == other.bundle_name)
+            and (self.trained_slice_ids == other.trained_slice_ids)
         )
 
     def __hash__(self):
@@ -134,6 +143,8 @@ def from_json(cls, payload: dict, client):
             reference_id=payload["ref_id"],
             metadata=payload["metadata"] or None,
             client=client,
+            tags=payload.get(MODEL_TAGS_KEY, None),
+            trained_slice_ids=payload.get(MODEL_TRAINED_SLICE_IDS_KEY, None),
         )
 
     def create_run(
@@ -242,7 +253,9 @@ def add_tags(self, tags: List[str]):
         )
 
         if response.ok:
-            self.tags.extend(tags)
+            for tag in tags:
+                if tag not in self.tags:
+                    self.tags.append(tag)
 
         return response.json()
 
@@ -269,3 +282,55 @@ def remove_tags(self, tags: List[str]):
             self.tags = list(filter(lambda t: t not in tags, self.tags))
 
         return response.json()
+
+    def add_trained_slice_ids(self, slice_ids: List[str]):
+        """Add trained slice id(s) to the model. ::
+
+            import nucleus
+            client = nucleus.NucleusClient("YOUR_SCALE_API_KEY")
+            model = client.list_models()[0]
+
+            model.add_trained_slice_ids(["slc_...", "slc_..."])
+
+        Args:
+            slice_ids: list of trained slice ids
+        """
+        response: requests.Response = self._client.make_request(
+            {MODEL_TRAINED_SLICE_IDS_KEY: slice_ids},
+            f"model/{self.id}/trainedSliceId",
+            requests_command=requests.post,
+            return_raw_response=True,
+        )
+
+        if response.ok:
+            for slice_id in slice_ids:
+                if slice_id not in self.trained_slice_ids:
+                    self.trained_slice_ids.append(slice_id)
+
+        return response.json()
+
+    def remove_trained_slice_ids(self, slide_ids: List[str]):
+        """Remove trained slice id(s) from the model. ::
+
+            import nucleus
+            client = nucleus.NucleusClient("YOUR_SCALE_API_KEY")
+            model = client.list_models()[0]
+
+            model.remove_trained_slice_ids(["slc_...", "slc_..."])
+
+        Args:
+            slice_ids: list of trained slice ids to remove
+        """
+        response: requests.Response = self._client.make_request(
+            {MODEL_TRAINED_SLICE_IDS_KEY: slide_ids},
+            f"model/{self.id}/trainedSliceId",
+            requests_command=requests.delete,
+            return_raw_response=True,
+        )
+
+        if response.ok:
+            self.trained_slice_ids = list(
+                filter(lambda t: t not in slide_ids, self.trained_slice_ids)
+            )
+
+        return response.json()

nucleus/payload_constructor.py

@@ -19,6 +19,7 @@
     MODEL_BUNDLE_NAME_KEY,
     MODEL_ID_KEY,
     MODEL_TAGS_KEY,
+    MODEL_TRAINED_SLICE_IDS_KEY,
     NAME_KEY,
     REFERENCE_ID_KEY,
     SCENES_KEY,
@@ -137,13 +138,16 @@ def construct_model_creation_payload(
     metadata: Optional[Dict],
     bundle_name: Optional[str],
     tags: Optional[List[str]],
+    trained_slice_ids: Optional[List[str]],
 ) -> dict:
     payload = {
         NAME_KEY: name,
         REFERENCE_ID_KEY: reference_id,
         METADATA_KEY: metadata if metadata else {},
     }
 
+    if trained_slice_ids:
+        payload[MODEL_TRAINED_SLICE_IDS_KEY] = trained_slice_ids
     if bundle_name:
         payload[MODEL_BUNDLE_NAME_KEY] = bundle_name
     if tags:

pyproject.toml

@@ -25,7 +25,7 @@ ignore = ["E501", "E741", "E731", "F401"]  # Easy ignore for getting it running
 
 [tool.poetry]
 name = "scale-nucleus"
-version = "0.16.17"
+version = "0.16.18"
 description = "The official Python client library for Nucleus, the Data Platform for AI"
 license =  "MIT"
 authors = ["Scale AI Nucleus Team <nucleusapi@scaleapi.com>"]