From 6651662b2cd915b17cd64b5de5a731aecb01f99a Mon Sep 17 00:00:00 2001 From: Meta Vens Date: Fri, 23 Oct 2026 17:40:08 +0200 Subject: [PATCH] Encoder: pool after LLM2Vec's query marker, live_identity reads its own versions, tests reach every field The header-on vector now matches Kimodo's encoder loaded from the repo ids (cosine 1.0; the bare header gave 0.9974) because the wrapped text carries LLM2Vec's instruction/query separator, which keeps the header out of the pooling. live_identity reads the installed versions itself, so the worker passes no version it could fake, and the tests pin installed_versions against substituted modules, the prompt form and device in the identity, and a text that differs past any prefix. Co-Authored-By: Claude Opus 5.5 --- scripts/kimodo-resident-worker.py | 4 +-- scripts/test_text_encoding.py | 50 +++++++++++++++---------------- scripts/text_encoding.py | 14 +++++---- 3 files changed, 34 insertions(+), 34 deletions(-) diff --git a/scripts/kimodo-resident-worker.py b/scripts/kimodo-resident-worker.py index eff7069..df9feb2 100644 --- a/scripts/kimodo-resident-worker.py +++ b/scripts/kimodo-resident-worker.py @@ -47,9 +47,7 @@ model_load_started_at = time.time() preset = TEXT_ENCODER_PRESETS["llm2vec"]["kwargs"] kimodo_commit = os.environ["VILI_KIMODO_COMMIT"] python_env = os.environ["VILI_PYTHON_ENV"] -encoder_identity = text_encoding.live_identity( - HF_HOME, preset, kimodo_commit, text_encoding.installed_versions(), -) +encoder_identity = text_encoding.live_identity(HF_HOME, preset, kimodo_commit) # The 8B encoder never lives here: misses go to a one-shot child (kimodo-text-encode.py), hits come from the store. text_encoder = text_encoding.CachedTextEncoder( text_encoding.open_store(os.path.join(os.environ["VILI_EMBEDDINGS_DIR"], "text-embeddings.cc")), diff --git a/scripts/test_text_encoding.py b/scripts/test_text_encoding.py index 2b547d6..49d6fa4 100644 --- a/scripts/test_text_encoding.py +++ b/scripts/test_text_encoding.py @@ -5,7 +5,9 @@ import subprocess import sys import tempfile import textwrap +import types import unittest +from unittest import mock from pathlib import Path sys.path.insert(0, str(Path(__file__).resolve().parent)) @@ -51,10 +53,13 @@ class Identity(unittest.TestCase): self.assertEqual(len(key), 64) int(key, 16) self.assertNotEqual(key, te.embedding_key("id", "text b")) + self.assertNotEqual(te.embedding_key("id", "a" * 300 + "x"), te.embedding_key("id", "a" * 300 + "y")) - def test_chat_header_is_the_llama3_user_turn(self): - self.assertEqual(te.chat_header(" hello "), "<|start_header_id|>user<|end_header_id|>" + chr(10) * 2 + "hello<|eot_id|>") - form = hashlib.sha256(("<|start_header_id|>user<|end_header_id|>" + chr(10) * 2 + "|<|eot_id|>").encode("utf-8")).hexdigest() + def test_chat_header_is_the_llama3_user_turn_pooled_after_the_separator(self): + header = "<|start_header_id|>user<|end_header_id|>" + chr(10) * 2 + self.assertEqual(te.chat_header("hello "), header + "!@#$%^&*()hello<|eot_id|>") + self.assertEqual(te.chat_header(" hello"), header + "!@#$%^&*() hello<|eot_id|>") + form = hashlib.sha256(("|".join((header, "!@#$%^&*()", "<|eot_id|>"))).encode("utf-8")).hexdigest() self.assertEqual(te.PROMPT_FORM, "llama3-chat-header:" + form) def test_every_identity_field_changes_the_key(self): @@ -76,11 +81,12 @@ class Identity(unittest.TestCase): for repo in repos: write_revision(repo, revision[repo]) - real_child = te.CHILD_PATH + real_child, real_versions = te.CHILD_PATH, te.installed_versions te.CHILD_PATH = loader try: def identity(preset=preset, commit="6bb5848", versions=versions): - return te.live_identity(home, preset, commit, versions) + te.installed_versions = lambda: versions + return te.live_identity(home, preset, commit) base = identity() self.assertGreater(base.index("assembly="), 200) @@ -98,12 +104,12 @@ class Identity(unittest.TestCase): loader.write_bytes(b"one") self.assertEqual(identity(), base) finally: - te.CHILD_PATH = real_child - self.assertEqual(len(changed), 9) + te.CHILD_PATH, te.installed_versions = real_child, real_versions + self.assertEqual(len(changed), 10) for field, other in changed.items(): self.assertNotEqual(te.embedding_key(base, "x"), te.embedding_key(other, "x"), field) - for field in ("assembly", "prompt", "device"): - self.assertIn(f"{field}=", base) + for part in (f"assembly={te.ASSEMBLY}", f"prompt={te.PROMPT_FORM}", "device=cpu", "dtype=bfloat16", "torch=2.5.1+cpu"): + self.assertIn(part, base) def test_assembly_prompt_and_device_change_the_key(self): fields = dict( @@ -112,23 +118,15 @@ class Identity(unittest.TestCase): kimodo_commit="6bb5848", loader_sha256="a" * 64, ) key = te.embedding_key(te.encoder_identity(**fields), "x") - for name, other in (("assembly", "llm2vec-A"), ("prompt_form", "none"), ("device", "cuda")): + for name, other in (("assembly", "llm2vec-A"), ("prompt_form", "none"), ("device", "other-device")): self.assertNotEqual(key, te.embedding_key(te.encoder_identity(**{**fields, name: other}), "x"), name) - def test_installed_versions_are_the_libraries_in_use(self): - import accelerate - import peft - import transformers - self.assertEqual(te.installed_versions(), { - "torch": te.torch.__version__, "transformers": transformers.__version__, "peft": peft.__version__, "accelerate": accelerate.__version__, - }) - - def test_snapshot_revision_is_read_from_the_mounted_cache(self): - with tempfile.TemporaryDirectory() as home: - refs = Path(home) / "hub" / "models--org--name" / "refs" - refs.mkdir(parents=True) - (refs / "main").write_text("deadbeef\n") - self.assertEqual(te.read_snapshot_revision(Path(home), "org/name"), "deadbeef") + def test_installed_versions_are_read_from_the_libraries_in_use(self): + fakes = {name: types.SimpleNamespace(__version__=f"{name}-v") for name in ("transformers", "peft", "accelerate")} + with mock.patch.dict(sys.modules, fakes), mock.patch.object(te.torch, "__version__", "torch-v"): + self.assertEqual(te.installed_versions(), { + "torch": "torch-v", "transformers": "transformers-v", "peft": "peft-v", "accelerate": "accelerate-v", + }) class SnapshotPath(unittest.TestCase): @@ -278,7 +276,7 @@ import sys, numpy as np from kimodo.model.load_model import TEXT_ENCODER_PRESETS from kimodo.model.llm2vec.llm2vec_wrapper import LLM2VecEncoder kwargs = TEXT_ENCODER_PRESETS["llm2vec"]["kwargs"] -encoder = LLM2VecEncoder(**kwargs, device="cpu") +encoder = LLM2VecEncoder(**{**kwargs, "device": "cpu"}) vector, _ = encoder(sys.argv[2]) np.save(sys.argv[1], vector.float().numpy().reshape(-1)) """ @@ -296,7 +294,7 @@ class ReferenceVector(unittest.TestCase): subprocess.run([sys.executable, "-c", REFERENCE_CODE, str(reference), ROCK], env=env, check=True) os.environ.update(env) ours = np.frombuffer(te.run_child([ROCK], Path(directory) / "answer.mp")["vectors"][0], dtype=np.float32) - expected = np.load(reference) + expected = np.load(reference) cosine = float(ours @ expected / (np.linalg.norm(ours) * np.linalg.norm(expected))) print("cosine to Kimodo from repo ids:", cosine) self.assertGreaterEqual(cosine, 0.9999) diff --git a/scripts/text_encoding.py b/scripts/text_encoding.py index e4ce0ac..4af2d4e 100644 --- a/scripts/text_encoding.py +++ b/scripts/text_encoding.py @@ -28,10 +28,13 @@ CHILD_PATH = Path(__file__).with_name("kimodo-text-encode.py") LLAMA_BASE = "meta-llama/Meta-Llama-3-8B-Instruct" # the base model both LLM2Vec adapters are applied to ASSEMBLY = "llm2vec-mntp-merged-then-supervised" # Llama-3 Instruct's chat header, applied by Vili's own code (motion-pipeline:ruling:encoder-chat-header-on): the -# text goes in as a user turn, as LLM2Vec does for this model when it is loaded from the repo id. +# text goes in as a user turn, as LLM2Vec does for this model when it is loaded from the repo id. The separator is +# LLM2Vec's instruction/query marker: the model attends over the whole turn, and pooling covers what follows the +# marker (the text and the end-of-turn token), not the header, exactly as Kimodo's encoder pools. CHAT_HEADER_PREFIX = "<|start_header_id|>user<|end_header_id|>\n\n" CHAT_HEADER_SUFFIX = "<|eot_id|>" -PROMPT_FORM = "llama3-chat-header:" + hashlib.sha256((CHAT_HEADER_PREFIX + "|" + CHAT_HEADER_SUFFIX).encode("utf-8")).hexdigest() +QUERY_SEPARATOR = "!@#$%^&*()" +PROMPT_FORM = "llama3-chat-header:" + hashlib.sha256("|".join((CHAT_HEADER_PREFIX, QUERY_SEPARATOR, CHAT_HEADER_SUFFIX)).encode("utf-8")).hexdigest() MAX_THREADS = 16 @@ -61,8 +64,8 @@ def snapshot_path(hf_home: Path, repo_id: str) -> Path: def chat_header(text: str) -> str: - """The text as the encoder reads it: a user turn in Llama-3 Instruct's chat format.""" - return CHAT_HEADER_PREFIX + text.strip() + CHAT_HEADER_SUFFIX + """The text as the encoder reads it: a user turn in Llama-3 Instruct's chat format, pooled after the separator.""" + return CHAT_HEADER_PREFIX + (QUERY_SEPARATOR + text).strip() + CHAT_HEADER_SUFFIX def encoder_identity(*, assembly: str, prompt_form: str, device: str, revisions: dict, dtype: str, torch_version: str, @@ -87,9 +90,10 @@ def installed_versions() -> dict: return {"torch": torch.__version__, "transformers": transformers.__version__, "peft": peft.__version__, "accelerate": accelerate.__version__} -def live_identity(hf_home, preset: dict, kimodo_commit: str, versions: dict) -> str: +def live_identity(hf_home, preset: dict, kimodo_commit: str) -> str: """The identity the worker keys vectors under: the three snapshot revisions the mounted cache names, the preset's dtype, the library versions, the Kimodo commit and the sha256 of the encoder child's own bytes.""" + versions = installed_versions() repos = (preset["base_model_name_or_path"], preset["peft_model_name_or_path"], LLAMA_BASE) return encoder_identity( assembly=ASSEMBLY, prompt_form=PROMPT_FORM, device="cpu",