Lean into model_required CLI decorator

Why these changes are being introduced:

Many of the  CLI commands will require an embedding class and model to work.
A decorator was created originally that injected a --model-uri CLI argument,
but it also provides a place to load the class itself and become more of a
middleware.

How this addresses that need:
Updates the model_required decorator to also load the embedding model class.
This DRY's up the CLI commands that use it and centralizes that logic and
conventions for the CLI argument, env vars, and whatnot.

Lastly, it is now required to include a 'model_path' when instantiating
a model class instance, and this location is used for both download and
load.

Side effects of this change:
* None

Relevant ticket(s):
* https://mitlibraries.atlassian.net/browse/USE-112

Author: ghukill

Dockerfile

@@ -19,12 +19,12 @@ COPY embeddings ./embeddings
 RUN uv pip install --system .
 
 # Download the model and include in the Docker image
-# NOTE: The env vars "TE_MODEL_URI" and "TE_MODEL_DOWNLOAD_PATH" are set here to support
-#  the downloading of the model into this image build, but persist in the container and
-#  effectively also set this as the default model.
+# NOTE: The env vars "TE_MODEL_URI" and "TE_MODEL_PATH" are set here to support
+#  the downloading of the model during image build, but also persist in the container and
+#  effectively set the default model.
 ENV HF_HUB_DISABLE_PROGRESS_BARS=true
 ENV TE_MODEL_URI=opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte
-ENV TE_MODEL_DOWNLOAD_PATH=/model
+ENV TE_MODEL_PATH=/model
 RUN python -m embeddings.cli --verbose download-model
 
 ENTRYPOINT ["python", "-m", "embeddings.cli"]

README.md

@@ -24,7 +24,7 @@ WORKSPACE=### Set to `dev` for local development, this will be set to `stage` an
 
 ```shell
 TE_MODEL_URI=# HuggingFace model URI
-TE_MODEL_DOWNLOAD_PATH=# Download location for model
+TE_MODEL_PATH=# Path where the model will be downloaded to and loaded from
 HF_HUB_DISABLE_PROGRESS_BARS=#boolean to use progress bars for HuggingFace model downloads; defaults to 'true' in deployed contexts
 ```
 
@@ -34,7 +34,7 @@ This CLI application is designed to create embeddings for input texts.  To do th
 
 To this end, there is a base embedding class `BaseEmbeddingModel` that is designed to be extended and customized for a particular embedding model.
 
-Once an embedding class has been created, the preferred approach is to set env vars `TE_MODEL_URI` and `TE_MODEL_DOWNLOAD_PATH` directly in the `Dockerfile` to a) download a local snapshot of the model during image build, and b) set this model as the default for the CLI.
+Once an embedding class has been created, the preferred approach is to set env vars `TE_MODEL_URI` and `TE_MODEL_PATH` directly in the `Dockerfile` to a) download a local snapshot of the model during image build, and b) set this model as the default for the CLI.
 
 This allows invoking the CLI without specifying a model URI or local location, allowing this model to serve as the default, e.g.:
 
@@ -61,18 +61,38 @@ Usage: embeddings ping [OPTIONS]
 ```text
 Usage: embeddings download-model [OPTIONS]
 
-  Download a model from HuggingFace and save as zip file.
+  Download a model from HuggingFace and save locally.
 
 Options:
-  --model-uri TEXT  HuggingFace model URI (e.g., 'org/model-name')  [required]
-  --output PATH     Output path for zipped model (e.g., '/path/to/model.zip')
-                    [required]
-  --help            Show this message and exit.
+  --model-uri TEXT   HuggingFace model URI (e.g., 'org/model-name')
+                     [required]
+  --model-path PATH  Path where the model will be downloaded to and loaded
+                     from, e.g. '/path/to/model'.  [required]
+  --help             Show this message and exit.
 ```
 
-### `create-embeddings`
+### `test-model-load`
 ```text
-TODO...
+Usage: embeddings test-model-load [OPTIONS]
+
+  Test loading of embedding class and local model based on env vars.
+
+  In a deployed context, the following env vars are expected:     -
+  TE_MODEL_URI     - TE_MODEL_PATH
+
+  With these set, the embedding class should be registered successfully and
+  initialized, and the model loaded from a local copy.
+
+  This CLI command is NOT used during normal workflows.  This is used primary
+  during development and after model downloading/loading changes to ensure the
+  model loads correctly.
+
+Options:
+  --model-uri TEXT   HuggingFace model URI (e.g., 'org/model-name')
+                     [required]
+  --model-path PATH  Path where the model will be downloaded to and loaded
+                     from, e.g. '/path/to/model'.  [required]
+  --help             Show this message and exit.
 ```
 
 

embeddings/cli.py

@@ -1,10 +1,10 @@
 import functools
 import logging
-import os
 import time
 from collections.abc import Callable
 from datetime import timedelta
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 import click
 
@@ -13,21 +13,8 @@
 
 logger = logging.getLogger(__name__)
 
-
-def model_required(f: Callable) -> Callable:
-    """Decorator for commands that require a specific model."""
-
-    @click.option(
-        "--model-uri",
-        envvar="TE_MODEL_URI",
-        required=True,
-        help="HuggingFace model URI (e.g., 'org/model-name')",
-    )
-    @functools.wraps(f)
-    def wrapper(*args: list, **kwargs: dict) -> Callable:
-        return f(*args, **kwargs)
-
-    return wrapper
+if TYPE_CHECKING:
+    from embeddings.models.base import BaseEmbeddingModel
 
 
 @click.group("embeddings")
@@ -60,6 +47,60 @@ def _log_command_elapsed_time() -> None:
     ctx.call_on_close(_log_command_elapsed_time)
 
 
+def model_required(f: Callable) -> Callable:
+    """Middleware decorator for commands that require an embedding model.
+
+    This decorator adds two CLI options:
+    - "--model-uri": defaults to environment variable "TE_MODEL_URI"
+    - "--model-path": defaults to environment variable "TE_MODEL_PATH"
+
+    The decorator intercepts these parameters, uses the model URI to identify and
+    instantiate the appropriate embedding model class with the provided model path,
+    and stores the model instance in the Click context at ctx.obj["model"].
+
+    Both model_uri and model_path parameters are consumed by the decorator and not
+    passed to the decorated command function.
+    """
+
+    @click.option(
+        "--model-uri",
+        envvar="TE_MODEL_URI",
+        required=True,
+        help="HuggingFace model URI (e.g., 'org/model-name')",
+    )
+    @click.option(
+        "--model-path",
+        required=True,
+        envvar="TE_MODEL_PATH",
+        type=click.Path(path_type=Path),
+        help=(
+            "Path where the model will be downloaded to and loaded from, "
+            "e.g. '/path/to/model'."
+        ),
+    )
+    @functools.wraps(f)
+    def wrapper(*args: tuple, **kwargs: dict[str, str | Path]) -> Callable:
+        # pop "model_uri" and "model_path" from CLI args
+        model_uri: str = str(kwargs.pop("model_uri"))
+        model_path: str | Path = str(kwargs.pop("model_path"))
+
+        # initialize embedding class
+        model_class = get_model_class(str(model_uri))
+        model: BaseEmbeddingModel = model_class(model_path)
+        logger.info(
+            f"Embedding class '{model.__class__.__name__}' "
+            f"initialized from model URI '{model_uri}'."
+        )
+
+        # save embedding class instance to Context
+        ctx: click.Context = args[0]  # type: ignore[assignment]
+        ctx.obj["model"] = model
+
+        return f(*args, **kwargs)
+
+    return wrapper
+
+
 @main.command()
 def ping() -> None:
     """Emit 'pong' to debug logs and stdout."""
@@ -68,53 +109,49 @@ def ping() -> None:
 
 
 @main.command()
+@click.pass_context
 @model_required
-@click.option(
-    "--output",
-    required=True,
-    envvar="TE_MODEL_DOWNLOAD_PATH",
-    type=click.Path(path_type=Path),
-    help="Output path for zipped model (e.g., '/path/to/model.zip')",
-)
-def download_model(model_uri: str, output: Path) -> None:
-    """Download a model from HuggingFace and save as zip file."""
-    # load embedding model class
-    model_class = get_model_class(model_uri)
-    model = model_class()
+def download_model(
+    ctx: click.Context,
+) -> None:
+    """Download a model from HuggingFace and save locally."""
+    model: BaseEmbeddingModel = ctx.obj["model"]
 
-    # download model assets
-    logger.info(f"Downloading model: {model_uri}")
-    result_path = model.download(output)
+    logger.info(f"Downloading model: {model.model_uri}")
+    result_path = model.download()
 
     message = f"Model downloaded and saved to: {result_path}"
     logger.info(message)
     click.echo(result_path)
 
 
 @main.command()
-def test_model_load() -> None:
+@click.pass_context
+@model_required
+def test_model_load(ctx: click.Context) -> None:
     """Test loading of embedding class and local model based on env vars.
 
     In a deployed context, the following env vars are expected:
         - TE_MODEL_URI
-        - TE_MODEL_DOWNLOAD_PATH
+        - TE_MODEL_PATH
 
     With these set, the embedding class should be registered successfully and initialized,
     and the model loaded from a local copy.
-    """
-    # load embedding model class
-    model_class = get_model_class(os.environ["TE_MODEL_URI"])
-    model = model_class()
 
-    model.load(os.environ["TE_MODEL_DOWNLOAD_PATH"])
+    This CLI command is NOT used during normal workflows.  This is used primary
+    during development and after model downloading/loading changes to ensure the model
+    loads correctly.
+    """
+    model: BaseEmbeddingModel = ctx.obj["model"]
+    model.load()
     click.echo("OK")
 
 
 @main.command()
+@click.pass_context
 @model_required
-def create_embeddings(_model_uri: str) -> None:
-    # TODO: docstring # noqa: FIX002
-    raise NotImplementedError
+def create_embedding(ctx: click.Context) -> None:
+    """Create a single embedding for a single input text."""
 
 
 if __name__ == "__main__":  # pragma: no cover

embeddings/models/base.py

@@ -12,6 +12,14 @@ class BaseEmbeddingModel(ABC):
 
     MODEL_URI: str  # Type hint to document the requirement
 
+    def __init__(self, model_path: str | Path) -> None:
+        """Initialize the embedding model with a model path.
+
+        Args:
+            model_path: Path where the model will be downloaded to and loaded from.
+        """
+        self.model_path = Path(model_path)
+
     def __init_subclass__(cls, **kwargs: dict) -> None:  # noqa: D105
         super().__init_subclass__(**kwargs)
 
@@ -28,17 +36,13 @@ def model_uri(self) -> str:
         return self.MODEL_URI
 
     @abstractmethod
-    def download(self, output_path: str | Path) -> Path:
-        """Download and prepare model, saving to output_path.
+    def download(self) -> Path:
+        """Download and prepare model, saving to self.model_path.
 
-        Args:
-            output_path: Path where the model zip should be saved.
+        Returns:
+            Path where the model was saved.
         """
 
     @abstractmethod
-    def load(self, model_path: str | Path) -> None:
-        """Load model from local, downloaded instance.
-
-        Args:
-            model_path: Path of local model directory.
-        """
+    def load(self) -> None:
+        """Load model from self.model_path."""

embeddings/models/os_neural_sparse_doc_v3_gte.py

@@ -30,24 +30,27 @@ class OSNeuralSparseDocV3GTE(BaseEmbeddingModel):
 
     MODEL_URI = "opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte"
 
-    def __init__(self) -> None:
-        """Initialize the model."""
-        super().__init__()
+    def __init__(self, model_path: str | Path) -> None:
+        """Initialize the model.
+
+        Args:
+            model_path: Path where the model will be downloaded to and loaded from.
+        """
+        super().__init__(model_path)
         self._model: PreTrainedModel | None = None
         self._tokenizer: DistilBertTokenizerFast | None = None
         self._special_token_ids: list | None = None
         self._id_to_token: list | None = None
 
-    def download(self, output_path: str | Path) -> Path:
-        """Download and prepare model, saving to output_path.
+    def download(self) -> Path:
+        """Download and prepare model, saving to self.model_path.
 
-        Args:
-            output_path: Path where the model should be saved.
+        Returns:
+            Path where the model was saved.
         """
         start_time = time.perf_counter()
 
-        output_path = Path(output_path)
-        logger.info(f"Downloading model: {self.model_uri}, saving to: {output_path}.")
+        logger.info(f"Downloading model: {self.model_uri}, saving to: {self.model_path}.")
 
         with tempfile.TemporaryDirectory() as temp_dir:
             temp_path = Path(temp_dir)
@@ -60,19 +63,21 @@ def download(self, output_path: str | Path) -> Path:
             self._patch_local_model_with_alibaba_new_impl(temp_path)
 
             # compress model directory as a zip file
-            if output_path.suffix.lower() == ".zip":
+            if self.model_path.suffix.lower() == ".zip":
                 logger.debug("Creating zip file of model contents.")
-                shutil.make_archive(str(output_path.with_suffix("")), "zip", temp_path)
+                shutil.make_archive(
+                    str(self.model_path.with_suffix("")), "zip", temp_path
+                )
 
             # copy to output directory without zipping
             else:
-                logger.debug(f"Copying model contents to {output_path}")
-                if output_path.exists():
-                    shutil.rmtree(output_path)
-                shutil.copytree(temp_path, output_path)
+                logger.debug(f"Copying model contents to {self.model_path}")
+                if self.model_path.exists():
+                    shutil.rmtree(self.model_path)
+                shutil.copytree(temp_path, self.model_path)
 
         logger.info(f"Model downloaded successfully, {time.perf_counter() - start_time}s")
-        return output_path
+        return self.model_path
 
     def _patch_local_model_with_alibaba_new_impl(self, model_temp_path: Path) -> None:
         """Patch downloaded model with required assets from Alibaba-NLP/new-impl.
@@ -124,28 +129,23 @@ def _patch_local_model_with_alibaba_new_impl(self, model_temp_path: Path) -> Non
 
             logger.debug("Dependency model Alibaba-NLP/new-impl downloaded and used.")
 
-    def load(self, model_path: str | Path) -> None:
-        """Load the model from the specified path.
-
-        Args:
-            model_path: Path to the model directory.
-        """
+    def load(self) -> None:
+        """Load the model from self.model_path."""
         start_time = time.perf_counter()
-        logger.info(f"Loading model from: {model_path}")
-        model_path = Path(model_path)
+        logger.info(f"Loading model from: {self.model_path}")
 
         # ensure model exists locally
-        if not model_path.exists():
-            raise FileNotFoundError(f"Model not found at path: {model_path}")
+        if not self.model_path.exists():
+            raise FileNotFoundError(f"Model not found at path: {self.model_path}")
 
         # load local model and tokenizer
         self._model = AutoModelForMaskedLM.from_pretrained(
-            model_path,
+            self.model_path,
             trust_remote_code=True,
             local_files_only=True,
         )
         self._tokenizer = AutoTokenizer.from_pretrained(  # type: ignore[no-untyped-call]
-            model_path,
+            self.model_path,
             local_files_only=True,
         )