Stub base model class and CLI download-model command

Why these changes are being introduced:

For this application to create embeddings, it will need some
structure to handle the downloading, loading, and use of
particular models.

Additionally, this application should provide the opinionated
functionality to download and zip up a model, either for use
within a Docker image, or local testing.

How this addresses that need:
* Class 'BaseEmbeddingModel' is created.  This base class will be
extended by actual models we intend to use.  This base class will
expect any child class to define methods for downloading, loading,
and using the model.

* Stubs a CLI command 'download-model' which we will use very
soon for building a Docker image that contains the model weights.

Side effects of this change:
* None at this time.

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

Author: ghukill

embeddings/cli.py

@@ -1,10 +1,12 @@
 import logging
 import time
 from datetime import timedelta
+from pathlib import Path
 
 import click
 
 from embeddings.config import configure_logger, configure_sentry
+from embeddings.models.registry import get_model_class
 
 logger = logging.getLogger(__name__)
 
@@ -46,6 +48,54 @@ def ping() -> None:
     click.echo("pong")
 
 
+@main.command()
+@click.option(
+    "--model-uri",
+    required=True,
+    help="HuggingFace model URI (e.g., 'org/model-name')",
+)
+@click.option(
+    "--output",
+    required=True,
+    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."""
+    try:
+        model_class = get_model_class(model_uri)
+    except ValueError as e:
+        logger.exception("Unknown model URI: %s", model_uri)
+        raise click.ClickException(str(e)) from e
+
+    logger.info("Downloading model: %s", model_uri)
+    model = model_class(model_uri)
+
+    try:
+        result_path = model.download(output)
+        logger.info("Model downloaded successfully to: %s", result_path)
+        click.echo(f"Model saved to: {result_path}")
+    except NotImplementedError as e:
+        logger.exception("Download not yet implemented for model: %s", model_uri)
+        raise click.ClickException(str(e)) from e
+    except Exception as e:
+        logger.exception("Failed to download model: %s", model_uri)
+        msg = f"Download failed: {e}"
+        raise click.ClickException(msg) from e
+
+
+@main.command()
+@click.option(
+    "--model-uri",
+    required=True,
+    help="HuggingFace model URI (e.g., 'org/model-name')",
+)
+def create_embeddings(model_uri: str) -> None:
+    """Create embeddings."""
+    logger.info("create-embeddings command called with model: %s", model_uri)
+    raise NotImplementedError
+
+
 if __name__ == "__main__":  # pragma: no cover
     logger = logging.getLogger("embeddings.main")
     main()

embeddings/models/__init__.py

@@ -0,0 +1 @@
+

embeddings/models/base.py

@@ -0,0 +1,23 @@
+"""Base class for embedding models."""
+
+from abc import ABC, abstractmethod
+from pathlib import Path
+
+
+class BaseEmbeddingModel(ABC):
+    """Abstract base class for embedding models.
+
+    Args:
+        model_uri: HuggingFace model identifier (e.g., 'org/model-name').
+    """
+
+    def __init__(self, model_uri: str) -> None:
+        self.model_uri = model_uri
+
+    @abstractmethod
+    def download(self, output_path: Path) -> Path:
+        """Download and prepare model, saving to output_path.
+
+        Args:
+            output_path: Path where the model zip should be saved.
+        """

embeddings/models/os_neural_sparse_doc_v3_gte.py

@@ -0,0 +1,24 @@
+"""OpenSearch Neural Sparse Doc v3 GTE model."""
+
+import logging
+from pathlib import Path
+
+from embeddings.models.base import BaseEmbeddingModel
+
+logger = logging.getLogger(__name__)
+
+
+class OSNeuralSparseDocV3GTE(BaseEmbeddingModel):
+    """OpenSearch Neural Sparse Encoding Doc v3 GTE model.
+
+    HuggingFace 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.
+
+        Args:
+            output_path: Path where the model zip should be saved.
+        """
+        logger.info(f"Downloading model: { self.model_uri}, saving to: {output_path}.")
+        raise NotImplementedError

embeddings/models/registry.py

@@ -0,0 +1,26 @@
+"""Registry mapping model URIs to model classes."""
+
+from embeddings.models.base import BaseEmbeddingModel
+from embeddings.models.os_neural_sparse_doc_v3_gte import OSNeuralSparseDocV3GTE
+
+MODEL_REGISTRY: dict[str, type[BaseEmbeddingModel]] = {
+    "opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte": (
+        OSNeuralSparseDocV3GTE
+    ),
+}
+
+
+def get_model_class(model_uri: str) -> type[BaseEmbeddingModel]:
+    """Get model class for given 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}"
+        raise ValueError(msg)
+    return MODEL_REGISTRY[model_uri]

pyproject.toml

@@ -9,6 +9,7 @@ requires-python = ">=3.12"
 
 dependencies = [
     "click>=8.2.1",
+    "huggingface-hub>=0.26.0",
     "sentry-sdk>=2.34.1",
     "timdex-dataset-api",
 ]
@@ -56,7 +57,8 @@ ignore = [
     "D101", 
     "D102",
     "D103",
-    "D104", 
+    "D104",
+    "G004",
     "PLR0912",
     "PLR0913", 
     "PLR0915",

tests/conftest.py

@@ -1,6 +1,11 @@
+import zipfile
+from pathlib import Path
+
 import pytest
 from click.testing import CliRunner
 
+from embeddings.models.base import BaseEmbeddingModel
+
 
 @pytest.fixture(autouse=True)
 def _test_env(monkeypatch):
@@ -11,3 +16,22 @@ def _test_env(monkeypatch):
 @pytest.fixture
 def runner():
     return CliRunner()
+
+
+class MockEmbeddingModel(BaseEmbeddingModel):
+    """Simple test model that doesn't hit external APIs."""
+
+    def download(self, output_path: Path) -> Path:
+        """Create a fake model zip file for testing."""
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+        with zipfile.ZipFile(output_path, "w") as zf:
+            zf.writestr("config.json", '{"model": "mock", "vocab_size": 30000}')
+            zf.writestr("pytorch_model.bin", b"fake model weights")
+            zf.writestr("tokenizer.json", '{"version": "1.0"}')
+        return output_path
+
+
+@pytest.fixture
+def mock_model():
+    """Fixture providing a MockEmbeddingModel instance."""
+    return MockEmbeddingModel("test/mock-model")

tests/test_cli.py

@@ -14,3 +14,30 @@ def test_cli_debug_logging(caplog, runner):
     assert "Logger 'root' configured with level=DEBUG" in caplog.text
     assert "pong" in caplog.text
     assert "pong" in result.output
+
+
+def test_download_model_unknown_uri(runner):
+    result = runner.invoke(
+        main, ["download-model", "--model-uri", "unknown/model", "--output", "out.zip"]
+    )
+    assert result.exit_code != 0
+    assert "Unknown model URI" in result.output
+
+
+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

@@ -0,0 +1,48 @@
+import zipfile
+
+import pytest
+
+from embeddings.models.registry import MODEL_REGISTRY, get_model_class
+
+
+def test_mock_model_instantiation(mock_model):
+    assert mock_model.model_uri == "test/mock-model"
+
+
+def test_mock_model_download_creates_zip(mock_model, tmp_path):
+    output_path = tmp_path / "test_model.zip"
+    result = mock_model.download(output_path)
+
+    assert result == output_path
+    assert output_path.exists()
+    assert zipfile.is_zipfile(output_path)
+
+
+def test_mock_model_download_contains_expected_files(mock_model, tmp_path):
+    output_path = tmp_path / "test_model.zip"
+    mock_model.download(output_path)
+
+    with zipfile.ZipFile(output_path, "r") as zf:
+        file_list = zf.namelist()
+        assert "config.json" in file_list
+        assert "pytorch_model.bin" in file_list
+        assert "tokenizer.json" in file_list
+
+
+def test_registry_contains_opensearch_model():
+    assert (
+        "opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte"
+        in MODEL_REGISTRY
+    )
+
+
+def test_get_model_class_returns_correct_class():
+    model_class = get_model_class(
+        "opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte"
+    )
+    assert model_class.__name__ == "OSNeuralSparseDocV3GTE"
+
+
+def test_get_model_class_raises_for_unknown_uri():
+    with pytest.raises(ValueError, match="Unknown model URI"):
+        get_model_class("unknown/model-uri")