From 4e57fbb97ef1dc03927c983d9ec8065e193c35f1 Mon Sep 17 00:00:00 2001 From: Oliver Holworthy <1216955+oliverholworthy@users.noreply.github.com> Date: Thu, 20 Aug 2026 20:48:32 +0100 Subject: [PATCH 1/2] fix(retrieval): support Sentence Transformers v6 metadata Signed-off-by: Oliver Holworthy <1216955+oliverholworthy@users.noreply.github.com> --- .../sentence_transformer_export.py | 23 +++++ .../test_sentence_transformer_export.py | 85 +++++++++++++++++++ uv.lock | 7 +- 3 files changed, 112 insertions(+), 3 deletions(-) diff --git a/nemo_automodel/_transformers/sentence_transformer_export.py b/nemo_automodel/_transformers/sentence_transformer_export.py index 9e5377adc6..f04800824b 100644 --- a/nemo_automodel/_transformers/sentence_transformer_export.py +++ b/nemo_automodel/_transformers/sentence_transformer_export.py @@ -64,9 +64,11 @@ "normalize": { "sentence_transformers.models.Normalize", "sentence_transformers.sentence_transformer.modules.normalize.Normalize", + "sentence_transformers.base.modules.normalize.Normalize", }, } _SENTENCE_TRANSFORMER_EXPORT_MODULE_TYPES = { + # v6 remaps these v5.4-era paths, keeping exported checkpoints loadable across v5.4+. "transformer": "sentence_transformers.base.modules.transformer.Transformer", "pooling": "sentence_transformers.sentence_transformer.modules.pooling.Pooling", "normalize": "sentence_transformers.sentence_transformer.modules.normalize.Normalize", @@ -172,6 +174,27 @@ def _load_sentence_transformer_wrapper_options( if modules[0].get("path") != "": raise ValueError("Sentence Transformers Transformer metadata must reference the checkpoint root.") + if len(modules) == 3: + normalize_path = modules[2].get("path") + if not isinstance(normalize_path, str) or not normalize_path: + raise ValueError("Sentence Transformers Normalize metadata must reference a module path.") + normalize_config = _load_sentence_transformer_json( + model_name_or_path, + os.path.join(normalize_path, "config.json"), + hf_kwargs, + ) + if normalize_config is not None: + if not isinstance(normalize_config, dict): + raise ValueError("Sentence Transformers Normalize config.json is invalid.") + normalize_input_name = normalize_config.get("module_input_name", "sentence_embedding") + normalize_output_name = normalize_config.get("module_output_name") + if normalize_output_name is None: + normalize_output_name = normalize_input_name + if normalize_input_name != "sentence_embedding" or normalize_output_name != "sentence_embedding": + raise ValueError( + "Sentence Transformers Normalize metadata must normalize the final sentence embedding in place." + ) + sentence_bert_config = _load_sentence_transformer_json( model_name_or_path, "sentence_bert_config.json", diff --git a/tests/unit_tests/_transformers/test_sentence_transformer_export.py b/tests/unit_tests/_transformers/test_sentence_transformer_export.py index 6f83bdadc0..664fdbfe50 100644 --- a/tests/unit_tests/_transformers/test_sentence_transformer_export.py +++ b/tests/unit_tests/_transformers/test_sentence_transformer_export.py @@ -452,6 +452,91 @@ def test_sentence_transformer_source_accepts_canonical_module_paths(tmp_path, mo assert options.l2_normalize is expected_normalize +def test_sentence_transformer_source_accepts_v6_normalize_module(tmp_path): + (tmp_path / "1_Pooling").mkdir() + (tmp_path / "2_Normalize").mkdir() + (tmp_path / "modules.json").write_text( + json.dumps( + [ + { + "idx": 0, + "path": "", + "type": "sentence_transformers.base.modules.transformer.Transformer", + }, + { + "idx": 1, + "path": "1_Pooling", + "type": "sentence_transformers.sentence_transformer.modules.pooling.Pooling", + }, + { + "idx": 2, + "path": "2_Normalize", + "type": "sentence_transformers.base.modules.normalize.Normalize", + }, + ] + ) + ) + (tmp_path / "1_Pooling" / "config.json").write_text(json.dumps({"pooling_mode": "mean"})) + (tmp_path / "2_Normalize" / "config.json").write_text( + json.dumps( + { + "module_input_name": "sentence_embedding", + "module_output_name": "sentence_embedding", + } + ) + ) + + options = sentence_transformer_export._load_sentence_transformer_wrapper_options(str(tmp_path), {}) + + assert options is not None + assert options.pooling == "avg" + assert options.l2_normalize is True + + +@pytest.mark.parametrize( + "normalize_config", + [ + { + "module_input_name": "token_embeddings", + "module_output_name": "token_embeddings", + }, + { + "module_input_name": "sentence_embedding", + "module_output_name": "custom_embedding", + }, + ], +) +def test_sentence_transformer_source_rejects_unrepresentable_v6_normalize_config(tmp_path, normalize_config): + (tmp_path / "1_Pooling").mkdir() + (tmp_path / "2_Normalize").mkdir() + (tmp_path / "modules.json").write_text( + json.dumps( + [ + { + "idx": 0, + "path": "", + "type": "sentence_transformers.base.modules.transformer.Transformer", + }, + { + "idx": 1, + "path": "1_Pooling", + "type": "sentence_transformers.sentence_transformer.modules.pooling.Pooling", + }, + { + "idx": 2, + "path": "2_Normalize", + "type": "sentence_transformers.base.modules.normalize.Normalize", + }, + ] + ) + ) + (tmp_path / "1_Pooling" / "config.json").write_text(json.dumps({"pooling_mode": "mean"})) + (tmp_path / "2_Normalize" / "config.json").write_text(json.dumps(normalize_config)) + + with pytest.raises(ValueError, match="final sentence embedding"): + sentence_transformer_export._load_sentence_transformer_wrapper_options(str(tmp_path), {}) + + @pytest.mark.parametrize( "module_types", [ diff --git a/uv.lock b/uv.lock index f84c62d080..f48e100cef 100644 --- a/uv.lock +++ b/uv.lock @@ -7020,7 +7020,7 @@ wheels = [ [[package]] name = "sentence-transformers" -version = "5.6.0" +version = "6.0.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "huggingface-hub" }, @@ -7029,6 +7029,7 @@ dependencies = [ { name = "scikit-learn" }, { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "scipy", version = "1.16.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "tokenizers" }, { name = "torch", version = "2.10.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin'" }, { name = "torch", version = "2.10.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "sys_platform != 'darwin' and sys_platform != 'linux'" }, { name = "torch", version = "2.10.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'linux'" }, @@ -7036,9 +7037,9 @@ dependencies = [ { name = "transformers" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/f9/56/d2cb00765a6b15c994a7fccf20f9032f16e8193ca49147cb5155166ad744/sentence_transformers-5.6.0.tar.gz", hash = "sha256:0e7164d051e416c1853ade7c274ff52af3f9da0f4be7f0b83d734c27699e1057", size = 453194, upload-time = "2026-06-16T14:01:56.42Z" } +sdist = { url = "https://files.pythonhosted.org/packages/e7/43/6b53e6a2098440ce21478742facbc058f1a66ba2cb80b24bdc64942e1e2c/sentence_transformers-6.0.0.tar.gz", hash = "sha256:9e8c2c24f3b1c7473cd5f519a3d3cff60daaeb95533b82d045ffb43ee5f2dac4", size = 575048, upload-time = "2026-08-18T13:33:49.919Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/76/c1/dc1582b79e9a2eb0cddf9559cd9bcdff084f541d6fe881fdd9d98630dba7/sentence_transformers-5.6.0-py3-none-any.whl", hash = "sha256:d2075b5e687a1611005e20ab04a6846994d51adfcf39610aed066af3c0c0b81f", size = 596411, upload-time = "2026-06-16T14:01:55.103Z" }, + { url = "https://files.pythonhosted.org/packages/04/fe/9d19b01fe87945f9455c617bf5d33dfbf29fe06ab6580bc0bea06080c788/sentence_transformers-6.0.0-py3-none-any.whl", hash = "sha256:b974ac67523ea2a955afa87b1024129305472bd884367dcc969261cb086790e9", size = 739640, upload-time = "2026-08-18T13:33:48.428Z" }, ] [[package]] From 40bfc8cdb2344062b60f232119b58ae15e14c833 Mon Sep 17 00:00:00 2001 From: Oliver Holworthy <1216955+oliverholworthy@users.noreply.github.com> Date: Fri, 21 Aug 2026 10:06:19 +0100 Subject: [PATCH 2/2] refactor(retrieval): validate sentence transformer module stack Signed-off-by: Oliver Holworthy <1216955+oliverholworthy@users.noreply.github.com> --- .../sentence_transformer_export.py | 44 ++++++++++++------- .../test_sentence_transformer_export.py | 28 +++++++++--- 2 files changed, 50 insertions(+), 22 deletions(-) diff --git a/nemo_automodel/_transformers/sentence_transformer_export.py b/nemo_automodel/_transformers/sentence_transformer_export.py index f04800824b..2a8b346c9e 100644 --- a/nemo_automodel/_transformers/sentence_transformer_export.py +++ b/nemo_automodel/_transformers/sentence_transformer_export.py @@ -67,6 +67,10 @@ "sentence_transformers.base.modules.normalize.Normalize", }, } +_SUPPORTED_SENTENCE_TRANSFORMER_MODULE_STACKS = { + ("transformer", "pooling"), + ("transformer", "pooling", "normalize"), +} _SENTENCE_TRANSFORMER_EXPORT_MODULE_TYPES = { # v6 remaps these v5.4-era paths, keeping exported checkpoints loadable across v5.4+. "transformer": "sentence_transformers.base.modules.transformer.Transformer", @@ -156,26 +160,34 @@ def _load_sentence_transformer_wrapper_options( if not isinstance(modules, list): raise ValueError("Sentence Transformers modules.json must contain a list of modules.") - expected_module_types = [ - _SENTENCE_TRANSFORMER_MODULE_TYPES["transformer"], - _SENTENCE_TRANSFORMER_MODULE_TYPES["pooling"], - ] - module_types = [module.get("type") if isinstance(module, dict) else None for module in modules] - if len(module_types) == 3: - expected_module_types.append(_SENTENCE_TRANSFORMER_MODULE_TYPES["normalize"]) - if len(module_types) not in (2, 3) or any( - module_type not in allowed_types - for module_type, allowed_types in zip(module_types, expected_module_types, strict=True) - ): + module_roles = [] + for module in modules: + module_type = module.get("type") if isinstance(module, dict) else None + module_role = next( + ( + role + for role, allowed_types in _SENTENCE_TRANSFORMER_MODULE_TYPES.items() + if module_type in allowed_types + ), + None, + ) + module_roles.append(module_role) + module_roles = tuple(module_roles) + if module_roles not in _SUPPORTED_SENTENCE_TRANSFORMER_MODULE_STACKS: raise ValueError( "Sentence Transformers metadata must use the exact supported module stack: " "Transformer, Pooling, and optional Normalize." ) - if modules[0].get("path") != "": + + modules_by_role = dict(zip(module_roles, modules, strict=True)) + transformer_module = modules_by_role["transformer"] + pooling_module = modules_by_role["pooling"] + normalize_module = modules_by_role.get("normalize") + if transformer_module.get("path") != "": raise ValueError("Sentence Transformers Transformer metadata must reference the checkpoint root.") - if len(modules) == 3: - normalize_path = modules[2].get("path") + if normalize_module is not None: + normalize_path = normalize_module.get("path") if not isinstance(normalize_path, str) or not normalize_path: raise ValueError("Sentence Transformers Normalize metadata must reference a module path.") normalize_config = _load_sentence_transformer_json( @@ -206,7 +218,7 @@ def _load_sentence_transformer_wrapper_options( "the NeMo inference path does not lowercase text." ) - pooling_path = modules[1].get("path") + pooling_path = pooling_module.get("path") if not isinstance(pooling_path, str) or not pooling_path: raise ValueError("Sentence Transformers Pooling metadata must reference a module path.") pooling_config = _load_sentence_transformer_json( @@ -269,7 +281,7 @@ def _load_sentence_transformer_wrapper_options( return SentenceTransformerWrapperOptions( pooling=matching_pooling[0], - l2_normalize=len(modules) == 3, + l2_normalize=normalize_module is not None, query_prompt=query_prompt, document_prompt=document_prompt, ) diff --git a/tests/unit_tests/_transformers/test_sentence_transformer_export.py b/tests/unit_tests/_transformers/test_sentence_transformer_export.py index 664fdbfe50..199423831f 100644 --- a/tests/unit_tests/_transformers/test_sentence_transformer_export.py +++ b/tests/unit_tests/_transformers/test_sentence_transformer_export.py @@ -371,18 +371,34 @@ def test_cache_hub_source_legal_assets_uses_exact_loaded_revision(monkeypatch): ) -def test_sentence_transformer_source_with_unsupported_module_is_rejected(tmp_path): - (tmp_path / "1_Pooling").mkdir() +@pytest.mark.parametrize( + "module_types", + [ + [ + "sentence_transformers.models.Transformer", + "sentence_transformers.models.Pooling", + "sentence_transformers.models.Dense", + ], + [ + "sentence_transformers.models.Transformer", + "sentence_transformers.models.Normalize", + "sentence_transformers.models.Pooling", + ], + ], +) +def test_sentence_transformer_source_with_unsupported_module_stack_is_rejected(tmp_path, module_types): (tmp_path / "modules.json").write_text( json.dumps( [ - {"idx": 0, "path": "", "type": "sentence_transformers.models.Transformer"}, - {"idx": 1, "path": "1_Pooling", "type": "sentence_transformers.models.Pooling"}, - {"idx": 2, "path": "2_Dense", "type": "sentence_transformers.models.Dense"}, + { + "idx": index, + "path": "" if index == 0 else f"{index}_Module", + "type": module_type, + } + for index, module_type in enumerate(module_types) ] ) ) - (tmp_path / "1_Pooling" / "config.json").write_text(json.dumps({"pooling_mode_mean_tokens": True})) with pytest.raises(ValueError, match="exact supported module stack"): sentence_transformer_export._load_sentence_transformer_wrapper_options(str(tmp_path), {})