File size: 5,073 Bytes
d69f154 | 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 | """Local release loader and optional text -> BGE-M3 -> Scripture Genome CLI.
Read the accompanying custom model code before allowing trust_remote_code.
INT8 is weight-only storage, dequantized to FP32 here, not INT8 execution.
"""
import argparse
import hashlib
import json
from pathlib import Path
import torch
from safetensors.torch import load_file
from transformers import AutoConfig, AutoModel, AutoTokenizer
def load_head(model_dir, precision="fp32", device="cpu"):
model_dir = Path(model_dir)
if precision not in {"fp32", "fp16", "int8"}:
raise ValueError("precision must be fp32, fp16, or int8")
config = AutoConfig.from_pretrained(model_dir, trust_remote_code=True, local_files_only=True)
gene_bytes = (model_dir / "gene_ids.txt").read_bytes()
if hashlib.sha256(gene_bytes).hexdigest() != config.gene_ids_sha256:
raise ValueError("gene_ids.txt does not match the model configuration")
genes = gene_bytes.decode("utf-8").splitlines()
if len(genes) != config.output_dim or len(set(genes)) != len(genes):
raise ValueError("Invalid gene count/order file")
if precision == "int8":
arrays = load_file(str(model_dir / "head_int8.safetensors"))
model = AutoModel.from_config(config, trust_remote_code=True)
expected = {f"{prefix}{i}" for i in (0, 1) for prefix in ("w", "b", "scale")}
if set(arrays) != expected:
raise ValueError("Unexpected INT8 tensor keys")
state = {}
for i, layer in ((0, 0), (1, 2)):
w, scale, bias = arrays[f"w{i}"], arrays[f"scale{i}"], arrays[f"b{i}"]
if w.dtype != torch.int8 or scale.shape != (w.shape[0],):
raise ValueError("Invalid INT8 weights or scales")
if not torch.isfinite(scale).all() or not (scale > 0).all() or not torch.isfinite(bias).all():
raise ValueError("Nonfinite weights or invalid quantization scales")
state[f"head.{layer}.weight"] = w.float() * scale.float()[:, None]
state[f"head.{layer}.bias"] = bias.float()
model.load_state_dict(state, strict=True)
else:
# FP16 storage is loaded into FP32 arithmetic here for CPU portability.
# To run FP16 on a compatible device, explicitly cast the head and input.
options = {"variant": "fp16"} if precision == "fp16" else {}
model = AutoModel.from_pretrained(
model_dir, trust_remote_code=True, local_files_only=True,
use_safetensors=True, dtype=torch.float32, **options,
)
return model.to(device).eval(), genes
def encode_texts(texts, head, tokenizer, encoder):
"""No retrieval: returns dense embeddings and 568-feature predictions."""
if not texts or any(not isinstance(t, str) or not t.strip() for t in texts):
raise ValueError("Provide a nonempty list of nonempty strings")
encoded = tokenizer(texts, padding=True, truncation=False, return_tensors="pt")
if encoded["input_ids"].shape[1] > head.config.max_length:
raise ValueError(f"Input exceeds {head.config.max_length} tokens; shorten it explicitly")
encoder_device = next(encoder.parameters()).device
head_parameter = next(head.parameters())
encoded = {name: tensor.to(encoder_device) for name, tensor in encoded.items()}
head.eval()
encoder.eval()
with torch.inference_mode():
hidden = encoder(**encoded).last_hidden_state
dense = torch.nn.functional.normalize(hidden[:, 0, :].float(), p=2, dim=-1)
genome = head(embeddings=dense.to(device=head_parameter.device, dtype=head_parameter.dtype)).genome
return {"dense": dense.cpu(), "genome": genome.float().cpu()}
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--model-dir", type=Path, default=Path(__file__).resolve().parent)
parser.add_argument("--precision", choices=("fp32", "fp16", "int8"), default="fp32")
parser.add_argument("--device", default="cpu")
parser.add_argument("--base-dir", type=Path, help="Optional locally extracted pinned BGE-M3 folder")
parser.add_argument("--text", action="append", required=True)
args = parser.parse_args()
head, genes = load_head(args.model_dir, args.precision, args.device)
if args.base_dir:
base, options = str(args.base_dir), {"local_files_only": True}
else:
base = head.config.encoder_name
options = {"revision": head.config.encoder_revision}
tokenizer = AutoTokenizer.from_pretrained(base, **options)
encoder = AutoModel.from_pretrained(base, **options).to(args.device).eval()
result = encode_texts(args.text, head, tokenizer, encoder)
print(json.dumps({
"texts": args.text, "dense_shape": list(result["dense"].shape),
"genome_shape": list(result["genome"].shape),
"checkpoint_sha256": head.config.checkpoint_sha256,
"precision": args.precision, "gene_ids": genes,
"genomes": result["genome"].tolist(),
}, ensure_ascii=False, allow_nan=False))
if __name__ == "__main__":
main()
|