feat(be): adapt force_train for multiclass workspaces

label_sleuth/app.py

@@ -1046,25 +1046,29 @@ def get_elements_to_label(workspace_id):
 @login_if_required
 @validate_category_id
 @validate_workspace_id
-def force_train_for_category(workspace_id):
+def force_train(workspace_id):
     """
     This call is used for manually triggering a new Iteration flow.
 
     :param workspace_id:
     :request_arg category_id:
     """
-    category_id = int(request.args['category_id'])
+    is_multiclass = request.args.get('mode') == WorkspaceModelType.MultiClass.name
+    category_id = request.args.get('category_id')
+    if category_id is not None:
+        category_id = int(category_id)    
+
     dataset_name = curr_app.orchestrator_api.get_dataset_name(workspace_id)
 
-    model_id = curr_app.orchestrator_api.train_if_recommended(workspace_id, category_id, force=True)
+    curr_app.orchestrator_api.train_if_recommended(workspace_id, category_id, force=True)
 
     labeling_counts = curr_app.orchestrator_api.get_label_counts(workspace_id, dataset_name, category_id)
-    logging.info(f"force training a new model in workspace '{workspace_id}' for category '{category_id}', "
-                 f"model id: {model_id}")
+    to_log_message = f"Force training a new model in workspace '{workspace_id}'"
+    if not is_multiclass: to_log_message += f" for category '{category_id}'"
+    logging.info(to_log_message)
 
     return jsonify({
-        "labeling_counts": labeling_counts,
-        "model_id": model_id
+        "labeling_counts": labeling_counts
     })
 
 

label_sleuth/test_app.py

@@ -753,7 +753,6 @@ def test_full_flow_multiclass(self):
         for element in res.get_json()["elements"]:
             self.assertEqual(0, element["model_predictions"], msg=f"element {element} expected prediction is 0 but got {element['model_predictions']}")
 
-
         # delete category 0
         res = self.client.delete(f"/workspace/{workspace_name}/category/0?mode=MultiClass", headers=HEADERS)
         self.assertEqual(200, res.status_code, msg="Failed to delete category 0 in multiclass workspace")
@@ -771,6 +770,9 @@ def test_full_flow_multiclass(self):
         self.assertEqual({'elements': [], 'hit_count': 0},
                          res.get_json(), msg="labeled elements for category 0 were deleted should not exist in this workspace")
 
+        # force model training
+        res = self.client.get(f"/workspace/{workspace_name}/force_train?mode=MultiClass", headers=HEADERS)
+        self.assertEqual(200, res.status_code, msg="Failed to force train in multiclass workspace")
 
         # delete workspace
         res = self.client.delete(f"/workspace/{workspace_name}",