Spaces:
Running
Running
File size: 2,883 Bytes
d289270 b8e2033 d289270 b8e2033 d289270 ba8d0f9 d289270 ba8d0f9 d289270 ba8d0f9 d289270 ba8d0f9 d289270 | 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 | """Pin/download checkpoints and adapt Transformers 5 tokenizer metadata to v4."""
import os
import json
from pathlib import Path
from huggingface_hub import snapshot_download
SPECS = [
("intfloat/multilingual-e5-small",
"fd1525a9fd15316a2d503bf26ab031a61d056e98",
"E5_MODEL_DIR", "/opt/models/multilingual-e5-small"),
("Qwen/Qwen3-Embedding-0.6B",
"97b0c614be4d77ee51c0cef4e5f07c00f9eb65b3",
"QWEN_MODEL_DIR", "/opt/models/qwen3-embedding-0.6b"),
("lightonai/mLateOn",
"edd378f",
"COLBERT_MODEL_DIR", "/opt/models/mlateon"),
("BAAI/bge-m3",
"5617a9f",
"BGE_M3_MODEL_DIR", "/opt/models/bge-m3"),
]
def adapt_mlateon_tokenizer(model_dir: str | Path) -> None:
"""Transformers v5 calls the generic fast tokenizer TokenizersBackend;
Transformers v4 calls it PreTrainedTokenizerFast. All vocabulary and
tokenizer.json bytes stay unchanged.
"""
folder = Path(model_dir)
config_path = folder / "tokenizer_config.json"
config = json.loads(config_path.read_text(encoding="utf-8"))
cls = config.get("tokenizer_class")
if cls != "TokenizersBackend":
raise RuntimeError(f"Unexpected mLateOn tokenizer_class {cls!r}")
if not (folder / "tokenizer.json").is_file():
raise RuntimeError("mLateOn tokenizer.json not found")
config["tokenizer_class"] = "PreTrainedTokenizerFast"
config_path.write_text(
json.dumps(config, indent=2, ensure_ascii=False) + "\n",
encoding="utf-8",
)
for repo_id, revision, env_name, default in SPECS:
# BGE-M3 has NO safetensors encoder: pytorch_model.bin and the
# sparse_linear.pt / colbert_linear.pt heads are all mandatory.
ignored = ["*.onnx", "*.h5", "*.msgpack", "*.ot",
"*.md", "*.tflite", "*.gguf"]
if repo_id != "BAAI/bge-m3":
ignored += ["*.bin", "*.pt"]
path = snapshot_download(
repo_id=repo_id, revision=revision,
local_dir=os.getenv(env_name, default),
ignore_patterns=ignored,
)
if repo_id == "BAAI/bge-m3":
for filename in ("pytorch_model.bin", "sparse_linear.pt",
"colbert_linear.pt"):
if not (Path(path) / filename).is_file():
raise RuntimeError(f"Missing BGE-M3 tri-head weight: {filename}")
print(f"Downloaded {repo_id}@{revision} to {path}", flush=True)
if repo_id == "lightonai/mLateOn":
adapt_mlateon_tokenizer(path)
# Fail the Docker build before release if the real tokenizer fails.
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(path, local_files_only=True)
if not tokenizer("Hello שלום привет", truncation=False)["input_ids"]:
raise RuntimeError("mLateOn tokenizer produced no tokens")
print("mLateOn fast-tokenizer build preflight passed", flush=True)
|