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/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. +``` diff --git a/embeddings/cli.py b/embeddings/cli.py index f68212b..3c459ee 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,140 @@ 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 EmbeddingInput # noqa: PLC0415 + + input_records = ( + EmbeddingInput( + 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_vector=[0.1, 0.2, 0.3], + embedding_token_weights={"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..de474b3 100644 --- a/embeddings/config.py +++ b/embeddings/config.py @@ -10,9 +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) - for handler in logging.root.handlers: - handler.addFilter(logging.Filter("embeddings")) + 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/embeddings/embedding.py b/embeddings/embedding.py new file mode 100644 index 0000000..d0c8187 --- /dev/null +++ b/embeddings/embedding.py @@ -0,0 +1,58 @@ +import datetime +import json +from dataclasses import asdict, dataclass, field + + +@dataclass +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 + 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_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 + run_id: str + run_record_offset: int + model_uri: str + embedding_strategy: str + embedding_vector: list[float] + embedding_token_weights: dict + + 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..79ae683 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, EmbeddingInput + 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: EmbeddingInput) -> Embedding: + """Create an Embedding for an EmbeddingInput. + + Args: + input_record: EmbeddingInput instance + """ + + def create_embeddings( + self, input_records: Iterator[EmbeddingInput] + ) -> 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..8b7eee0 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, EmbeddingInput 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: EmbeddingInput) -> 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..6f0d7df 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, EmbeddingInput +from embeddings.models import registry from embeddings.models.base import BaseEmbeddingModel logger = logging.getLogger(__name__) @@ -43,6 +45,17 @@ def download(self) -> Path: def load(self) -> None: logger.info("Model loaded successfully, 1.5s") + def create_embedding(self, input_record: EmbeddingInput) -> 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_vector=[0.1, 0.2, 0.3], + embedding_token_weights={"coffee": 0.9, "seattle": 0.5}, + ) + @pytest.fixture def mock_model(tmp_path): @@ -50,6 +63,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_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): diff --git a/tests/test_models.py b/tests/test_models.py index 9fc8f24..d8f1734 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -2,37 +2,57 @@ import pytest +from embeddings.embedding import EmbeddingInput 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 = EmbeddingInput( + 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_vector == [0.1, 0.2, 0.3] + assert embedding.embedding_token_weights == {"coffee": 0.9, "seattle": 0.5} + + def test_registry_contains_opensearch_model(): assert ( "opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte" @@ -64,3 +84,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 = [ + EmbeddingInput( + timdex_record_id="id-1", + run_id="run-1", + run_record_offset=0, + embedding_strategy="full_record", + text="text 1", + ), + EmbeddingInput( + 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" },