| """Build models/nanbeige4.2-3b-uno: Uno's SDAR modeling adapted to Nanbeige4.2 (no qk-norm, 2 loops).""" |
| import json, os, re, shutil, sys |
| from pathlib import Path |
| from huggingface_hub import snapshot_download |
|
|
| out = Path(sys.argv[1]) |
| nb = Path(snapshot_download("Nanbeige/Nanbeige4.2-3B", revision="b82e54bd609793562a75cbf9337970a93369eab5")) |
| sdar = Path(snapshot_download("s-sahoo/uno-qwen3-8B", revision="b8a7577b3223bdcf2b3af0f2fc6e95258b3bbc29")) |
| out.mkdir(parents=True, exist_ok=True) |
|
|
| for f in nb.iterdir(): |
| if f.suffix in (".safetensors", ".json", ".model", ".txt", ".jinja") and f.name != "config.json": |
| dst = out / f.name |
| if not dst.exists(): |
| dst.symlink_to(f.resolve()) |
| for name in ("configuration_sdar.py", "fused_linear_diffusion_cross_entropy.py"): |
| shutil.copy(sdar / name, out / name) |
|
|
| src = (sdar / "modeling_sdar.py").read_text() |
| |
| for nm in ("q_norm", "k_norm"): |
| old = f" self.{nm} = SDARRMSNorm(self.head_dim, eps=config.rms_norm_eps)\n" |
| assert old in src, nm |
| src = src.replace(old, f" self.{nm} = (SDARRMSNorm(self.head_dim, eps=config.rms_norm_eps)\n" |
| f" if getattr(config, \"use_qk_norm\", True) else nn.Identity())\n") |
| |
| head = " for decoder_layer in self.layers[: self.config.num_hidden_layers]:\n" |
| tail = " hidden_states = self.norm(hidden_states)\n\n # add hidden states from the last decoder layer\n" |
| i, j = src.index(head), src.index(tail) |
| body = "".join(" " + ln if ln.strip() else ln for ln in src[i:j].splitlines(True)) |
| loop = (' num_loops = getattr(self.config, "num_loops", 1)\n' |
| ' if num_loops > 1 and use_cache:\n' |
| ' raise NotImplementedError("KV caching across loops is not supported (training-only).")\n' |
| ' for loop_idx in range(num_loops):\n' + body + |
| ' if loop_idx < num_loops - 1:\n' |
| ' hidden_states = self.norm(hidden_states)\n\n') |
| src = src[:i] + loop + src[j:] |
| (out / "modeling_sdar.py").write_text(src) |
|
|
| ref = json.loads((nb / "config.json").read_text()) |
| cfg = json.loads((sdar / "config.json").read_text()) |
| for k in ("vocab_size", "hidden_size", "intermediate_size", "num_hidden_layers", "num_attention_heads", |
| "num_key_value_heads", "head_dim", "rms_norm_eps", "rope_theta", "max_position_embeddings", |
| "bos_token_id", "eos_token_id", "pad_token_id", "tie_word_embeddings", "num_loops"): |
| cfg[k] = ref[k] |
| cfg.update(use_qk_norm=False, max_window_layers=ref["num_hidden_layers"], mask_token_id=int(sys.argv[2]), |
| block_size=8, use_cache=False, noise="uniform") |
| cfg.pop("rope_scaling", None) |
| (out / "config.json").write_text(json.dumps(cfg, indent=2)) |
| print("built", out) |
|
|
