File size: 5,127 Bytes
534cc76 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | """Load pinned BioLORD weights and cache normalized text embeddings."""
import argparse
import json
import os
from pathlib import Path
import torch
from .provenance import TEXT_ID, TEXT_REVISION, TEXT_IDENTITY, VERSION, sha256
class TextEncoder:
"""Load frozen BioLORD on demand and cache normalized sentence vectors on CPU."""
def __init__(self, cache=None, device="cpu", model_path=None):
self.device, self.model_path, self.model = device, model_path, None
self.embeddings = {}
if cache is not None and Path(cache).exists():
payload = torch.load(cache, map_location="cpu", weights_only=True)
if payload.get("identity") != TEXT_IDENTITY:
raise ValueError("Text bank does not match pinned BioLORD processing")
for text, vector in payload["embeddings"].items():
if vector.shape != (768,) or not torch.isfinite(vector).all():
raise ValueError("Invalid text bank vector")
self.embeddings[text] = vector.float()
def encode(self, strings):
"""Return separate [N,768] CPU vectors; reject texts above the token limit."""
if not isinstance(strings, list) or not strings or any(not isinstance(s, str) or not s.strip() for s in strings):
raise ValueError("Encode a nonempty list of separate, nonempty text strings")
missing = list(dict.fromkeys(s for s in strings if s not in self.embeddings))
if missing:
if self.model is None:
from sentence_transformers import SentenceTransformer
local = self.model_path or (
Path("third-party/biolord")
if Path("third-party/biolord").exists()
else None
)
if local is not None:
local = Path(local)
source = json.loads((local / "source.json").read_text())
if (
source.get("model_id") != TEXT_ID
or source.get("revision") != TEXT_REVISION
):
raise ValueError(
"Local text weights do not match pinned BioLORD"
)
for name, expected in source["files"].items():
path = (local / name.replace("\\", "/")).resolve()
if (
not path.is_relative_to(local.resolve())
or sha256(path) != expected
):
raise ValueError(
"Local text weight/config checksum changed"
)
self.model_path = local
self.model = SentenceTransformer(
str(local) if local else TEXT_ID,
revision=None if local else TEXT_REVISION,
device=self.device,
trust_remote_code=False,
cache_folder=os.environ.get("HF_HOME", "cache/huggingface"),
)
self.model.eval().requires_grad_(False)
for text in missing:
if (
len(self.model.tokenizer.encode(text, add_special_tokens=True))
> self.model.max_seq_length
):
raise ValueError(
"Text exceeds BioLORD context; do not silently truncate decisive wording"
)
with torch.inference_mode():
vectors = self.model.encode(
missing,
normalize_embeddings=True,
convert_to_tensor=True,
show_progress_bar=False,
batch_size=32,
)
for text, vector in zip(missing, vectors):
if vector.shape != (768,) or not torch.isfinite(vector).all():
raise ValueError("Text encoder produced an invalid vector")
self.embeddings[text] = vector.detach().cpu().float().clone()
return torch.stack([self.embeddings[s] for s in strings])
def save(self, path):
"""Save cached vectors with their encoder identity."""
Path(path).parent.mkdir(parents=True, exist_ok=True)
torch.save({"identity": TEXT_IDENTITY, "embeddings": self.embeddings}, path)
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--texts", type=Path, required=True)
parser.add_argument("--output", type=Path, default=Path(f"cache/{VERSION}/text.pt"))
parser.add_argument("--device", default="cpu")
args = parser.parse_args()
encoder = TextEncoder(args.output, args.device)
encoder.encode(json.loads(args.texts.read_text()))
encoder.save(args.output)
print(json.dumps({"texts": len(encoder.embeddings), "identity": TEXT_IDENTITY}))
if __name__ == "__main__":
main()
|