| |
| |
| |
| |
| @@ -151,19 +151,20 @@ def _load_tokenizer(tok_dir: str, cfg: Dict) -> Any: |
| """ |
| from transformers import AutoTokenizer |
|
|
| + remote = bool(cfg.get("trust_remote_code", False)) |
| if not os.path.isdir(tok_dir): |
| - return AutoTokenizer.from_pretrained(cfg.get("encoder")) |
| + return AutoTokenizer.from_pretrained(cfg.get("encoder"), trust_remote_code=remote) |
|
|
| try: |
| stamp = (os.path.abspath(tok_dir), os.path.getmtime(os.path.join(tok_dir, "tokenizer_config.json"))) |
| except OSError: |
| - return AutoTokenizer.from_pretrained(tok_dir) |
| + return AutoTokenizer.from_pretrained(tok_dir, trust_remote_code=remote) |
|
|
| with _TOKENIZERS_LOCK: |
| cached = _TOKENIZERS.get(stamp) |
| if cached is not None: |
| return cached |
| - tokenizer = AutoTokenizer.from_pretrained(tok_dir) |
| + tokenizer = AutoTokenizer.from_pretrained(tok_dir, trust_remote_code=remote) |
| _TOKENIZERS[stamp] = tokenizer |
| return tokenizer |
|
|
| |
| |
| |
| |
| @@ -201,12 +201,26 @@ def _apply_rope_config(ecfg) -> None: |
| def build_model(cfg: Dict, encoder_dir: Optional[str] = None, pretrained: bool = True) -> DecisionModel: |
| from transformers import AutoConfig, AutoModel |
|
|
| + # Encoders whose modeling code lives on the Hub (e.g. NeoBERT) must opt in via the config. |
| + remote = bool(cfg.get("trust_remote_code", False)) |
| if not pretrained or (encoder_dir and os.path.exists(encoder_dir)): |
| - ecfg = AutoConfig.from_pretrained(encoder_dir or cfg["encoder"]) |
| + ecfg = AutoConfig.from_pretrained(encoder_dir or cfg["encoder"], trust_remote_code=remote) |
| _apply_rope_config(ecfg) |
| - enc = AutoModel.from_config(ecfg, attn_implementation="sdpa") |
| + enc = AutoModel.from_config(ecfg, attn_implementation="sdpa", trust_remote_code=remote) |
| else: |
| - enc = AutoModel.from_pretrained(cfg["encoder"], attn_implementation="sdpa") |
| + enc = AutoModel.from_pretrained(cfg["encoder"], attn_implementation="sdpa", trust_remote_code=remote) |
| + # NeoBERT computes its RoPE tables as non-persistent buffers in __init__; transformers 5 builds models on |
| + # the meta device, so they come back as uninitialised memory and every forward pass returns NaN. |
| + if hasattr(enc, "freqs_cos"): |
| + import sys |
| + |
| + precompute = sys.modules[type(enc).__module__].precompute_freqs |
| + c = enc.config |
| + enc.freqs_cos, enc.freqs_sin = precompute(c.hidden_size // c.num_attention_heads, c.max_length) |
| + # Some encoder configs pin a half-precision dtype (NeoBERT ships bfloat16), which transformers 5 honours. Keep |
| + # master weights in float32: bf16 weights round most optimizer updates to zero, and CPU inference in bf16 is ~8x slower. |
| + # Mixed precision is applied by autocast at train/inference time instead. |
| + enc = enc.float() |
| return DecisionModel(enc, cfg.get("head_layers", 2), len(cfg.get("act_costs", {})) + 1) |
|
|
|
|
|
|
