Spaces:
Sleeping
Sleeping
File size: 2,336 Bytes
74679d4 | 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 | """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)
|