add model version for prediction (#363)

openfoodfacts · Nov 21, 2024 · 1db1a63 · 1db1a63
1 parent fb59153
commit 1db1a63
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/ml_utils/ml_utils_cli/cli/apps/projects.py b/ml_utils/ml_utils_cli/cli/apps/projects.py
@@ -213,6 +213,10 @@ def add_prediction(
                 help="Raise an error if image download fails"
             ),
         ] = True,
+    model_version: Annotated[
+        Optional[str],
+        typer.Option(help="Model version to use for the prediction"),
+    ] = None,
 ):
     """Add predictions as pre-annotations to Label Studio tasks,
     for an object detection model running on Triton Inference Server."""
@@ -297,6 +301,7 @@ def add_prediction(
                 ls.predictions.create(
                     task=task.id,
                     result=label_studio_result,
+                    model_version=model_version,
                 )