Spaces:
Running
Running
| """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) | |