USE 113 - Rework model registration and loading

README.md

@@ -22,8 +22,43 @@ WORKSPACE=### Set to `dev` for local development, this will be set to `stage` an
 
 ### Optional
 
-_None yet at this time._
+```shell
+TE_MODEL_URI=# HuggingFace model URI
+TE_MODEL_DOWNLOAD_PATH=# Download location for model
+HF_HUB_DISABLE_PROGRESS_BARS=#boolean to use progress bars for HuggingFace model downloads; defaults to 'true' in deployed contexts
+```
+
+## CLI Commands
+
+For local development, all CLI commands should be invoked with the following format to pickup environment variables from `.env`:
+
+```shell
+uv run --env-file .env embeddings <COMMAND> <ARGS>
+```
 
+### `ping`
+```text
+Usage: embeddings ping [OPTIONS]
 
+  Emit 'pong' to debug logs and stdout.
+```
+
+### `download-model`
+```text
+Usage: embeddings download-model [OPTIONS]
+
+  Download a model from HuggingFace and save as zip file.
+
+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.
+```
+
+### `create-embeddings`
+```text
+TODO...
+```
 
 

embeddings/cli.py

@@ -1,5 +1,7 @@
+import functools
 import logging
 import time
+from collections.abc import Callable
 from datetime import timedelta
 from pathlib import Path
 
@@ -11,6 +13,22 @@
 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
+
+
 @click.group("embeddings")
 @click.option(
     "-v",
@@ -49,22 +67,19 @@ def ping() -> None:
 
 
 @main.command()
-@click.option(
-    "--model-uri",
-    required=True,
-    help="HuggingFace model URI (e.g., 'org/model-name')",
-)
+@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(model_uri)
+    model = model_class()
 
     # download model assets
     logger.info(f"Downloading model: {model_uri}")
@@ -76,11 +91,7 @@ def download_model(model_uri: str, output: Path) -> None:
 
 
 @main.command()
-@click.option(
-    "--model-uri",
-    required=True,
-    help="HuggingFace model URI (e.g., 'org/model-name')",
-)
+@model_required
 def create_embeddings(_model_uri: str) -> None:
     # TODO: docstring # noqa: FIX002
     raise NotImplementedError

embeddings/models/base.py

@@ -7,12 +7,25 @@
 class BaseEmbeddingModel(ABC):
     """Abstract base class for embedding models.
 
-    Args:
-        model_uri: HuggingFace model identifier (e.g., 'org/model-name').
+    All child classes must set class level attribute MODEL_URI.
     """
 
-    def __init__(self, model_uri: str) -> None:
-        self.model_uri = model_uri
+    MODEL_URI: str  # Type hint to document the requirement
+
+    def __init_subclass__(cls, **kwargs: dict) -> None:  # noqa: D105
+        super().__init_subclass__(**kwargs)
+
+        # require class level MODEL_URI to be set
+        if not hasattr(cls, "MODEL_URI"):
+            msg = f"{cls.__name__} must define 'MODEL_URI' class attribute"
+            raise TypeError(msg)
+        if not isinstance(cls.MODEL_URI, str):
+            msg = f"{cls.__name__} must override 'MODEL_URI' with a valid string"
+            raise TypeError(msg)
+
+    @property
+    def model_uri(self) -> str:
+        return self.MODEL_URI
 
     @abstractmethod
     def download(self, output_path: Path) -> Path:

embeddings/models/os_neural_sparse_doc_v3_gte.py

@@ -14,6 +14,8 @@ class OSNeuralSparseDocV3GTE(BaseEmbeddingModel):
     HuggingFace URI: opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte
     """
 
+    MODEL_URI = "opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte"
+
     def download(self, output_path: Path) -> Path:
         """Download and prepare model, saving to output_path.
 

embeddings/models/registry.py

@@ -7,25 +7,23 @@
 
 logger = logging.getLogger(__name__)
 
+MODEL_CLASSES = [OSNeuralSparseDocV3GTE]
+
 MODEL_REGISTRY: dict[str, type[BaseEmbeddingModel]] = {
-    "opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte": (
-        OSNeuralSparseDocV3GTE
-    ),
+    model.MODEL_URI: model for model in MODEL_CLASSES
 }
 
 
 def get_model_class(model_uri: str) -> type[BaseEmbeddingModel]:
-    """Get model class for given URI.
+    """Get an embedding model class via the HuggingFace model URI.
 
     Args:
         model_uri: HuggingFace model identifier.
-
-    Returns:
-        Model class for the given URI.
     """
     if model_uri not in MODEL_REGISTRY:
         available = ", ".join(sorted(MODEL_REGISTRY.keys()))
         msg = f"Unknown model URI: {model_uri}. Available models: {available}"
         logger.error(msg)
         raise ValueError(msg)
+
     return MODEL_REGISTRY[model_uri]

pyproject.toml

@@ -18,6 +18,7 @@ dependencies = [
 dev = [
     "black>=25.1.0",
     "coveralls>=4.0.1",
+    "ipython>=9.6.0",
     "mypy>=1.17.1",
     "pip-audit>=2.9.0",
     "pre-commit>=4.3.0",
@@ -58,13 +59,15 @@ ignore = [
     "D102",
     "D103",
     "D104",
+    "EM102",
     "G004",
     "PLR0912",
     "PLR0913", 
     "PLR0915",
     "S321",
     "TD002",
     "TD003",
+    "TRY003",
 ]
 
 # allow autofix behavior for specified rules

tests/conftest.py

@@ -21,6 +21,8 @@ def runner():
 class MockEmbeddingModel(BaseEmbeddingModel):
     """Simple test model that doesn't hit external APIs."""
 
+    MODEL_URI = "test/mock-model"
+
     def download(self, output_path: Path) -> Path:
         """Create a fake model zip file for testing."""
         output_path.parent.mkdir(parents=True, exist_ok=True)
@@ -34,4 +36,4 @@ def download(self, output_path: Path) -> Path:
 @pytest.fixture
 def mock_model():
     """Fixture providing a MockEmbeddingModel instance."""
-    return MockEmbeddingModel("test/mock-model")
+    return MockEmbeddingModel()

tests/test_cli.py

@@ -22,22 +22,3 @@ def test_download_model_unknown_uri(caplog, runner):
     )
     assert result.exit_code != 0
     assert "Unknown model URI" in caplog.text
-
-
-def test_download_model_not_implemented(caplog, runner):
-    caplog.set_level("INFO")
-    result = runner.invoke(
-        main,
-        [
-            "download-model",
-            "--model-uri",
-            "opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte",
-            "--output",
-            "out.zip",
-        ],
-    )
-    assert (
-        "Downloading model: opensearch-project/"
-        "opensearch-neural-sparse-encoding-doc-v3-gte, saving to: out.zip."
-    ) in caplog.text
-    assert result.exit_code != 0

tests/test_models.py

@@ -2,6 +2,7 @@
 
 import pytest
 
+from embeddings.models.base import BaseEmbeddingModel
 from embeddings.models.registry import MODEL_REGISTRY, get_model_class
 
 
@@ -46,3 +47,17 @@ def test_get_model_class_returns_correct_class():
 def test_get_model_class_raises_for_unknown_uri():
     with pytest.raises(ValueError, match="Unknown model URI"):
         get_model_class("unknown/model-uri")
+
+
+def test_subclass_without_model_uri_raises_type_error():
+    with pytest.raises(TypeError, match="must define 'MODEL_URI' class attribute"):
+
+        class InvalidModel(BaseEmbeddingModel):
+            pass
+
+
+def test_subclass_with_non_string_model_uri_raises_type_error():
+    with pytest.raises(TypeError, match="must override 'MODEL_URI' with a valid string"):
+
+        class InvalidModel(BaseEmbeddingModel):
+            MODEL_URI = 123