c3684b733b
The encoder child is CPU only (GPU placement, GpuBusy and plan_placement are deleted) and assembles LLM2Vec B: Llama base, MNTP merged, then the supervised adapter. The cache identity names assembly, device, dtype, torch/transformers/peft/accelerate versions, Kimodo commit and the child's own sha256 beside the three revisions. VILI_THREADS (1..16) caps torch in the worker and the child. The worker binds port 0 and reports it in its ready line; errors name a class and place, never the request. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
79 lines
3.2 KiB
Python
79 lines
3.2 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 LLM2Vec assembly B on CPU in bfloat16, encodes each text alone as
|
|
LLM2VecEncoder does, 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 # 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(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.
|
|
"""
|
|
tokenizer = AutoTokenizer.from_pretrained(LLAMA_BASE)
|
|
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)
|
|
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(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([model.prepare_for_tokenization(model._convert_to_str("", 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())
|