Use new resolver

Author: mnoszczak

labelbox/orm/model.py

@@ -371,6 +371,7 @@ class Entity(metaclass=EntityMeta):
         LabelingFrontendOptions: Type[labelbox.LabelingFrontendOptions]
         Label: Type[labelbox.Label]
         MEAPredictionImport: Type[labelbox.MEAPredictionImport]
+        MALPredictionImport: Type[labelbox.MALPredictionImport]
         Invite: Type[labelbox.Invite]
         InviteLimit: Type[labelbox.InviteLimit]
         ProjectRole: Type[labelbox.ProjectRole]

labelbox/schema/annotation_import.py

@@ -401,6 +401,32 @@ def create_from_url(cls, client: "labelbox.Client", project_id: str,
         else:
             raise ValueError(f"Url {url} is not reachable")
 
+    @classmethod
+    def create_for_model_run_data_rows(cls, client: "labelbox.Client",
+                                       model_run_id: str,
+                                       data_row_ids: List[str], project_id: str,
+                                       name: str) -> "MALPredictionImport":
+        """
+        Create an MAL prediction import job from a list of data row ids of a specific model run
+
+        Args:
+            client: Labelbox Client for executing queries
+            data_row_ids: A list of data row ids
+            model_run_id: model run id
+        Returns:
+            MALPredictionImport
+        """
+        query_str = cls._get_model_run_data_rows_mutation()
+        return cls(
+            client,
+            client.execute(query_str,
+                           params={
+                               "dataRowIds": data_row_ids,
+                               "modelRunId": model_run_id,
+                               "projectId": project_id,
+                               "name": name
+                           })["createMalPredictionImportForModelRunDataRows"])
+
     @classmethod
     def from_name(cls,
                   client: "labelbox.Client",
@@ -445,6 +471,17 @@ def _get_url_mutation(cls) -> str:
             }) {%s}
         }""" % query.results_query_part(cls)
 
+    @classmethod
+    def _get_model_run_data_rows_mutation(cls) -> str:
+        return """mutation createMalPredictionImportForModelRunDataRows($projectId : ID!, $dataRowIds: [ID!]!, $name: String!, $modelRunId: ID!) {
+            createMalPredictionImportForModelRunDataRows(data: {
+                name: $name
+                modelRunId: $modelRunId
+                dataRowIds: $dataRowIds
+                projectId: $projectId
+            }) {%s}
+        }""" % query.results_query_part(cls)
+
     @classmethod
     def _get_file_mutation(cls) -> str:
         return """mutation createMALPredictionImportByFilePyApi($projectId : ID!, $name: String!, $file: Upload!, $contentLength: Int!) {

labelbox/schema/model_run.py

@@ -12,7 +12,6 @@
 from labelbox.orm.query import results_query_part
 from labelbox.orm.model import Field, Relationship, Entity
 from labelbox.orm.db_object import DbObject, experimental
-from labelbox.data.annotation_types import LabelList
 
 if TYPE_CHECKING:
     from labelbox import MEAPredictionImport
@@ -128,11 +127,11 @@ def _wait_until_done(self, status_fn, timeout_seconds=120, sleep_time=5):
     def upsert_predictions_and_send_to_project(
         self,
         name: str,
-        predictions: LabelList,
+        predictions: Union[str, Path, Iterable[Dict]],
         project_id: str,
         priority: Optional[int] = 5,
     ) -> 'MEAPredictionImport':  # type: ignore
-        """ Upload predictions and creates a batch import to project.
+        """ Upload predictions and create a batch import to project.
         Args:
             name (str): name of the AnnotationImport job as well as the name of the batch import
             predictions (Iterable):
@@ -142,16 +141,21 @@ def upsert_predictions_and_send_to_project(
         Returns:
             (AnnotationImport, Project)
         """
-        data_rows_set = set(
-            map(lambda x: x.data_row.id, predictions)[:DATAROWS_IMPORT_LIMIT])
-        data_rows = list(data_rows_set)
+        kwargs = dict(client=self.client, model_run_id=self.uid, name=name)
         project = self.client.get_project(project_id)
-        batch = project.create_batch(name, data_rows, priority)
-
-        predictions_for_data_rows = filter(
-            lambda x: x.data_row.id in data_rows_set, predictions)
-
-        return (self.add_predictions(name, predictions_for_data_rows), batch)
+        import_job = self.add_predictions(name, predictions)
+        prediction_statuses = import_job.statuses()
+        mea_to_mal_data_rows_set = set([
+            row['dataRow']['id']
+            for row in prediction_statuses
+            if row['status'] == 'SUCCESS'
+        ])
+        mea_to_mal_data_rows = list(
+            mea_to_mal_data_rows_set)[:DATAROWS_IMPORT_LIMIT]
+        batch = project.create_batch(name, mea_to_mal_data_rows, priority)
+        mal_prediction_import = Entity.MALPredictionImport.create_for_model_run_data_rows(
+            data_row_ids=mea_to_mal_data_rows, project_id=project_id, **kwargs)
+        return mea_prediction_import, batch, mal_prediction_import
 
     def add_predictions(
         self,