From ce181dfeb2aa2d31c4abec11f7cf7c865251c016 Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Thu, 30 Oct 2025 11:50:26 -0400 Subject: [PATCH 1/5] Stub CLI command methods to create embeddings How this addresses that need: * CLI command create-embeddings created * args and some functionality in place * WIP comments and DEBUG code temporarily added to demonstrate how it will work * class RecordText added to encapsulate text that is ready for an embedding * this will support future functionality of pre-embedding "strategies" applied to records * class Embedding created to encapsulate the embedding result * this captures the TIMDEX record the embedding was assocaited with, and the model + strategy used to prepare the text Side effects of this change: * None Relevant ticket(s): * https://mitlibraries.atlassian.net/browse/USE-112 --- .gitignore | 2 + embeddings/cli.py | 137 +++++++++++++++++- embeddings/config.py | 2 - embeddings/embedding.py | 51 +++++++ embeddings/models/base.py | 22 +++ .../models/os_neural_sparse_doc_v3_gte.py | 4 + pyproject.toml | 7 + tests/conftest.py | 19 +++ tests/test_cli.py | 68 ++++++--- tests/test_models.py | 67 +++++++-- tests/test_os_neural_sparse_doc_v3_gte.py | 2 +- uv.lock | 14 ++ 12 files changed, 362 insertions(+), 33 deletions(-) create mode 100644 embeddings/embedding.py diff --git a/.gitignore b/.gitignore index 526ebcd..e2261d0 100644 --- a/.gitignore +++ b/.gitignore @@ -155,3 +155,5 @@ cython_debug/ .DS_Store output/ .vscode/ + +CLAUDE.md \ No newline at end of file diff --git a/embeddings/cli.py b/embeddings/cli.py index f68212b..3e061b8 100644 --- a/embeddings/cli.py +++ b/embeddings/cli.py @@ -7,6 +7,8 @@ from typing import TYPE_CHECKING import click +import jsonlines +from timdex_dataset_api import TIMDEXDataset from embeddings.config import configure_logger, configure_sentry from embeddings.models.registry import get_model_class @@ -150,8 +152,139 @@ def test_model_load(ctx: click.Context) -> None: @main.command() @click.pass_context @model_required -def create_embedding(ctx: click.Context) -> None: - """Create a single embedding for a single input text.""" +@click.option( + "-d", + "--dataset-location", + required=True, + type=click.Path(), + help="TIMDEX dataset location, e.g. 's3://timdex/dataset', to read records from.", +) +@click.option( + "--run-id", + required=True, + type=str, + help="TIMDEX ETL run id.", +) +@click.option( + "--run-record-offset", + required=True, + type=int, + default=0, + help="TIMDEX ETL run record offset to start from, default = 0.", +) +@click.option( + "--record-limit", + required=True, + type=int, + default=None, + help="Limit number of records after --run-record-offset, default = None (unlimited).", +) +@click.option( + "--strategy", + type=str, # WIP: establish an enum of supported strategies + required=True, + multiple=True, + help="Pre-embedding record transformation strategy to use. Repeatable.", +) +@click.option( + "--output-jsonl", + required=False, + type=str, + default=None, + help="Optionally write embeddings to local JSONLines file (primarily for testing).", +) +def create_embeddings( + ctx: click.Context, + dataset_location: str, + run_id: str, + run_record_offset: int, + record_limit: int, + strategy: list[str], + output_jsonl: str, +) -> None: + """Create embeddings for TIMDEX records.""" + model: BaseEmbeddingModel = ctx.obj["model"] + + # init TIMDEXDataset + timdex_dataset = TIMDEXDataset(dataset_location) + + # query TIMDEX dataset for an iterator of records + timdex_records = timdex_dataset.read_dicts_iter( + columns=[ + "timdex_record_id", + "run_id", + "run_record_offset", + "transformed_record", + ], + run_id=run_id, + where=f"""run_record_offset >= {run_record_offset}""", + limit=record_limit, + action="index", + ) + + # create an iterator of InputTexts applying all requested strategies to all records + # WIP NOTE: this will leverage some kind of pre-embedding transformer class(es) that + # create texts based on the requested strategies (e.g. "full record"), which are + # captured in --strategy CLI args + # WIP NOTE: the following simulates that... + # DEBUG ------------------------------------------------------------------------------ + import json # noqa: PLC0415 + + from embeddings.embedding import RecordText # noqa: PLC0415 + + input_records = ( + RecordText( + timdex_record_id=timdex_record["timdex_record_id"], + run_id=timdex_record["run_id"], + run_record_offset=timdex_record["run_record_offset"], + embedding_strategy=_strategy, + text=json.dumps(timdex_record["transformed_record"].decode()), + ) + for timdex_record in timdex_records + for _strategy in strategy + ) + # DEBUG ------------------------------------------------------------------------------ + + # create an iterator of Embeddings via the embedding model + # WIP NOTE: this will use the embedding class .create_embeddings() bulk method + # WIP NOTE: the following simulates that... + # DEBUG ------------------------------------------------------------------------------ + from embeddings.embedding import Embedding # noqa: PLC0415 + + embeddings = ( + Embedding( + timdex_record_id=input_record.timdex_record_id, + run_id=input_record.run_id, + run_record_offset=input_record.run_record_offset, + embedding_strategy=input_record.embedding_strategy, + model_uri=model.model_uri, + embedding={"coffee": 0.9, "seattle": 0.5}, + ) + for input_record in input_records + ) + # DEBUG ------------------------------------------------------------------------------ + + # if requested, write embeddings to a local JSONLines file + if output_jsonl: + with jsonlines.open( + output_jsonl, + mode="w", + dumps=lambda obj: json.dumps( + obj, + default=str, + ), + ) as writer: + for embedding in embeddings: + writer.write(embedding.to_dict()) + + # else, default writing embeddings back to TIMDEX dataset + else: + # WIP NOTE: write via anticipated timdex_dataset.embeddings.write(...) + # NOTE: will likely use an imported TIMDEXEmbedding class from TDA, which the + # Embedding instance will nearly 1:1 map to. + raise NotImplementedError + + logger.info("Embeddings creation complete.") if __name__ == "__main__": # pragma: no cover diff --git a/embeddings/config.py b/embeddings/config.py index 7d47b73..21e094e 100644 --- a/embeddings/config.py +++ b/embeddings/config.py @@ -11,8 +11,6 @@ def configure_logger(logger: logging.Logger, *, verbose: bool) -> str: "%(message)s" ) logger.setLevel(logging.DEBUG) - for handler in logging.root.handlers: - handler.addFilter(logging.Filter("embeddings")) else: logging.basicConfig( format="%(asctime)s %(levelname)s %(name)s.%(funcName)s(): %(message)s" diff --git a/embeddings/embedding.py b/embeddings/embedding.py new file mode 100644 index 0000000..596c567 --- /dev/null +++ b/embeddings/embedding.py @@ -0,0 +1,51 @@ +import datetime +import json +from dataclasses import asdict, dataclass, field + + +@dataclass +class RecordText: + """Input record for creating an embedding for. + + Args: + (timdex_record_id, run_id, run_record_offset): composite key for TIMDEX record + embedding_strategy: strategy used to create text for embedding + text: text to embed, created from the TIMDEX record via the embedding_strategy + """ + + timdex_record_id: str + run_id: str + run_record_offset: int + embedding_strategy: str + text: str + + +@dataclass +class Embedding: + """Encapsulates a single embedding. + + Args: + (timdex_record_id, run_id, run_record_offset): composite key for TIMDEX record + model_uri: model URI used to create the embedding + embedding_strategy: strategy used to create text for embedding + embedding: model embedding created from text + """ + + timdex_record_id: str + run_id: str + run_record_offset: int + model_uri: str + embedding_strategy: str + embedding: dict | list[float] + + timestamp: datetime.datetime = field( + default_factory=lambda: datetime.datetime.now(datetime.UTC) + ) + + def to_dict(self) -> dict: + """Marshal to dictionary.""" + return asdict(self) + + def to_json(self) -> str: + """Serialize to JSON.""" + return json.dumps(self.to_dict(), default=str) diff --git a/embeddings/models/base.py b/embeddings/models/base.py index 34d55fe..ac08cc5 100644 --- a/embeddings/models/base.py +++ b/embeddings/models/base.py @@ -1,8 +1,11 @@ """Base class for embedding models.""" from abc import ABC, abstractmethod +from collections.abc import Iterator from pathlib import Path +from embeddings.embedding import Embedding, RecordText + class BaseEmbeddingModel(ABC): """Abstract base class for embedding models. @@ -46,3 +49,22 @@ def download(self) -> Path: @abstractmethod def load(self) -> None: """Load model from self.model_path.""" + + @abstractmethod + def create_embedding(self, input_record: RecordText) -> Embedding: + """Create an Embedding for an RecordText. + + Args: + input_record: RecordText instance + """ + + def create_embeddings( + self, input_records: Iterator[RecordText] + ) -> Iterator[Embedding]: + """Yield Embeddings for an iterator of InputRecords. + + Args: + input_records: iterator of InputRecords + """ + for input_text in input_records: + yield self.create_embedding(input_text) diff --git a/embeddings/models/os_neural_sparse_doc_v3_gte.py b/embeddings/models/os_neural_sparse_doc_v3_gte.py index 8e72f23..37fae99 100644 --- a/embeddings/models/os_neural_sparse_doc_v3_gte.py +++ b/embeddings/models/os_neural_sparse_doc_v3_gte.py @@ -11,6 +11,7 @@ from huggingface_hub import snapshot_download from transformers import AutoModelForMaskedLM, AutoTokenizer +from embeddings.embedding import Embedding, RecordText from embeddings.models.base import BaseEmbeddingModel if TYPE_CHECKING: @@ -161,3 +162,6 @@ def load(self) -> None: self._id_to_token[token_id] = token logger.info(f"Model loaded successfully, {time.perf_counter()-start_time}s") + + def create_embedding(self, input_record: RecordText) -> Embedding: + raise NotImplementedError diff --git a/pyproject.toml b/pyproject.toml index 3e0bb33..cec37ea 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,6 +10,7 @@ requires-python = ">=3.12" dependencies = [ "click>=8.2.1", "huggingface-hub>=0.26.0", + "jsonlines>=4.0.0", "sentry-sdk>=2.34.1", "timdex-dataset-api", "torch>=2.9.0", @@ -39,6 +40,11 @@ exclude = [ "output/" ] +[[tool.mypy.overrides]] +module = ["timdex_dataset_api.*"] +follow_untyped_imports = true + + [tool.pytest.ini_options] log_level = "INFO" @@ -88,6 +94,7 @@ fixture-parentheses = false "tests/**/*" = [ "ANN", "ARG001", + "PLR2004", "S101", ] diff --git a/tests/conftest.py b/tests/conftest.py index 3a5fbf3..ec98e47 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -6,6 +6,8 @@ import pytest from click.testing import CliRunner +from embeddings.embedding import Embedding, RecordText +from embeddings.models import registry from embeddings.models.base import BaseEmbeddingModel logger = logging.getLogger(__name__) @@ -43,6 +45,16 @@ def download(self) -> Path: def load(self) -> None: logger.info("Model loaded successfully, 1.5s") + def create_embedding(self, input_record: RecordText) -> Embedding: + return Embedding( + timdex_record_id=input_record.timdex_record_id, + run_id=input_record.run_id, + run_record_offset=input_record.run_record_offset, + embedding_strategy=input_record.embedding_strategy, + model_uri=self.model_uri, + embedding={"coffee": 0.9, "seattle": 0.5}, + ) + @pytest.fixture def mock_model(tmp_path): @@ -50,6 +62,13 @@ def mock_model(tmp_path): return MockEmbeddingModel(tmp_path / "model") +@pytest.fixture +def register_mock_model(monkeypatch): + """Register MockEmbeddingModel in the model registry.""" + monkeypatch.setitem(registry.MODEL_REGISTRY, "test/mock-model", MockEmbeddingModel) + monkeypatch.setenv("TE_MODEL_PATH", "/fake/path") + + @pytest.fixture def neural_sparse_doc_v3_gte_fake_model_directory(tmp_path): """Create a fake downloaded model directory with required files.""" diff --git a/tests/test_cli.py b/tests/test_cli.py index 63a466c..f360be6 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1,6 +1,4 @@ from embeddings.cli import main -from embeddings.models import registry -from tests.conftest import MockEmbeddingModel def test_cli_default_logging(caplog, runner): @@ -27,10 +25,10 @@ def test_download_model_unknown_uri(caplog, runner): assert "Unknown model URI" in caplog.text -def test_model_required_decorator_with_cli_option(caplog, monkeypatch, runner, tmp_path): +def test_model_required_decorator_with_cli_option( + caplog, register_mock_model, runner, tmp_path +): """Test decorator successfully initializes model from --model-uri option.""" - monkeypatch.setitem(registry.MODEL_REGISTRY, "test/mock-model", MockEmbeddingModel) - output_path = tmp_path / "model.zip" result = runner.invoke( main, @@ -51,9 +49,10 @@ def test_model_required_decorator_with_cli_option(caplog, monkeypatch, runner, t assert output_path.exists() -def test_model_required_decorator_with_env_var(caplog, monkeypatch, runner, tmp_path): +def test_model_required_decorator_with_env_var( + caplog, monkeypatch, register_mock_model, runner, tmp_path +): """Test decorator successfully initializes model from TE_MODEL_URI env var.""" - monkeypatch.setitem(registry.MODEL_REGISTRY, "test/mock-model", MockEmbeddingModel) monkeypatch.setenv("TE_MODEL_URI", "test/mock-model") output_path = tmp_path / "model.zip" @@ -76,11 +75,9 @@ def test_model_required_decorator_missing_parameter(runner): def test_model_required_decorator_stores_model_in_context( - caplog, monkeypatch, runner, tmp_path + caplog, register_mock_model, runner, tmp_path ): """Test decorator stores model instance in ctx.obj['model'].""" - monkeypatch.setitem(registry.MODEL_REGISTRY, "test/mock-model", MockEmbeddingModel) - output_path = tmp_path / "model.zip" result = runner.invoke( main, @@ -99,10 +96,10 @@ def test_model_required_decorator_stores_model_in_context( assert output_path.exists() -def test_model_required_decorator_log_message(caplog, monkeypatch, runner, tmp_path): +def test_model_required_decorator_log_message( + caplog, register_mock_model, runner, tmp_path +): """Test decorator logs correct initialization message.""" - monkeypatch.setitem(registry.MODEL_REGISTRY, "test/mock-model", MockEmbeddingModel) - output_path = tmp_path / "model.zip" result = runner.invoke( main, @@ -122,11 +119,10 @@ def test_model_required_decorator_log_message(caplog, monkeypatch, runner, tmp_p ) -def test_model_required_decorator_works_across_commands(caplog, monkeypatch, runner): +def test_model_required_decorator_works_across_commands( + caplog, register_mock_model, runner +): """Test decorator works for multiple commands (test_model_load).""" - monkeypatch.setitem(registry.MODEL_REGISTRY, "test/mock-model", MockEmbeddingModel) - monkeypatch.setenv("TE_MODEL_PATH", "/fake/path") - result = runner.invoke(main, ["test-model-load", "--model-uri", "test/mock-model"]) assert result.exit_code == 0 @@ -135,3 +131,41 @@ def test_model_required_decorator_works_across_commands(caplog, monkeypatch, run "'test/mock-model'" in caplog.text ) assert "OK" in result.output + + +def test_create_embeddings_requires_dataset_location(register_mock_model, runner): + result = runner.invoke(main, ["create-embeddings", "--model-uri", "test/mock-model"]) + assert result.exit_code != 0 + assert "--dataset-location" in result.output + + +def test_create_embeddings_requires_run_id(register_mock_model, runner): + result = runner.invoke( + main, + [ + "create-embeddings", + "--model-uri", + "test/mock-model", + "--dataset-location", + "s3://test", + ], + ) + assert result.exit_code != 0 + assert "Missing option '--run-id'" in result.output + + +def test_create_embeddings_requires_strategy(register_mock_model, runner): + result = runner.invoke( + main, + [ + "create-embeddings", + "--model-uri", + "test/mock-model", + "--dataset-location", + "s3://test", + "--run-id", + "run-1", + ], + ) + assert result.exit_code != 0 + assert "Missing option '--strategy'" in result.output diff --git a/tests/test_models.py b/tests/test_models.py index 9fc8f24..6948a6c 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -2,37 +2,56 @@ import pytest +from embeddings.embedding import RecordText from embeddings.models.base import BaseEmbeddingModel from embeddings.models.registry import MODEL_REGISTRY, get_model_class -from tests.conftest import MockEmbeddingModel def test_mock_model_instantiation(mock_model): assert mock_model.model_uri == "test/mock-model" -def test_mock_model_download_creates_zip(tmp_path): - output_path = tmp_path / "test_model.zip" - mock_model = MockEmbeddingModel(output_path) +def test_mock_model_download_creates_zip(mock_model): result = mock_model.download() - assert result == output_path - assert output_path.exists() - assert zipfile.is_zipfile(output_path) + assert result == mock_model.model_path + assert mock_model.model_path.exists() + assert zipfile.is_zipfile(mock_model.model_path) -def test_mock_model_download_contains_expected_files(tmp_path): - output_path = tmp_path / "test_model.zip" - mock_model = MockEmbeddingModel(output_path) +def test_mock_model_download_contains_expected_files(mock_model): mock_model.download() - with zipfile.ZipFile(output_path, "r") as zf: + with zipfile.ZipFile(mock_model.model_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_mock_model_load(caplog, mock_model): + mock_model.load() + assert "Model loaded successfully" in caplog.text + + +def test_mock_model_create_embedding(mock_model): + input_record = RecordText( + timdex_record_id="test-id", + run_id="test-run", + run_record_offset=42, + embedding_strategy="full_record", + text="test text", + ) + embedding = mock_model.create_embedding(input_record) + + assert embedding.timdex_record_id == "test-id" + assert embedding.run_id == "test-run" + assert embedding.run_record_offset == 42 + assert embedding.embedding_strategy == "full_record" + assert embedding.model_uri == "test/mock-model" + assert embedding.embedding == {"coffee": 0.9, "seattle": 0.5} + + def test_registry_contains_opensearch_model(): assert ( "opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte" @@ -64,3 +83,29 @@ def test_subclass_with_non_string_model_uri_raises_type_error(): class InvalidModel(BaseEmbeddingModel): MODEL_URI = 123 + + +def test_base_model_create_embeddings_calls_create_embedding(mock_model): + input_records = [ + RecordText( + timdex_record_id="id-1", + run_id="run-1", + run_record_offset=0, + embedding_strategy="full_record", + text="text 1", + ), + RecordText( + timdex_record_id="id-2", + run_id="run-1", + run_record_offset=1, + embedding_strategy="full_record", + text="text 2", + ), + ] + + # create_embeddings should iterate and call create_embedding + embeddings = list(mock_model.create_embeddings(iter(input_records))) + + assert len(embeddings) == 2 # two input records + assert embeddings[0].timdex_record_id == "id-1" + assert embeddings[1].timdex_record_id == "id-2" diff --git a/tests/test_os_neural_sparse_doc_v3_gte.py b/tests/test_os_neural_sparse_doc_v3_gte.py index 3e093eb..462ae85 100644 --- a/tests/test_os_neural_sparse_doc_v3_gte.py +++ b/tests/test_os_neural_sparse_doc_v3_gte.py @@ -1,6 +1,6 @@ """Tests for OSNeuralSparseDocV3GTE embedding model.""" -# ruff: noqa: SLF001, PLR2004 +# ruff: noqa: SLF001 import json from pathlib import Path diff --git a/uv.lock b/uv.lock index 072bb83..1dfd0db 100644 --- a/uv.lock +++ b/uv.lock @@ -559,6 +559,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/31/b4/b9b800c45527aadd64d5b442f9b932b00648617eb5d63d2c7a6587b7cafc/jmespath-1.0.1-py3-none-any.whl", hash = "sha256:02e2e4cc71b5bcab88332eebf907519190dd9e6e82107fa7f83b1003a6252980", size = 20256, upload-time = "2022-06-17T18:00:10.251Z" }, ] +[[package]] +name = "jsonlines" +version = "4.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "attrs" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/35/87/bcda8e46c88d0e34cad2f09ee2d0c7f5957bccdb9791b0b934ec84d84be4/jsonlines-4.0.0.tar.gz", hash = "sha256:0c6d2c09117550c089995247f605ae4cf77dd1533041d366351f6f298822ea74", size = 11359, upload-time = "2023-09-01T12:34:44.187Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f8/62/d9ba6323b9202dd2fe166beab8a86d29465c41a0288cbe229fac60c1ab8d/jsonlines-4.0.0-py3-none-any.whl", hash = "sha256:185b334ff2ca5a91362993f42e83588a360cf95ce4b71a73548502bda52a7c55", size = 8701, upload-time = "2023-09-01T12:34:42.563Z" }, +] + [[package]] name = "license-expression" version = "30.4.4" @@ -1622,6 +1634,7 @@ source = { editable = "." } dependencies = [ { name = "click" }, { name = "huggingface-hub" }, + { name = "jsonlines" }, { name = "sentry-sdk" }, { name = "timdex-dataset-api" }, { name = "torch" }, @@ -1644,6 +1657,7 @@ dev = [ requires-dist = [ { name = "click", specifier = ">=8.2.1" }, { name = "huggingface-hub", specifier = ">=0.26.0" }, + { name = "jsonlines", specifier = ">=4.0.0" }, { name = "sentry-sdk", specifier = ">=2.34.1" }, { name = "timdex-dataset-api", git = "https://github.com/MITLibraries/timdex-dataset-api" }, { name = "torch", specifier = ">=2.9.0" }, From 0024b8f69ed479bda8359cd6043027342c2ea32f Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Thu, 30 Oct 2025 15:15:33 -0400 Subject: [PATCH 2/5] Verbose mode targets desired logging familes Why these changes are being introduced: Previously, when --verbose was set for the CLI all loggers inherited this. In other applications, we have used a 'WARNING_ONLY_LOGGERS' env var that would limit them to WARNING level. This worked, but was perhaps not ideal. Without that env var, it's a bit of whack-a-mole to figure out which loggers to quiet. How this addresses that need: Instead of defaulting all loggers to DEBUG in verbose mode, we target only libraries we expect this application to log in DEBUG. By default, all other logger families will still have WARNING. This may be a pattern we want to explore in other repositories. Potentially even further inverting the pattern and supporting a 'DEBUG_LOGGERS' env var list that would explicitly toggle on DEBUG logging for those libraries. That would allow troubleshooting in deployed environments just by setting an env var. This is NOT applied in this commit, but noting for future consideration. Side effects of this change: * Both 'embeddings' and 'timdex_dataset_api' are logged as DEBUG in verbose mode, but only those libraries. Relevant ticket(s): * https://mitlibraries.atlassian.net/browse/USE-112 --- embeddings/config.py | 3 ++- tests/test_config.py | 4 ++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/embeddings/config.py b/embeddings/config.py index 21e094e..de474b3 100644 --- a/embeddings/config.py +++ b/embeddings/config.py @@ -10,7 +10,8 @@ def configure_logger(logger: logging.Logger, *, verbose: bool) -> str: format="%(asctime)s %(levelname)s %(name)s.%(funcName)s() line %(lineno)d: " "%(message)s" ) - logger.setLevel(logging.DEBUG) + logging.getLogger("embeddings").setLevel(logging.DEBUG) + logging.getLogger("timdex_dataset_api").setLevel(logging.DEBUG) else: logging.basicConfig( format="%(asctime)s %(levelname)s %(name)s.%(funcName)s(): %(message)s" diff --git a/tests/test_config.py b/tests/test_config.py index 4ca42ab..ad8606d 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -12,11 +12,11 @@ def test_configure_logger_not_verbose(): def test_configure_logger_verbose(): - logger = logging.getLogger(__name__) + logger = logging.getLogger("embeddings.cli") result = configure_logger(logger, verbose=True) debug_log_level = 10 assert logger.getEffectiveLevel() == debug_log_level - assert result == "Logger 'tests.test_config' configured with level=DEBUG" + assert result == "Logger 'embeddings.cli' configured with level=DEBUG" def test_configure_sentry_no_env_variable(monkeypatch): From d433f171cdac680424f0713f01f96a84425a045f Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Mon, 3 Nov 2025 09:09:57 -0500 Subject: [PATCH 3/5] Rename RecordText to EmbeddingInput Why these changes are being introduced: Code review suggested that 'RecordText' was a confusing name for the object that we prepare to then create an embedding from. How this addresses that need: Renamign to 'EmbeddingInput' makes it crystal clear that we are preparing an object that will be used to create an embedding. Side effects of this change: * None Relevant ticket(s): * https://mitlibraries.atlassian.net/browse/USE-112 --- embeddings/cli.py | 4 ++-- embeddings/embedding.py | 8 ++++++-- embeddings/models/base.py | 10 +++++----- embeddings/models/os_neural_sparse_doc_v3_gte.py | 4 ++-- tests/conftest.py | 4 ++-- tests/test_models.py | 8 ++++---- 6 files changed, 21 insertions(+), 17 deletions(-) diff --git a/embeddings/cli.py b/embeddings/cli.py index 3e061b8..1470199 100644 --- a/embeddings/cli.py +++ b/embeddings/cli.py @@ -230,10 +230,10 @@ def create_embeddings( # DEBUG ------------------------------------------------------------------------------ import json # noqa: PLC0415 - from embeddings.embedding import RecordText # noqa: PLC0415 + from embeddings.embedding import EmbeddingInput # noqa: PLC0415 input_records = ( - RecordText( + EmbeddingInput( timdex_record_id=timdex_record["timdex_record_id"], run_id=timdex_record["run_id"], run_record_offset=timdex_record["run_record_offset"], diff --git a/embeddings/embedding.py b/embeddings/embedding.py index 596c567..f9434cb 100644 --- a/embeddings/embedding.py +++ b/embeddings/embedding.py @@ -4,8 +4,12 @@ @dataclass -class RecordText: - """Input record for creating an embedding for. +class EmbeddingInput: + """Encapsulates the inputs for an embedding. + + When creating an embedding, we need to note what TIMDEX record the embedding is + associated with and what strategy was used to prepare the embedding input text from + the record itself. Args: (timdex_record_id, run_id, run_record_offset): composite key for TIMDEX record diff --git a/embeddings/models/base.py b/embeddings/models/base.py index ac08cc5..79ae683 100644 --- a/embeddings/models/base.py +++ b/embeddings/models/base.py @@ -4,7 +4,7 @@ from collections.abc import Iterator from pathlib import Path -from embeddings.embedding import Embedding, RecordText +from embeddings.embedding import Embedding, EmbeddingInput class BaseEmbeddingModel(ABC): @@ -51,15 +51,15 @@ def load(self) -> None: """Load model from self.model_path.""" @abstractmethod - def create_embedding(self, input_record: RecordText) -> Embedding: - """Create an Embedding for an RecordText. + def create_embedding(self, input_record: EmbeddingInput) -> Embedding: + """Create an Embedding for an EmbeddingInput. Args: - input_record: RecordText instance + input_record: EmbeddingInput instance """ def create_embeddings( - self, input_records: Iterator[RecordText] + self, input_records: Iterator[EmbeddingInput] ) -> Iterator[Embedding]: """Yield Embeddings for an iterator of InputRecords. diff --git a/embeddings/models/os_neural_sparse_doc_v3_gte.py b/embeddings/models/os_neural_sparse_doc_v3_gte.py index 37fae99..8b7eee0 100644 --- a/embeddings/models/os_neural_sparse_doc_v3_gte.py +++ b/embeddings/models/os_neural_sparse_doc_v3_gte.py @@ -11,7 +11,7 @@ from huggingface_hub import snapshot_download from transformers import AutoModelForMaskedLM, AutoTokenizer -from embeddings.embedding import Embedding, RecordText +from embeddings.embedding import Embedding, EmbeddingInput from embeddings.models.base import BaseEmbeddingModel if TYPE_CHECKING: @@ -163,5 +163,5 @@ def load(self) -> None: logger.info(f"Model loaded successfully, {time.perf_counter()-start_time}s") - def create_embedding(self, input_record: RecordText) -> Embedding: + def create_embedding(self, input_record: EmbeddingInput) -> Embedding: raise NotImplementedError diff --git a/tests/conftest.py b/tests/conftest.py index ec98e47..63a54ab 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -6,7 +6,7 @@ import pytest from click.testing import CliRunner -from embeddings.embedding import Embedding, RecordText +from embeddings.embedding import Embedding, EmbeddingInput from embeddings.models import registry from embeddings.models.base import BaseEmbeddingModel @@ -45,7 +45,7 @@ def download(self) -> Path: def load(self) -> None: logger.info("Model loaded successfully, 1.5s") - def create_embedding(self, input_record: RecordText) -> Embedding: + def create_embedding(self, input_record: EmbeddingInput) -> Embedding: return Embedding( timdex_record_id=input_record.timdex_record_id, run_id=input_record.run_id, diff --git a/tests/test_models.py b/tests/test_models.py index 6948a6c..1c2ad0d 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -2,7 +2,7 @@ import pytest -from embeddings.embedding import RecordText +from embeddings.embedding import EmbeddingInput from embeddings.models.base import BaseEmbeddingModel from embeddings.models.registry import MODEL_REGISTRY, get_model_class @@ -35,7 +35,7 @@ def test_mock_model_load(caplog, mock_model): def test_mock_model_create_embedding(mock_model): - input_record = RecordText( + input_record = EmbeddingInput( timdex_record_id="test-id", run_id="test-run", run_record_offset=42, @@ -87,14 +87,14 @@ class InvalidModel(BaseEmbeddingModel): def test_base_model_create_embeddings_calls_create_embedding(mock_model): input_records = [ - RecordText( + EmbeddingInput( timdex_record_id="id-1", run_id="run-1", run_record_offset=0, embedding_strategy="full_record", text="text 1", ), - RecordText( + EmbeddingInput( timdex_record_id="id-2", run_id="run-1", run_record_offset=1, From de351a1024d2a4c42eb155a456c0657941e601b5 Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Mon, 3 Nov 2025 09:14:05 -0500 Subject: [PATCH 4/5] Update README --- README.md | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/README.md b/README.md index 0521d50..1295ad6 100644 --- a/README.md +++ b/README.md @@ -95,4 +95,28 @@ Options: --help Show this message and exit. ``` +### `create-embeddings` +```text +Usage: embeddings create-embeddings [OPTIONS] + + Create embeddings for TIMDEX records. +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] + -d, --dataset-location PATH TIMDEX dataset location, e.g. + 's3://timdex/dataset', to read records from. + [required] + --run-id TEXT TIMDEX ETL run id. [required] + --run-record-offset INTEGER TIMDEX ETL run record offset to start from, + default = 0. [required] + --record-limit INTEGER Limit number of records after --run-record- + offset, default = None (unlimited). [required] + --strategy TEXT Pre-embedding record transformation strategy to + use. Repeatable. [required] + --output-jsonl TEXT Optionally write embeddings to local JSONLines + file (primarily for testing). + --help Show this message and exit. +``` From d4e6e5450914a16660609bc5272f4a27dc6b498e Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Mon, 3 Nov 2025 09:21:04 -0500 Subject: [PATCH 5/5] Store sparse vector and decoded token weights Why these changes are being introduced: Formerly, our 'Embedding' class only had an 'embedding' property for the output. However, for our first model in the pipeline, opensearch-project/ opensearch-neural-sparse-encoding-doc-v3-gte, it produces two representations of the embedding that are useful to store: a sparse vector and decoded token weights. How this addresses that need: Updates the 'Embedding' class to explicitly store both representations of the embedding. We may decide that we don't store both, or some futures models may not produce decoded token weights of any kind, but this matches our first proposed model and pipeline. Better to be explicit and opinionated in these early days, then adjust later if needed. Side effects of this change: * None Relevant ticket(s): * https://mitlibraries.atlassian.net/browse/USE-112 --- embeddings/cli.py | 3 ++- embeddings/embedding.py | 7 +++++-- tests/conftest.py | 3 ++- tests/test_models.py | 3 ++- 4 files changed, 11 insertions(+), 5 deletions(-) diff --git a/embeddings/cli.py b/embeddings/cli.py index 1470199..3c459ee 100644 --- a/embeddings/cli.py +++ b/embeddings/cli.py @@ -258,7 +258,8 @@ def create_embeddings( run_record_offset=input_record.run_record_offset, embedding_strategy=input_record.embedding_strategy, model_uri=model.model_uri, - embedding={"coffee": 0.9, "seattle": 0.5}, + embedding_vector=[0.1, 0.2, 0.3], + embedding_token_weights={"coffee": 0.9, "seattle": 0.5}, ) for input_record in input_records ) diff --git a/embeddings/embedding.py b/embeddings/embedding.py index f9434cb..d0c8187 100644 --- a/embeddings/embedding.py +++ b/embeddings/embedding.py @@ -32,7 +32,9 @@ class Embedding: (timdex_record_id, run_id, run_record_offset): composite key for TIMDEX record model_uri: model URI used to create the embedding embedding_strategy: strategy used to create text for embedding - embedding: model embedding created from text + embedding_vector: vector representation of embedding + embedding_token_weights: decoded token:weight pairs from sparse vector + - only applicable to models that produce this output """ timdex_record_id: str @@ -40,7 +42,8 @@ class Embedding: run_record_offset: int model_uri: str embedding_strategy: str - embedding: dict | list[float] + embedding_vector: list[float] + embedding_token_weights: dict timestamp: datetime.datetime = field( default_factory=lambda: datetime.datetime.now(datetime.UTC) diff --git a/tests/conftest.py b/tests/conftest.py index 63a54ab..6f0d7df 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -52,7 +52,8 @@ def create_embedding(self, input_record: EmbeddingInput) -> Embedding: run_record_offset=input_record.run_record_offset, embedding_strategy=input_record.embedding_strategy, model_uri=self.model_uri, - embedding={"coffee": 0.9, "seattle": 0.5}, + embedding_vector=[0.1, 0.2, 0.3], + embedding_token_weights={"coffee": 0.9, "seattle": 0.5}, ) diff --git a/tests/test_models.py b/tests/test_models.py index 1c2ad0d..d8f1734 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -49,7 +49,8 @@ def test_mock_model_create_embedding(mock_model): assert embedding.run_record_offset == 42 assert embedding.embedding_strategy == "full_record" assert embedding.model_uri == "test/mock-model" - assert embedding.embedding == {"coffee": 0.9, "seattle": 0.5} + assert embedding.embedding_vector == [0.1, 0.2, 0.3] + assert embedding.embedding_token_weights == {"coffee": 0.9, "seattle": 0.5} def test_registry_contains_opensearch_model():