Merge pull request #162 from awu-labelbox/awu/mea-create-model

Add create model and create model run

Author: awu-labelbox

labelbox/__init__.py

@@ -19,3 +19,4 @@
 from labelbox.schema.role import Role, ProjectRole
 from labelbox.schema.invite import Invite, InviteLimit
 from labelbox.schema.model_run import ModelRun
+from labelbox.schema.model import Model

labelbox/client.py

@@ -533,3 +533,19 @@ def get_models(self, where=None):
             An iterable of Models (typically a PaginatedCollection).
         """
         return self._get_all(Model, where, filter_deleted=False)
+
+    def create_model(self, name, ontology_id):
+        """ Creates a Model object on the server.
+
+        >>> model = client.create_model(<model_name>, <ontology_id>)
+
+        Args:
+            name (string): Name of the model
+            ontology_id (string): ID of the related ontology
+        Returns:
+            A new Model object.
+        Raises:
+            InvalidAttributeError: If the Model type does not contain
+                any of the attribute names given in kwargs.
+        """
+        return self._create(Model, {"name": name, "ontology_id": ontology_id})

labelbox/schema/__init__.py

@@ -3,16 +3,17 @@
 import labelbox.schema.benchmark
 import labelbox.schema.data_row
 import labelbox.schema.dataset
+import labelbox.schema.invite
 import labelbox.schema.label
 import labelbox.schema.labeling_frontend
+import labelbox.schema.model
+import labelbox.schema.model_run
+import labelbox.schema.ontology
 import labelbox.schema.organization
+import labelbox.schema.prediction
 import labelbox.schema.project
 import labelbox.schema.review
+import labelbox.schema.role
 import labelbox.schema.task
 import labelbox.schema.user
 import labelbox.schema.webhook
-import labelbox.schema.prediction
-import labelbox.schema.ontology
-import labelbox.schema.invite
-import labelbox.schema.role
-import labelbox.schema.model_run

labelbox/schema/model.py

@@ -1,5 +1,6 @@
+from labelbox.orm import query
 from labelbox.orm.db_object import DbObject
-from labelbox.orm.model import Field, Relationship
+from labelbox.orm.model import Entity, Field, Relationship
 
 
 class Model(DbObject):
@@ -16,3 +17,26 @@ class Model(DbObject):
 
     model_runs = Relationship.ToMany("ModelRun", False)
     ontology = Relationship.ToOne("Ontology", False)
+
+    def create_model_run(self, name):
+        """ Creates a model run belonging to this model.
+
+        Args:
+            name (string): The name for the model run.
+        Returns:
+            ModelRun, the created model run.
+        """
+        name_param = "name"
+        model_id_param = "modelId"
+        ModelRun = Entity.ModelRun
+        query_str = """mutation CreateModelRunPyApi($%s: String!, $%s: ID!) {
+            createModelRun(data: {name: $%s, modelId: $%s}) {%s}}""" % (
+                name_param, model_id_param,
+                name_param, model_id_param,
+                query.results_query_part(ModelRun)
+            )
+        res = self.client.execute(query_str, {
+            name_param: name,
+            model_id_param: self.uid
+        })
+        return ModelRun(self.client, res["createModelRun"])

tests/integration/test_model.py

@@ -0,0 +1,23 @@
+from labelbox import Model
+
+def test_model(client, configured_project, rand_gen):
+    before = list(client.get_models())
+    for m in before:
+        assert isinstance(m, Model)
+
+    ontology = configured_project.ontology
+
+    data = {"name": rand_gen(str), "ontology_id": ontology.uid}
+    model = client.create_model(data["name"], data["ontology_id"])
+    assert model.name == data["name"]
+    assert model.ontology.id == data["ontology_id"]
+
+    after = list(client.get_models())
+    assert len(after) == len(before) + 1
+    assert model in after
+
+    model = client.get_model(model.uid)
+    assert model.name == data["name"]
+    assert model.ontology.id == data["ontology_id"]
+
+