diff --git a/scripts/kimodo-text-encode.py b/scripts/kimodo-text-encode.py index 6bd9e2a..f86ed00 100644 --- a/scripts/kimodo-text-encode.py +++ b/scripts/kimodo-text-encode.py @@ -15,7 +15,7 @@ import msgpack import torch sys.path.insert(0, str(Path(__file__).resolve().parent)) -from text_encoding import ASSEMBLY, EMBEDDING_DIM, LLAMA_BASE, apply_thread_cap # noqa: E402 +from text_encoding import ASSEMBLY, EMBEDDING_DIM, LLAMA_BASE, apply_thread_cap, snapshot_path # noqa: E402 from peft import PeftModel # noqa: E402 from transformers import AutoConfig, AutoTokenizer # noqa: E402 @@ -24,20 +24,23 @@ from kimodo.model.llm2vec.llm2vec import LLM2Vec # noqa: E402 from kimodo.model.load_model import TEXT_ENCODER_PRESETS # noqa: E402 -def load_encoder(mntp_id, supervised_id, dtype): +def load_encoder(hf_home, mntp_id, supervised_id, dtype): """Assembly B: the Llama base with the MNTP adapter merged in, then the supervised adapter on top. - Kimodo's own LLM2VecEncoder loads the MNTP repo as the base and never applies the MNTP weights under the - pinned transformers (assembly A); this assembly is ours (motion-pipeline:ruling:encoder-assembly-b). The - identity carries the Kimodo commit and this file's sha256, so a changed loader misses the cache. + Kimodo's own LLM2VecEncoder loads the MNTP repo as the base (assembly A); this assembly is ours + (motion-pipeline:ruling:encoder-assembly-b). Everything loads from the local snapshot directories, as the + research's reference vector did: the model's _name_or_path is then a path, not the repo id, so + LLM2Vec.prepare_for_tokenization does not wrap the text in Llama's chat header, which is what separates B's + vector from A's. The identity carries the Kimodo commit and this file's sha256, so a changed loader misses. """ - tokenizer = AutoTokenizer.from_pretrained(LLAMA_BASE) + llama, mntp, supervised = (snapshot_path(hf_home, repo) for repo in (LLAMA_BASE, mntp_id, supervised_id)) + tokenizer = AutoTokenizer.from_pretrained(llama) tokenizer.pad_token = tokenizer.eos_token tokenizer.padding_side = "left" - model_class = LLM2Vec._get_model_class(AutoConfig.from_pretrained(LLAMA_BASE).__class__.__name__, enable_bidirectional=True) - model = model_class.from_pretrained(LLAMA_BASE, torch_dtype=dtype) - model = PeftModel.from_pretrained(model, mntp_id).merge_and_unload() - model = PeftModel.from_pretrained(model, supervised_id) + model_class = LLM2Vec._get_model_class(AutoConfig.from_pretrained(llama).__class__.__name__, enable_bidirectional=True) + model = model_class.from_pretrained(llama, torch_dtype=dtype) + model = PeftModel.from_pretrained(model, mntp).merge_and_unload() + model = PeftModel.from_pretrained(model, supervised) return LLM2Vec(model=model, tokenizer=tokenizer) @@ -48,7 +51,7 @@ def main() -> int: preset = TEXT_ENCODER_PRESETS["llm2vec"]["kwargs"] started = time.perf_counter() - model = load_encoder(preset["base_model_name_or_path"], preset["peft_model_name_or_path"], getattr(torch, preset["dtype"])) + model = load_encoder(os.environ["HF_HOME"], preset["base_model_name_or_path"], preset["peft_model_name_or_path"], getattr(torch, preset["dtype"])) model.eval() for p in model.parameters(): p.requires_grad = False diff --git a/scripts/test_text_encoding.py b/scripts/test_text_encoding.py index 899ffe6..ddaad14 100644 --- a/scripts/test_text_encoding.py +++ b/scripts/test_text_encoding.py @@ -72,6 +72,17 @@ class Identity(unittest.TestCase): self.assertEqual(te.read_snapshot_revision(Path(home), "org/name"), "deadbeef") +class SnapshotPath(unittest.TestCase): + def test_resolves_the_revision_the_mounted_cache_names(self): + with tempfile.TemporaryDirectory() as home: + repo = Path(home) / "hub" / "models--org--name" + (repo / "refs").mkdir(parents=True) + (repo / "refs" / "main").write_text("deadbeef +") + (repo / "snapshots" / "deadbeef").mkdir(parents=True) + self.assertEqual(te.snapshot_path(Path(home), "org/name"), repo / "snapshots" / "deadbeef") + + class CachedEncoder(unittest.TestCase): def test_one_child_call_for_a_requests_misses_and_none_for_hits(self): with tempfile.TemporaryDirectory() as directory: @@ -167,7 +178,7 @@ class ThreadCap(unittest.TestCase): def test_the_child_applies_the_cap_before_it_loads_anything(self): source = (Path(te.__file__).with_name("kimodo-text-encode.py")).read_text() - self.assertLess(source.index("apply_thread_cap(os.environ)"), source.index("load_encoder(preset")) + self.assertLess(source.index("apply_thread_cap(os.environ)"), source.index("load_encoder(os.environ")) class ChildBoundary(unittest.TestCase): diff --git a/scripts/text_encoding.py b/scripts/text_encoding.py index 26676cd..26c947a 100644 --- a/scripts/text_encoding.py +++ b/scripts/text_encoding.py @@ -50,6 +50,11 @@ def read_snapshot_revision(hf_home: Path, repo_id: str) -> str: return ref.read_text().strip() +def snapshot_path(hf_home: Path, repo_id: str) -> Path: + """The directory of the snapshot read_snapshot_revision names: assembly B loads from these, never from repo ids.""" + return Path(hf_home) / "hub" / ("models--" + repo_id.replace("/", "--")) / "snapshots" / read_snapshot_revision(hf_home, repo_id) + + def encoder_identity(*, assembly: str, device: str, revisions: dict, dtype: str, torch_version: str, transformers_version: str, peft_version: str, accelerate_version: str, kimodo_commit: str, loader_sha256: str) -> str: """The canonical string a vector is keyed under: every input that can change it, in a fixed field order.