76ab8308eb
text_encoding.py holds the CultCache embedding store (gamecult.vili.text_embedding.v0 through cultcache-py), the encoder identity and key, plan_placement with its typed gpu-busy refusal, and CachedTextEncoder, which sends a request's misses to one child call. kimodo-text-encode.py is that child: it assembles Kimodo's LLM2Vec encoder in bf16 with its own loader (operator ruling encoder-own-loader; peft's offload path raises TypeError in Kimodo's loader), encodes each text alone, writes a MessagePack answer and exits. GPU reserve 3 GiB: four cold runs on Raven's 16 GB card produced identical vectors with no CUDA error. The worker and daemon are not wired yet. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
117 lines
5.1 KiB
Python
117 lines
5.1 KiB
Python
#!/usr/bin/env python3
|
|
"""One-shot Kimodo text encoder: texts on stdin, vectors in --out, then exit so the OS reclaims RAM and VRAM.
|
|
|
|
Reads {"texts": [...]} as JSON on stdin. Loads Kimodo's LLM2Vec preset in bfloat16 with the layers that do not
|
|
fit in free VRAM held in RAM, encodes each text alone as LLM2VecEncoder does, and writes a MessagePack answer.
|
|
There is no CPU fallback; --device cpu exists only for the reference probe.
|
|
"""
|
|
import argparse
|
|
import json
|
|
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 CHILD_EXIT_GPU_BUSY, EMBEDDING_DIM, GpuBusy, plan_placement # noqa: E402
|
|
|
|
from peft import PeftModel # noqa: E402
|
|
from transformers import AutoConfig, AutoTokenizer # noqa: E402
|
|
|
|
from kimodo.model.llm2vec.llm2vec import LLM2Vec, batch_to_device # noqa: E402
|
|
from kimodo.model.load_model import TEXT_ENCODER_PRESETS # noqa: E402
|
|
|
|
|
|
def load_encoder(base_id, peft_id, dtype, peft_kwargs, **model_kwargs):
|
|
"""Assemble the encoder in the order and with the ids of LLM2Vec.from_pretrained.
|
|
|
|
MUST TRACK Kimodo's loader: kimodo/model/llm2vec/llm2vec.py, LLM2Vec.from_pretrained, Kimodo commit 6bb5848
|
|
(nv-tlabs/kimodo). It is mirrored, not called, because that method hands PeftModel.from_pretrained no kwargs, and
|
|
with layers offloaded to RAM peft then takes a path that raises TypeError in accelerate's get_balanced_memory
|
|
(peft 0.18.1, accelerate 1.13.0, transformers 5.1.0). Here every PeftModel.from_pretrained gets peft_kwargs
|
|
(device_map='sequential'). Operator ruling motion-pipeline:ruling:encoder-own-loader. The ids' llm2vec_config.json
|
|
is not a local path, so Kimodo's loader applies no encoder config; neither does this.
|
|
"""
|
|
tokenizer = AutoTokenizer.from_pretrained(base_id)
|
|
tokenizer.pad_token = tokenizer.eos_token
|
|
tokenizer.padding_side = "left"
|
|
model_class = LLM2Vec._get_model_class(AutoConfig.from_pretrained(base_id).__class__.__name__, enable_bidirectional=True)
|
|
model = model_class.from_pretrained(base_id, torch_dtype=dtype, **model_kwargs)
|
|
if hasattr(model, "peft_config"):
|
|
model = PeftModel.from_pretrained(model, base_id, **peft_kwargs).merge_and_unload()
|
|
model = PeftModel.from_pretrained(model, peft_id, **peft_kwargs)
|
|
return LLM2Vec(model=model, tokenizer=tokenizer)
|
|
|
|
|
|
def mem_available() -> int:
|
|
for line in Path("/proc/meminfo").read_text().splitlines():
|
|
if line.startswith("MemAvailable:"):
|
|
return int(line.split()[1]) * 1024
|
|
raise RuntimeError("MemAvailable is not in /proc/meminfo")
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--out", required=True)
|
|
parser.add_argument("--device", choices=["gpu", "cpu"], default="gpu")
|
|
args = parser.parse_args()
|
|
texts = json.load(sys.stdin)["texts"]
|
|
preset = TEXT_ENCODER_PRESETS["llm2vec"]["kwargs"]
|
|
|
|
load_kwargs, peft_kwargs = {}, {}
|
|
if args.device == "gpu":
|
|
free_vram, _ = torch.cuda.mem_get_info(0)
|
|
try:
|
|
max_memory = plan_placement(free_vram, mem_available())
|
|
load_kwargs = {"device_map": "auto", "max_memory": max_memory}
|
|
peft_kwargs = {"device_map": "sequential", "max_memory": max_memory}
|
|
except GpuBusy as busy:
|
|
Path(args.out).write_bytes(msgpack.packb({"error": busy.code, "freeVramBytes": busy.free_vram}, use_bin_type=True))
|
|
return CHILD_EXIT_GPU_BUSY
|
|
|
|
started = time.perf_counter()
|
|
model = load_encoder(
|
|
preset["base_model_name_or_path"],
|
|
preset["peft_model_name_or_path"],
|
|
getattr(torch, preset["dtype"]),
|
|
peft_kwargs,
|
|
**load_kwargs,
|
|
)
|
|
model.eval()
|
|
for p in model.parameters():
|
|
p.requires_grad = False
|
|
load_seconds = time.perf_counter() - started
|
|
|
|
placement = {"gpuBytes": 0, "cpuBytes": 0}
|
|
for p in model.parameters():
|
|
placement["gpuBytes" if p.device.type == "cuda" else "cpuBytes"] += p.numel() * p.element_size()
|
|
input_device = "cuda:0" if args.device == "gpu" else "cpu"
|
|
|
|
vectors, encode_seconds = [], []
|
|
for text in texts:
|
|
began = time.perf_counter()
|
|
# The path LLM2Vec.encode takes for one text, without its .to(device) that would undo the device map.
|
|
features = model.tokenize([model.prepare_for_tokenization(model._convert_to_str("", text))])
|
|
features = batch_to_device(features, input_device)
|
|
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)
|
|
|
|
Path(args.out).write_bytes(msgpack.packb({
|
|
"vectors": vectors,
|
|
"encodeSeconds": encode_seconds,
|
|
"loadSeconds": load_seconds,
|
|
"placement": placement,
|
|
"peakVramBytes": torch.cuda.max_memory_allocated(0) if torch.cuda.is_available() else 0,
|
|
"peakRssBytes": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * 1024,
|
|
"device": args.device,
|
|
}, use_bin_type=True))
|
|
return 0
|
|
|
|
|
|
sys.exit(main())
|