From ac20bf4cf024bca752a7df08278fe7b8dd9ff0a6 Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Fri, 24 Oct 2025 16:16:01 -0400 Subject: [PATCH 1/2] Rework model registration and loading Why these changes are being introduced: Previously, a model's URI was only codified in the registry.py file, and not directly obvious from a class itself aside from docstrings. How this addresses that need: Reworks the base class to require a MODEL_URI attribute. This attribute is used to generate a model_uri:class dictionary that is used for registration of models. Side effects of this change: * None Relevant ticket(s): * https://mitlibraries.atlassian.net/browse/USE-113 --- README.md | 37 +++- embeddings/cli.py | 33 ++-- embeddings/models/base.py | 21 ++- .../models/os_neural_sparse_doc_v3_gte.py | 2 + embeddings/models/registry.py | 12 +- pyproject.toml | 3 + tests/conftest.py | 4 +- tests/test_cli.py | 19 -- uv.lock | 169 ++++++++++++++++++ 9 files changed, 257 insertions(+), 43 deletions(-) diff --git a/README.md b/README.md index 27327b1..7026d28 100644 --- a/README.md +++ b/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 +``` +### `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... +``` diff --git a/embeddings/cli.py b/embeddings/cli.py index 82930cc..8611292 100644 --- a/embeddings/cli.py +++ b/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,14 +67,11 @@ 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')", ) @@ -64,7 +79,7 @@ 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 diff --git a/embeddings/models/base.py b/embeddings/models/base.py index b0cd7dc..51e6cd7 100644 --- a/embeddings/models/base.py +++ b/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: diff --git a/embeddings/models/os_neural_sparse_doc_v3_gte.py b/embeddings/models/os_neural_sparse_doc_v3_gte.py index c27f064..c6fcd10 100644 --- a/embeddings/models/os_neural_sparse_doc_v3_gte.py +++ b/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. diff --git a/embeddings/models/registry.py b/embeddings/models/registry.py index 5abca53..e3d4e18 100644 --- a/embeddings/models/registry.py +++ b/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] diff --git a/pyproject.toml b/pyproject.toml index e10c5c4..be5445b 100644 --- a/pyproject.toml +++ b/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,6 +59,7 @@ ignore = [ "D102", "D103", "D104", + "EM102", "G004", "PLR0912", "PLR0913", @@ -65,6 +67,7 @@ ignore = [ "S321", "TD002", "TD003", + "TRY003", ] # allow autofix behavior for specified rules diff --git a/tests/conftest.py b/tests/conftest.py index a4670c3..5ec25e9 100644 --- a/tests/conftest.py +++ b/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() diff --git a/tests/test_cli.py b/tests/test_cli.py index d489c72..e51eb35 100644 --- a/tests/test_cli.py +++ b/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 diff --git a/uv.lock b/uv.lock index 0c5ae56..e973a57 100644 --- a/uv.lock +++ b/uv.lock @@ -2,6 +2,15 @@ version = 1 revision = 3 requires-python = ">=3.12" +[[package]] +name = "asttokens" +version = "3.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/4a/e7/82da0a03e7ba5141f05cce0d302e6eed121ae055e0456ca228bf693984bc/asttokens-3.0.0.tar.gz", hash = "sha256:0dcd8baa8d62b0c1d118b399b2ddba3c4aff271d0d7a9e0d4c1681c79035bbc7", size = 61978, upload-time = "2024-11-30T04:30:14.439Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/25/8a/c46dcc25341b5bce5472c718902eb3d38600a903b14fa6aeecef3f21a46f/asttokens-3.0.0-py3-none-any.whl", hash = "sha256:e3078351a059199dd5138cb1c706e6430c05eff2ff136af5eb4790f9d28932e2", size = 26918, upload-time = "2024-11-30T04:30:10.946Z" }, +] + [[package]] name = "attrs" version = "25.4.0" @@ -290,6 +299,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/53/f1/f3be2e9820a2c26fa77622223e91f9c504e1581830930d477e06146073f4/cyclonedx_python_lib-9.1.0-py3-none-any.whl", hash = "sha256:55693fca8edaecc3363b24af14e82cc6e659eb1e8353e58b587c42652ce0fb52", size = 374968, upload-time = "2025-02-27T17:23:37.766Z" }, ] +[[package]] +name = "decorator" +version = "5.2.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/43/fa/6d96a0978d19e17b68d634497769987b16c8f4cd0a7a05048bec693caa6b/decorator-5.2.1.tar.gz", hash = "sha256:65f266143752f734b0a7cc83c46f4618af75b8c5911b00ccb61d0ac9b6da0360", size = 56711, upload-time = "2025-02-24T04:41:34.073Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4e/8c/f3147f5c4b73e7550fe5f9352eaa956ae838d5c51eb58e7a25b9f3e2643b/decorator-5.2.1-py3-none-any.whl", hash = "sha256:d316bb415a2d9e2d2b3abcc4084c6502fc09240e292cd76a76afc106a1c8e04a", size = 9190, upload-time = "2025-02-24T04:41:32.565Z" }, +] + [[package]] name = "defusedxml" version = "0.7.1" @@ -354,6 +372,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/2a/a2/e90242f53f7ae41554419b1695b4820b364df87c8350aa420b60b20cab92/duckdb_engine-0.17.0-py3-none-any.whl", hash = "sha256:3aa72085e536b43faab635f487baf77ddc5750069c16a2f8d9c6c3cb6083e979", size = 49676, upload-time = "2025-03-29T09:49:15.564Z" }, ] +[[package]] +name = "executing" +version = "2.2.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/cc/28/c14e053b6762b1044f34a13aab6859bbf40456d37d23aa286ac24cfd9a5d/executing-2.2.1.tar.gz", hash = "sha256:3632cc370565f6648cc328b32435bd120a1e4ebb20c77e3fdde9a13cd1e533c4", size = 1129488, upload-time = "2025-09-01T09:48:10.866Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c1/ea/53f2148663b321f21b5a606bd5f191517cf40b7072c0497d3c92c4a13b1e/executing-2.2.1-py2.py3-none-any.whl", hash = "sha256:760643d3452b4d777d295bb167ccc74c64a81df23fb5e08eff250c425a4b2017", size = 28317, upload-time = "2025-09-01T09:48:08.5Z" }, +] + [[package]] name = "filelock" version = "3.20.0" @@ -466,6 +493,51 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, ] +[[package]] +name = "ipython" +version = "9.6.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "decorator" }, + { name = "ipython-pygments-lexers" }, + { name = "jedi" }, + { name = "matplotlib-inline" }, + { name = "pexpect", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "prompt-toolkit" }, + { name = "pygments" }, + { name = "stack-data" }, + { name = "traitlets" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/2a/34/29b18c62e39ee2f7a6a3bba7efd952729d8aadd45ca17efc34453b717665/ipython-9.6.0.tar.gz", hash = "sha256:5603d6d5d356378be5043e69441a072b50a5b33b4503428c77b04cb8ce7bc731", size = 4396932, upload-time = "2025-09-29T10:55:53.948Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/48/c5/d5e07995077e48220269c28a221e168c91123ad5ceee44d548f54a057fc0/ipython-9.6.0-py3-none-any.whl", hash = "sha256:5f77efafc886d2f023442479b8149e7d86547ad0a979e9da9f045d252f648196", size = 616170, upload-time = "2025-09-29T10:55:47.676Z" }, +] + +[[package]] +name = "ipython-pygments-lexers" +version = "1.1.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pygments" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ef/4c/5dd1d8af08107f88c7f741ead7a40854b8ac24ddf9ae850afbcf698aa552/ipython_pygments_lexers-1.1.1.tar.gz", hash = "sha256:09c0138009e56b6854f9535736f4171d855c8c08a563a0dcd8022f78355c7e81", size = 8393, upload-time = "2025-01-17T11:24:34.505Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d9/33/1f075bf72b0b747cb3288d011319aaf64083cf2efef8354174e3ed4540e2/ipython_pygments_lexers-1.1.1-py3-none-any.whl", hash = "sha256:a9462224a505ade19a605f71f8fa63c2048833ce50abc86768a0d81d876dc81c", size = 8074, upload-time = "2025-01-17T11:24:33.271Z" }, +] + +[[package]] +name = "jedi" +version = "0.19.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "parso" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/72/3a/79a912fbd4d8dd6fbb02bf69afd3bb72cf0c729bb3063c6f4498603db17a/jedi-0.19.2.tar.gz", hash = "sha256:4770dc3de41bde3966b02eb84fbcf557fb33cce26ad23da12c742fb50ecb11f0", size = 1231287, upload-time = "2024-11-11T01:41:42.873Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c0/5a/9cac0c82afec3d09ccd97c8b6502d48f165f9124db81b4bcb90b4af974ee/jedi-0.19.2-py2.py3-none-any.whl", hash = "sha256:a8ef22bde8490f57fe5c7681a3c83cb58874daf72b4784de3cce5b6ef6edb5b9", size = 1572278, upload-time = "2024-11-11T01:41:40.175Z" }, +] + [[package]] name = "jmespath" version = "1.0.1" @@ -499,6 +571,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/94/54/e7d793b573f298e1c9013b8c4dade17d481164aa517d1d7148619c2cedbf/markdown_it_py-4.0.0-py3-none-any.whl", hash = "sha256:87327c59b172c5011896038353a81343b6754500a08cd7a4973bb48c6d578147", size = 87321, upload-time = "2025-08-11T12:57:51.923Z" }, ] +[[package]] +name = "matplotlib-inline" +version = "0.2.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "traitlets" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c7/74/97e72a36efd4ae2bccb3463284300f8953f199b5ffbc04cbbb0ec78f74b1/matplotlib_inline-0.2.1.tar.gz", hash = "sha256:e1ee949c340d771fc39e241ea75683deb94762c8fa5f2927ec57c83c4dffa9fe", size = 8110, upload-time = "2025-10-23T09:00:22.126Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/af/33/ee4519fa02ed11a94aef9559552f3b17bb863f2ecfe1a35dc7f548cde231/matplotlib_inline-0.2.1-py3-none-any.whl", hash = "sha256:d56ce5156ba6085e00a9d54fead6ed29a9c47e215cd1bba2e976ef39f5710a76", size = 9516, upload-time = "2025-10-23T09:00:20.675Z" }, +] + [[package]] name = "mdurl" version = "0.1.2" @@ -730,6 +814,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/70/44/5191d2e4026f86a2a109053e194d3ba7a31a2d10a9c2348368c63ed4e85a/pandas-2.3.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:3869faf4bd07b3b66a9f462417d0ca3a9df29a9f6abd5d0d0dbab15dac7abe87", size = 13202175, upload-time = "2025-09-29T23:31:59.173Z" }, ] +[[package]] +name = "parso" +version = "0.8.5" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d4/de/53e0bcf53d13e005bd8c92e7855142494f41171b34c2536b86187474184d/parso-0.8.5.tar.gz", hash = "sha256:034d7354a9a018bdce352f48b2a8a450f05e9d6ee85db84764e9b6bd96dafe5a", size = 401205, upload-time = "2025-08-23T15:15:28.028Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/16/32/f8e3c85d1d5250232a5d3477a2a28cc291968ff175caeadaf3cc19ce0e4a/parso-0.8.5-py2.py3-none-any.whl", hash = "sha256:646204b5ee239c396d040b90f9e272e9a8017c630092bf59980beb62fd033887", size = 106668, upload-time = "2025-08-23T15:15:25.663Z" }, +] + [[package]] name = "pathspec" version = "0.12.1" @@ -739,6 +832,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cc/20/ff623b09d963f88bfde16306a54e12ee5ea43e9b597108672ff3a408aad6/pathspec-0.12.1-py3-none-any.whl", hash = "sha256:a0d503e138a4c123b27490a4f7beda6a01c6f288df0e4a8b79c7eb0dc7b4cc08", size = 31191, upload-time = "2023-12-10T22:30:43.14Z" }, ] +[[package]] +name = "pexpect" +version = "4.9.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ptyprocess" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/42/92/cc564bf6381ff43ce1f4d06852fc19a2f11d180f23dc32d9588bee2f149d/pexpect-4.9.0.tar.gz", hash = "sha256:ee7d41123f3c9911050ea2c2dac107568dc43b2d3b0c7557a33212c398ead30f", size = 166450, upload-time = "2023-11-25T09:07:26.339Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9e/c3/059298687310d527a58bb01f3b1965787ee3b40dce76752eda8b44e9a2c5/pexpect-4.9.0-py2.py3-none-any.whl", hash = "sha256:7236d1e080e4936be2dc3e326cec0af72acf9212a7e1d060210e70a47e253523", size = 63772, upload-time = "2023-11-25T06:56:14.81Z" }, +] + [[package]] name = "pip" version = "25.2" @@ -827,6 +932,36 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5b/a5/987a405322d78a73b66e39e4a90e4ef156fd7141bf71df987e50717c321b/pre_commit-4.3.0-py2.py3-none-any.whl", hash = "sha256:2b0747ad7e6e967169136edffee14c16e148a778a54e4f967921aa1ebf2308d8", size = 220965, upload-time = "2025-08-09T18:56:13.192Z" }, ] +[[package]] +name = "prompt-toolkit" +version = "3.0.52" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "wcwidth" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a1/96/06e01a7b38dce6fe1db213e061a4602dd6032a8a97ef6c1a862537732421/prompt_toolkit-3.0.52.tar.gz", hash = "sha256:28cde192929c8e7321de85de1ddbe736f1375148b02f2e17edd840042b1be855", size = 434198, upload-time = "2025-08-27T15:24:02.057Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/84/03/0d3ce49e2505ae70cf43bc5bb3033955d2fc9f932163e84dc0779cc47f48/prompt_toolkit-3.0.52-py3-none-any.whl", hash = "sha256:9aac639a3bbd33284347de5ad8d68ecc044b91a762dc39b7c21095fcd6a19955", size = 391431, upload-time = "2025-08-27T15:23:59.498Z" }, +] + +[[package]] +name = "ptyprocess" +version = "0.7.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/20/e5/16ff212c1e452235a90aeb09066144d0c5a6a8c0834397e03f5224495c4e/ptyprocess-0.7.0.tar.gz", hash = "sha256:5c5d0a3b48ceee0b48485e0c26037c0acd7d29765ca3fbb5cb3831d347423220", size = 70762, upload-time = "2020-12-28T15:15:30.155Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/22/a6/858897256d0deac81a172289110f31629fc4cee19b6f01283303e18c8db3/ptyprocess-0.7.0-py2.py3-none-any.whl", hash = "sha256:4b41f3967fce3af57cc7e94b888626c18bf37a083e3651ca8feeb66d492fef35", size = 13993, upload-time = "2020-12-28T15:15:28.35Z" }, +] + +[[package]] +name = "pure-eval" +version = "0.2.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/cd/05/0a34433a064256a578f1783a10da6df098ceaa4a57bbeaa96a6c0352786b/pure_eval-0.2.3.tar.gz", hash = "sha256:5f4e983f40564c576c7c8635ae88db5956bb2229d7e9237d03b3c0b0190eaf42", size = 19752, upload-time = "2024-07-21T12:58:21.801Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8e/37/efad0257dc6e593a18957422533ff0f87ede7c9c6ea010a2177d738fb82f/pure_eval-0.2.3-py3-none-any.whl", hash = "sha256:1db8e35b67b3d218d818ae653e27f06c3aa420901fa7b081ca98cbedc874e0d0", size = 11842, upload-time = "2024-07-21T12:58:20.04Z" }, +] + [[package]] name = "py-serializable" version = "2.1.0" @@ -1104,6 +1239,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/9c/5e/6a29fa884d9fb7ddadf6b69490a9d45fded3b38541713010dad16b77d015/sqlalchemy-2.0.44-py3-none-any.whl", hash = "sha256:19de7ca1246fbef9f9d1bff8f1ab25641569df226364a0e40457dc5457c54b05", size = 1928718, upload-time = "2025-10-10T15:29:45.32Z" }, ] +[[package]] +name = "stack-data" +version = "0.6.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "asttokens" }, + { name = "executing" }, + { name = "pure-eval" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/28/e3/55dcc2cfbc3ca9c29519eb6884dd1415ecb53b0e934862d3559ddcb7e20b/stack_data-0.6.3.tar.gz", hash = "sha256:836a778de4fec4dcd1dcd89ed8abff8a221f58308462e1c4aa2a3cf30148f0b9", size = 44707, upload-time = "2023-09-30T13:58:05.479Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f1/7b/ce1eafaf1a76852e2ec9b22edecf1daa58175c090266e9f6c64afcd81d91/stack_data-0.6.3-py3-none-any.whl", hash = "sha256:d5558e0c25a4cb0853cddad3d77da9891a08cb85dd9f9f91b9f8cd66e511e695", size = 24521, upload-time = "2023-09-30T13:58:03.53Z" }, +] + [[package]] name = "timdex-dataset-api" version = "3.5.0" @@ -1133,6 +1282,7 @@ dependencies = [ dev = [ { name = "black" }, { name = "coveralls" }, + { name = "ipython" }, { name = "mypy" }, { name = "pip-audit" }, { name = "pre-commit" }, @@ -1152,6 +1302,7 @@ requires-dist = [ dev = [ { name = "black", specifier = ">=25.1.0" }, { name = "coveralls", specifier = ">=4.0.1" }, + { name = "ipython", specifier = ">=9.6.0" }, { name = "mypy", specifier = ">=1.17.1" }, { name = "pip-audit", specifier = ">=2.9.0" }, { name = "pre-commit", specifier = ">=4.3.0" }, @@ -1180,6 +1331,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d0/30/dc54f88dd4a2b5dc8a0279bdd7270e735851848b762aeb1c1184ed1f6b14/tqdm-4.67.1-py3-none-any.whl", hash = "sha256:26445eca388f82e72884e0d580d5464cd801a3ea01e63e5601bdff9ba6a48de2", size = 78540, upload-time = "2024-11-24T20:12:19.698Z" }, ] +[[package]] +name = "traitlets" +version = "5.14.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/eb/79/72064e6a701c2183016abbbfedaba506d81e30e232a68c9f0d6f6fcd1574/traitlets-5.14.3.tar.gz", hash = "sha256:9ed0579d3502c94b4b3732ac120375cda96f923114522847de4b3bb98b96b6b7", size = 161621, upload-time = "2024-04-19T11:11:49.746Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/00/c0/8f5d070730d7836adc9c9b6408dec68c6ced86b304a9b26a14df072a6e8c/traitlets-5.14.3-py3-none-any.whl", hash = "sha256:b74e89e397b1ed28cc831db7aea759ba6640cb3de13090ca145426688ff1ac4f", size = 85359, upload-time = "2024-04-19T11:11:46.763Z" }, +] + [[package]] name = "typing-extensions" version = "4.15.0" @@ -1220,3 +1380,12 @@ sdist = { url = "https://files.pythonhosted.org/packages/a4/d5/b0ccd381d55c8f45d wheels = [ { url = "https://files.pythonhosted.org/packages/27/73/d9a94da0e9d470a543c1b9d3ccbceb0f59455983088e727b8a1824ed90fb/virtualenv-20.35.3-py3-none-any.whl", hash = "sha256:63d106565078d8c8d0b206d48080f938a8b25361e19432d2c9db40d2899c810a", size = 5981061, upload-time = "2025-10-10T21:23:30.433Z" }, ] + +[[package]] +name = "wcwidth" +version = "0.2.14" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/24/30/6b0809f4510673dc723187aeaf24c7f5459922d01e2f794277a3dfb90345/wcwidth-0.2.14.tar.gz", hash = "sha256:4d478375d31bc5395a3c55c40ccdf3354688364cd61c4f6adacaa9215d0b3605", size = 102293, upload-time = "2025-09-22T16:29:53.023Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/af/b5/123f13c975e9f27ab9c0770f514345bd406d0e8d3b7a0723af9d43f710af/wcwidth-0.2.14-py2.py3-none-any.whl", hash = "sha256:a7bb560c8aee30f9957e5f9895805edd20602f2d7f720186dfd906e82b4982e1", size = 37286, upload-time = "2025-09-22T16:29:51.641Z" }, +] From 847f6711686bbc5111c0aae2f8b4fd192c7ba900 Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Fri, 24 Oct 2025 16:26:18 -0400 Subject: [PATCH 2/2] Add MODEL_URI unit tests --- tests/test_models.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/tests/test_models.py b/tests/test_models.py index e291d8f..caeba00 100644 --- a/tests/test_models.py +++ b/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