Spaces:
Sleeping
Sleeping
| """Tiny random model with logit-lens heads, for local UI development on CPU: the real dssd | |
| engine runs end to end (exact verification, copies, head, star), on gibberish text.""" | |
| import torch | |
| from transformers import Qwen3Config, Qwen3ForCausalLM | |
| from dssd import DSSD | |
| from dssd.heads import ExitHead | |
| LAYERS = [2, 4, 6] | |
| _CHARS = [chr(c) for c in range(32, 127)] + ["\n", "\t"] # 97 ids | |
| _TREE = {"copy": {"m": 11, "n_max": 4, "n_min": 1}, "head_budget": 4, | |
| "star": {"head": 0, "C": 4, "tau": [0.0, 0.0, 0.0, 0.0], "gamma": 0.0}} | |
| _POLICY = {"type": "threshold", "thresholds": {"0": 0.0, "2": 0.0}, "cum_gamma": {"0": 0.0}, "max_draft": 15, | |
| "tree": _TREE} | |
| class CharTokenizer: | |
| eos_token_id = None | |
| def encode(self, text, add_special_tokens=False): | |
| return [_CHARS.index(c) if c in _CHARS else _CHARS.index("?") for c in text] | |
| def decode(self, ids, skip_special_tokens=True): | |
| return "".join(_CHARS[i % len(_CHARS)] for i in ids) | |
| def tiny_engine() -> DSSD: | |
| torch.manual_seed(0) | |
| cfg = Qwen3Config(vocab_size=len(_CHARS), hidden_size=64, intermediate_size=128, num_hidden_layers=8, | |
| num_attention_heads=4, num_key_value_heads=2, head_dim=16, max_position_embeddings=2048, | |
| tie_word_embeddings=False) | |
| m = Qwen3ForCausalLM(cfg).eval() | |
| with torch.no_grad(): | |
| for p in m.parameters(): | |
| if p.dim() > 1: | |
| p.mul_(4.0) | |
| for layer in m.model.layers[5:]: # deep layers do nothing: the heads are often right | |
| layer.self_attn.o_proj.weight.zero_() | |
| layer.mlp.down_proj.weight.zero_() | |
| heads = [] | |
| for _ in LAYERS: | |
| full = ExitHead(64, cfg.vocab_size, cfg.rms_norm_eps) | |
| with torch.no_grad(): | |
| full.norm.weight.copy_(m.model.norm.weight) | |
| full.proj.weight.copy_(m.lm_head.weight) | |
| heads.append(ExitHead.from_full(full, torch.arange(cfg.vocab_size), reject=True)) | |
| config = {"base_model": "tiny", "default_domain": "web", | |
| "domains": {"web": {"policy": _POLICY, "stop": None}, | |
| "gsm8k": {"policy": _POLICY, "stop": "gsm8k"}, | |
| "code": {"policy": _POLICY, "stop": "humaneval"}}} | |
| return DSSD.from_components(m, CharTokenizer(), heads, LAYERS, config, max_len=2048) | |