From 721ea2bd7a572b69c6dcb1228b691c57d947c0c5 Mon Sep 17 00:00:00 2001
From: luandalmazo <luandalmazo@gmail.com>
Date: Wed, 1 Jul 2026 10:58:47 -0300
Subject: [PATCH] add _log_predictions

---
 datamint/mlflow/flavors/model.py              | 36 ++++++++++++++++++-
 .../01_fracatlas_classification.ipynb         |  5 +--
 2 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/datamint/mlflow/flavors/model.py b/datamint/mlflow/flavors/model.py
index 0797e6dc..6f62c6c3 100644
--- a/datamint/mlflow/flavors/model.py
+++ b/datamint/mlflow/flavors/model.py
@@ -214,8 +214,42 @@ def predict(  # pyright: ignore[reportIncompatibleMethodOverride]
         Routes to the appropriate handler based on ``params['mode']``.
         **Do not override** — implement :meth:`predict_default` (or other
         ``predict_*`` hooks) instead.
+
+        Pass ``params={'log_predictions': True, 'model_name': ...}`` to also upload
+        the resulting annotations to the Datamint server, tagged ``source='model_deploy'``.
+        This is meant for local inference against a model loaded via
+        :func:`datamint.mlflow.flavors.load_model` (testing a registered model outside
+        of ``Trainer.fit()``), as opposed to predictions made automatically during
+        training, which are tagged ``source='model_pipeline'`` instead.
         """
-        return self._router.dispatch(model_input, params or {})
+        params = dict(params or {})
+        log_predictions = params.pop('log_predictions', False)
+        model_name = params.pop('model_name', None)
+
+        result = self._router.dispatch(model_input, params)
+
+        if log_predictions:
+            self._log_predictions(model_input, result, model_name=model_name)
+
+        return result
+
+    def _log_predictions(
+        self,
+        resources: list[BaseResource],
+        predictions: PredictionResult,
+        model_name: str | None,
+    ) -> None:
+        """Upload predictions as annotations tagged ``source='model_deploy'``. """
+        for resource, preds in zip(resources, predictions):
+            if not preds:
+                continue
+            try:
+                annotations_api = resource._api._api_instance.annotations
+                annotations_api.upload_predictions(resource, preds, model_name=model_name, source='model_deploy')
+            except Exception as e:
+                logger.warning(
+                    "Failed to log predictions for resource %s: %s", getattr(resource, 'id', resource), e
+                )
 
     def get_supported_modes(self) -> list[str]:
         """Return the list of prediction modes supported by this model."""
diff --git a/notebooks/06_end_to_end/slice_based/01_fracatlas_classification.ipynb b/notebooks/06_end_to_end/slice_based/01_fracatlas_classification.ipynb
index 7badb8be..52d019f8 100644
--- a/notebooks/06_end_to_end/slice_based/01_fracatlas_classification.ipynb
+++ b/notebooks/06_end_to_end/slice_based/01_fracatlas_classification.ipynb
@@ -540,9 +540,10 @@
     "# Load the registered model\n",
     "model_loaded = datamint_flavor.load_model(f'models:/{PROJECT_NAME}/latest')\n",
     "\n",
-    "# Predict on a resource from the dataset\n",
+    "# Predict on a resource from the dataset, and log the prediction back to Datamint\n",
+    "# as an annotation tagged source='model_deploy' (set log_predictions=True)\n",
     "r = trainer.dataset[0]['resource']\n",
-    "prediction = model_loaded.predict([r])\n",
+    "prediction = model_loaded.predict([r], params={'log_predictions': True, 'model_name': PROJECT_NAME})\n",
     "print(prediction)"
    ]
   },