{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "556c43fc",
   "metadata": {},
   "source": [
    "# Model Registry Tutorial\n",
    "\n",
    "`api.models` is a facade over Datamint's MLflow-backed model registry. It wraps MLflow's\n",
    "`RegisteredModel` / `ModelVersion` objects in plain Python objects (`Model`, `ModelVersion`),\n",
    "so you can register, list, and inspect models without learning MLflow's object model.\n",
    "\n",
    "This notebook covers:\n",
    "- Registering a model with `api.models.create()`\n",
    "- Listing and finding models with `api.models.get_list()` / `get_by_name()`\n",
    "- Filtering to deployed models with `only_deployed=True`\n",
    "- Inspecting a model's versions with `model.get_versions()` / `get_latest_version()`\n",
    "- Reading what a version was trained for: `get_task_type()`, `get_supported_modes()`, `get_annotation_specs()`\n",
    "- Reading training metrics with `get_metrics()`\n",
    "- Checking deployment status with `is_deployed()`\n",
    "- Finding which projects a model is associated with, with `model.get_projects()`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff7cc0e3",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "Update `datamint` first if needed, then configure your API key."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d323bd4",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -U datamint --quiet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef4c457d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datamint import Api\n",
    "\n",
    "api = Api()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "547f3483",
   "metadata": {},
   "source": [
    "## Register A Model\n",
    "\n",
    "Registering ahead of training gives you a stable name to reference later. `create()` returns the\n",
    "existing model instead of raising when one with this name is already registered\n",
    "(`exists_ok=True` by default)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05e434aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_NAME = \"tutorial_model_registry_demo\"\n",
    "\n",
    "model = api.models.create(MODEL_NAME, description=\"Model created for the model registry tutorial\")\n",
    "model.name, model.description"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cc966eb",
   "metadata": {},
   "source": [
    "## List And Find Models\n",
    "\n",
    "`get_list()` returns every registered model; `get_by_name()` returns a single one, or `None` if it\n",
    "doesn't exist."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2571c3f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_models = api.models.get_list()\n",
    "print(f\"Registered models: {[m.name for m in all_models]}\")\n",
    "\n",
    "found = api.models.get_by_name(MODEL_NAME)\n",
    "missing = api.models.get_by_name(\"does-not-exist\")\n",
    "found.name, missing"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c30b3598",
   "metadata": {},
   "source": [
    "## Filter To Deployed Models\n",
    "\n",
    "Pass `only_deployed=True` to skip models that don't have a deployed image yet. See\n",
    "`05_deployment/01_deploy_registered_model.ipynb` for how to deploy one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c605a73",
   "metadata": {},
   "outputs": [],
   "source": [
    "deployed_models = api.models.get_list(only_deployed=True)\n",
    "print(f\"Deployed models: {[m.name for m in deployed_models]}\")\n",
    "\n",
    "model.is_deployed()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fab1e7a0",
   "metadata": {},
   "source": [
    "## Inspect Versions Of A Trained Model\n",
    "\n",
    "The rest of this notebook needs a model with at least one version behind it, typically one\n",
    "registered by a Datamint trainer (see `06_end_to_end`) or by\n",
    "`04_experiment_tracking/01_mlflow_manual_logging.ipynb`. Replace `MODEL_NAME` below with one\n",
    "you've already trained."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b0203fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_NAME = \"FracAtlas_adapted\"  # replace with a model you've already registered/trained\n",
    "\n",
    "model = api.models.get_by_name(MODEL_NAME)\n",
    "if model is None:\n",
    "    raise ValueError(f\"Model '{MODEL_NAME}' was not found. Train one first, e.g. via the 06_end_to_end notebooks.\")\n",
    "\n",
    "versions = model.get_versions()\n",
    "print(f\"{MODEL_NAME} has {len(versions)} version(s)\")\n",
    "\n",
    "latest = model.get_latest_version()\n",
    "latest.version, latest.run_id"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3045992",
   "metadata": {},
   "source": [
    "## What A Version Was Trained For\n",
    "\n",
    "If the version was logged with the `datamint` MLflow flavor (true for anything trained through a\n",
    "Datamint trainer), you can read the task type, supported prediction modes, and annotation specs\n",
    "straight from the model artifact, without inspecting the training run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ef28bd9",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Task type:\", latest.get_task_type())\n",
    "print(\"Supported modes:\", latest.get_supported_modes())\n",
    "print(\"Annotation specs:\", latest.get_annotation_specs())\n",
    "\n",
    "# Model.get_supported_modes() is a shortcut that delegates to the latest version\n",
    "model.get_supported_modes() == latest.get_supported_modes()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cdfdcb81",
   "metadata": {},
   "source": [
    "## Training Metrics\n",
    "\n",
    "`get_metrics()` reads the metrics logged for this version's training run. It returns `{}` instead\n",
    "of raising when there's no training run behind the version, for example a model registered from\n",
    "outside Datamint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f984040b",
   "metadata": {},
   "outputs": [],
   "source": [
    "latest.get_metrics()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee63e20b",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.get_metrics()  # shortcut, uses the latest version"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4281df3",
   "metadata": {},
   "source": [
    "## Find Which Projects A Model Belongs To\n",
    "\n",
    "A registered model can be associated with one or more Datamint projects. `model.get_projects()`\n",
    "returns the full `Project` objects, resolved via `api.projects`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac7a0fe2",
   "metadata": {},
   "outputs": [],
   "source": [
    "projects = model.get_projects()\n",
    "print(f\"{MODEL_NAME} belongs to {len(projects)} project(s)\")\n",
    "\n",
    "for project in projects:\n",
    "    print(project.id, project.name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5edb693f",
   "metadata": {},
   "source": [
    "## Next Steps\n",
    "\n",
    "Once you have a model version you're happy with, alias it and deploy it, see\n",
    "`05_deployment/01_deploy_registered_model.ipynb`. Registered models are also created\n",
    "automatically when you pass `--ai-model <name>` to `datamint upload` with a name that doesn't\n",
    "exist yet."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}