4fe1fc963c
The child wraps every text with text_encoding.chat_header and tokenizes it itself, so the prompt form no longer rides on the model's _name_or_path. The identity names the assembly truthfully (header form included, as a hash of the template) and re-keys the cache. The worker builds its identity through live_identity and installed_versions, which the tests reach; the identity test now uses live-length revisions, so a truncated identity or key fails. An env-gated test compares a fresh child vector with Kimodo's own encoder loaded from the repo ids. Supersedes encoder-assembly-b (motion-pipeline:ruling:encoder-chat-header-on). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
78 lines
3.3 KiB
Python
78 lines
3.3 KiB
Python
#!/usr/bin/env python3
|
|
"""One-shot Kimodo text encoder: texts on stdin, vectors in --out, then exit so the OS reclaims its RAM.
|
|
|
|
Reads {"texts": [...]} as JSON on stdin. Loads the LLM2Vec encoder on CPU in bfloat16, wraps each text in the Llama-3 chat header, encodes it alone, and writes a MessagePack answer. The torch thread budget is VILI_THREADS.
|
|
"""
|
|
import json
|
|
import os
|
|
import resource
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
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, chat_header, snapshot_path # noqa: E402
|
|
|
|
from peft import PeftModel # noqa: E402
|
|
from transformers import AutoConfig, AutoTokenizer # noqa: E402
|
|
|
|
from kimodo.model.llm2vec.llm2vec import LLM2Vec # noqa: E402
|
|
from kimodo.model.load_model import TEXT_ENCODER_PRESETS # noqa: E402
|
|
|
|
|
|
def load_encoder(hf_home, mntp_id, supervised_id, dtype):
|
|
"""The Llama base with the MNTP adapter merged in, then the supervised adapter on top (motion-pipeline:ruling:encoder-own-loader).
|
|
|
|
Everything loads from the local snapshot directories, so the model's _name_or_path is a path and LLM2Vec never
|
|
adds the chat header itself; main() wraps every text with chat_header, so the prompt form is ours and no
|
|
side effect of the load decides it. The identity carries the Kimodo commit and this file's sha256."""
|
|
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).__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)
|
|
|
|
|
|
def main() -> int:
|
|
apply_thread_cap(os.environ)
|
|
texts = json.load(sys.stdin)["texts"]
|
|
out = Path(sys.argv[sys.argv.index("--out") + 1])
|
|
preset = TEXT_ENCODER_PRESETS["llm2vec"]["kwargs"]
|
|
|
|
started = time.perf_counter()
|
|
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
|
|
load_seconds = time.perf_counter() - started
|
|
|
|
vectors, encode_seconds = [], []
|
|
for text in texts:
|
|
began = time.perf_counter()
|
|
features = model.tokenize([chat_header(text)])
|
|
with torch.no_grad():
|
|
vector = model.forward(features).detach().cpu().to(torch.float32)[0]
|
|
assert vector.shape == (EMBEDDING_DIM,)
|
|
vectors.append(vector.numpy().tobytes())
|
|
encode_seconds.append(time.perf_counter() - began)
|
|
|
|
out.write_bytes(msgpack.packb({
|
|
"vectors": vectors,
|
|
"encodeSeconds": encode_seconds,
|
|
"loadSeconds": load_seconds,
|
|
"placement": {"cpuBytes": sum(p.numel() * p.element_size() for p in model.parameters())},
|
|
"peakRssBytes": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * 1024,
|
|
"device": "cpu",
|
|
"assembly": ASSEMBLY,
|
|
}, use_bin_type=True))
|
|
return 0
|
|
|
|
|
|
sys.exit(main())
|