arthu1 commited on
Commit
91b4b1c
·
verified ·
1 Parent(s): eb9fa4a

Upload folder using huggingface_hub

Browse files
README.md ADDED
@@ -0,0 +1,59 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: apache-2.0
3
+ tags:
4
+ - north-ml
5
+ - flare-1
6
+ - deltanet
7
+ - hybrid-attention
8
+ thumbnail: banner.png
9
+ ---
10
+
11
+ # Fiber Lite — North ML Flare-1 (~1.6B)
12
+
13
+ The mid Flare-1 tier: Qwen3.5-derived gated **DeltaNet + GQA hybrid**
14
+ backbone (hidden 2048) with **AgentRegisters** — persistent learned state
15
+ slots for agent-style memory. This is the full-precision training
16
+ checkpoint (HF `transformers` format). The production 8-bit Kestrel build
17
+ with the learned router, Fiber memory trunk, and search pipeline is at
18
+ [North-ML1/Fiber](https://huggingface.co/North-ML1/Fiber).
19
+
20
+ ## Benchmarks (measured, same harness — Modal L4, greedy, Nov-2025)
21
+
22
+ | Model | ARC-C | MMLU | GSM8K | HumanEval |
23
+ |---|---|---|---|---|
24
+ | ink (725M) | 63.8% | 32.0% | 30.0% | 37.5% |
25
+ | **fiber-lite (1.6B)** | 72.9% | 38.0% | 16.0% | 37.5% |
26
+ | fiber-pro (1.6B) | 80.9% | 40.0% | 14.0% | 48.4% |
27
+ | Qwen3-0.6B (raw) | 49.2% | 29.0% | 4.0% | 4.7% |
28
+ | Qwen3-1.7B (raw) | 72.9% | 39.5% | 43.0% | 68.8% |
29
+ | SmolLM3-3B (raw) | 77.9% | 36.0% | 51.0% | 71.9% |
30
+
31
+ Subsets: ARC-C n=199, MMLU n=200, GSM8K n=100, HumanEval n=64. Honest read:
32
+ ties Qwen3-1.7B on ARC-C and MMLU at the same size class; trails on math
33
+ and code — math reasoning is the known weak spot.
34
+
35
+ ## Specs
36
+
37
+ | | |
38
+ |---|---|
39
+ | Parameters | ~1.63B (1.41B backbone + 158M Fiber trunk + registers + bridge) |
40
+ | Layers | 24 |
41
+ | Hidden | 2048 |
42
+ | Attention | GQA 8 heads / 2 KV, DeltaNet:GQA = 3:1 |
43
+ | Context | 262,144 tokens |
44
+ | Vocab | 65,536 |
45
+
46
+ ## Interactive chat
47
+
48
+ ```bash
49
+ pip install torch transformers safetensors huggingface_hub
50
+ python3 chat.py --model North-ML1/fiber-lite
51
+ ```
52
+
53
+ `chat.py` (in this repo) auto-downloads the weights, attaches the
54
+ AgentRegisters, and starts a REPL with optional two-phase thinking.
55
+
56
+ ## Links
57
+
58
+ - Architecture report: [Flare-1 paper](https://huggingface.co/spaces/North-ML1/flare-1-paper)
59
+ - Runtime artifacts: [Ink-1](https://huggingface.co/North-ML1/Ink-1) · [Fiber](https://huggingface.co/North-ML1/Fiber) · [Fiber-pro](https://huggingface.co/North-ML1/Fiber-pro)
agent_register.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:794a2d3bbd503bb521b6dca8d4e3ff38f0b2ee935e348a6bac39ecf0c3ce8a34
3
+ size 1188461
chat.py ADDED
@@ -0,0 +1,152 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Interactive chat with a North ML Flare model (HF-format checkpoint).
2
+
3
+ python3 chat.py --model North-ML1/fiber-pro
4
+ python3 chat.py --model North-ML1/ink --no-think
5
+
6
+ Requires: pip install torch transformers safetensors huggingface_hub
7
+ Commands: /think on|off /new /quit
8
+ """
9
+ import argparse, json, os, sys, threading, importlib.util
10
+
11
+ IM_S = "<" + "|im_start" + "|>"
12
+ IM_E = "<" + "|im_end" + "|>"
13
+ TH_S = "<" + "think" + ">"
14
+ TH_E = "<" + "/think" + ">"
15
+ EOT = "<" + "|endoftext" + "|>"
16
+
17
+
18
+ def pick_device():
19
+ import torch
20
+ if torch.cuda.is_available():
21
+ return "cuda", torch.bfloat16
22
+ if torch.backends.mps.is_available():
23
+ return "mps", torch.float16
24
+ return "cpu", torch.float32
25
+
26
+
27
+ def load(repo):
28
+ import torch
29
+ from huggingface_hub import snapshot_download
30
+ from transformers import AutoModelForCausalLM, AutoTokenizer
31
+
32
+ local = snapshot_download(repo) if not os.path.isdir(repo) else repo
33
+ device, dtype = pick_device()
34
+ tok = AutoTokenizer.from_pretrained(local)
35
+ model = AutoModelForCausalLM.from_pretrained(local, dtype=dtype).to(device).eval()
36
+
37
+ arch_p = os.path.join(local, "starlight_arch.py")
38
+ reg_p = os.path.join(local, "agent_register.pt")
39
+ meta_p = os.path.join(local, "starlight.json")
40
+ if all(os.path.exists(p) for p in (arch_p, reg_p, meta_p)):
41
+ spec = importlib.util.spec_from_file_location("starlight_arch", arch_p)
42
+ arch = importlib.util.module_from_spec(spec)
43
+ spec.loader.exec_module(arch)
44
+ meta = json.load(open(meta_p))
45
+ regs = arch.attach_registers(model, arch.hidden_size_of(model),
46
+ slots=int(meta.get("slots", 16)))
47
+ regs.load_state_dict(torch.load(reg_p, map_location="cpu")["registers"])
48
+ regs.to(device=device, dtype=dtype)
49
+ print("registers attached (%s slots)" % meta.get("slots", 16))
50
+ print("loaded on %s (%s)" % (device, dtype))
51
+ return model, tok, device
52
+
53
+
54
+ def render(history, think):
55
+ s = ""
56
+ for role, text in history:
57
+ s += IM_S + role + "\n" + text + IM_E + "\n"
58
+ s += IM_S + "assistant\n"
59
+ if think:
60
+ s += TH_S + "\n"
61
+ return s
62
+
63
+
64
+ def gen_stream(model, tok, device, input_ids, cap, eos_ids):
65
+ """Generate with streaming stdout. Returns (text, full_ids)."""
66
+ import torch
67
+ from transformers import TextIteratorStreamer
68
+ streamer = TextIteratorStreamer(tok, skip_special_tokens=True, skip_prompt=True)
69
+ holder = {}
70
+
71
+ def _run():
72
+ holder["out"] = model.generate(
73
+ input_ids=input_ids, max_new_tokens=cap, do_sample=False,
74
+ repetition_penalty=1.05, eos_token_id=eos_ids,
75
+ pad_token_id=tok.eos_token_id, streamer=streamer)
76
+
77
+ t = threading.Thread(target=_run)
78
+ t.start()
79
+ text = ""
80
+ for chunk in streamer:
81
+ text += chunk
82
+ sys.stdout.write(chunk)
83
+ sys.stdout.flush()
84
+ t.join()
85
+ return text, holder["out"]
86
+
87
+
88
+ def chat_turn(model, tok, device, history, think):
89
+ import torch
90
+ im_end = tok.convert_tokens_to_ids(IM_E)
91
+ eot = tok.convert_tokens_to_ids(EOT)
92
+ th_e = tok.convert_tokens_to_ids(TH_E)
93
+ eos = list({i for i in (tok.eos_token_id, im_end, eot) if i is not None})
94
+
95
+ prompt = render(history, think)
96
+ ids = tok(prompt, add_special_tokens=False, return_tensors="pt").input_ids.to(device)
97
+
98
+ if think:
99
+ sys.stdout.write("\n[think] ")
100
+ th_text, out1 = gen_stream(model, tok, device, ids, 220, eos + [th_e])
101
+ gen_ids = out1[0].tolist()
102
+ if th_e not in gen_ids[ids.shape[1]:]:
103
+ extra = tok(TH_E + "\n", add_special_tokens=False, return_tensors="pt")
104
+ ids2 = torch.cat([out1, extra.input_ids.to(device)], dim=1)
105
+ else:
106
+ ids2 = out1
107
+ sys.stdout.write("\n[answer] ")
108
+ ans_text, _ = gen_stream(model, tok, device, ids2, 500, eos)
109
+ sys.stdout.write("\n")
110
+ return ans_text
111
+ ans_text, _ = gen_stream(model, tok, device, ids, 500, eos)
112
+ sys.stdout.write("\n")
113
+ return ans_text
114
+
115
+
116
+ def main():
117
+ ap = argparse.ArgumentParser()
118
+ ap.add_argument("--model", default="North-ML1/fiber-pro")
119
+ ap.add_argument("--no-think", action="store_true")
120
+ a = ap.parse_args()
121
+
122
+ print("downloading %s ..." % a.model)
123
+ model, tok, device = load(a.model)
124
+ think = not a.no_think
125
+ history = []
126
+ print("\nReady. Ask anything. /think on|off /new /quit\n")
127
+ while True:
128
+ try:
129
+ q = input("you> ").strip()
130
+ except (EOFError, KeyboardInterrupt):
131
+ print()
132
+ break
133
+ if not q:
134
+ continue
135
+ if q in ("/quit", "/q", "exit"):
136
+ break
137
+ if q == "/new":
138
+ history = []
139
+ print("(cleared)\n")
140
+ continue
141
+ if q.startswith("/think"):
142
+ think = "off" not in q
143
+ print("think =", "on" if think else "off", "\n")
144
+ continue
145
+ history.append(("user", q))
146
+ ans = chat_turn(model, tok, device, history, think)
147
+ history.append(("assistant", ans.strip()))
148
+ sys.stdout.write("\n")
149
+
150
+
151
+ if __name__ == "__main__":
152
+ main()
chat_template.jinja ADDED
@@ -0,0 +1,154 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {%- set image_count = namespace(value=0) %}
2
+ {%- set video_count = namespace(value=0) %}
3
+ {%- macro render_content(content, do_vision_count, is_system_content=false) %}
4
+ {%- if content is string %}
5
+ {{- content }}
6
+ {%- elif content is iterable and content is not mapping %}
7
+ {%- for item in content %}
8
+ {%- if 'image' in item or 'image_url' in item or item.type == 'image' %}
9
+ {%- if is_system_content %}
10
+ {{- raise_exception('System message cannot contain images.') }}
11
+ {%- endif %}
12
+ {%- if do_vision_count %}
13
+ {%- set image_count.value = image_count.value + 1 %}
14
+ {%- endif %}
15
+ {%- if add_vision_id %}
16
+ {{- 'Picture ' ~ image_count.value ~ ': ' }}
17
+ {%- endif %}
18
+ {{- '<|vision_start|><|image_pad|><|vision_end|>' }}
19
+ {%- elif 'video' in item or item.type == 'video' %}
20
+ {%- if is_system_content %}
21
+ {{- raise_exception('System message cannot contain videos.') }}
22
+ {%- endif %}
23
+ {%- if do_vision_count %}
24
+ {%- set video_count.value = video_count.value + 1 %}
25
+ {%- endif %}
26
+ {%- if add_vision_id %}
27
+ {{- 'Video ' ~ video_count.value ~ ': ' }}
28
+ {%- endif %}
29
+ {{- '<|vision_start|><|video_pad|><|vision_end|>' }}
30
+ {%- elif 'text' in item %}
31
+ {{- item.text }}
32
+ {%- else %}
33
+ {{- raise_exception('Unexpected item type in content.') }}
34
+ {%- endif %}
35
+ {%- endfor %}
36
+ {%- elif content is none or content is undefined %}
37
+ {{- '' }}
38
+ {%- else %}
39
+ {{- raise_exception('Unexpected content type.') }}
40
+ {%- endif %}
41
+ {%- endmacro %}
42
+ {%- if not messages %}
43
+ {{- raise_exception('No messages provided.') }}
44
+ {%- endif %}
45
+ {%- if tools and tools is iterable and tools is not mapping %}
46
+ {{- '<|im_start|>system\n' }}
47
+ {{- "# Tools\n\nYou have access to the following functions:\n\n<tools>" }}
48
+ {%- for tool in tools %}
49
+ {{- "\n" }}
50
+ {{- tool | tojson }}
51
+ {%- endfor %}
52
+ {{- "\n</tools>" }}
53
+ {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n</IMPORTANT>' }}
54
+ {%- if messages[0].role == 'system' %}
55
+ {%- set content = render_content(messages[0].content, false, true)|trim %}
56
+ {%- if content %}
57
+ {{- '\n\n' + content }}
58
+ {%- endif %}
59
+ {%- endif %}
60
+ {{- '<|im_end|>\n' }}
61
+ {%- else %}
62
+ {%- if messages[0].role == 'system' %}
63
+ {%- set content = render_content(messages[0].content, false, true)|trim %}
64
+ {{- '<|im_start|>system\n' + content + '<|im_end|>\n' }}
65
+ {%- endif %}
66
+ {%- endif %}
67
+ {%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
68
+ {%- for message in messages[::-1] %}
69
+ {%- set index = (messages|length - 1) - loop.index0 %}
70
+ {%- if ns.multi_step_tool and message.role == "user" %}
71
+ {%- set content = render_content(message.content, false)|trim %}
72
+ {%- if not(content.startswith('<tool_response>') and content.endswith('</tool_response>')) %}
73
+ {%- set ns.multi_step_tool = false %}
74
+ {%- set ns.last_query_index = index %}
75
+ {%- endif %}
76
+ {%- endif %}
77
+ {%- endfor %}
78
+ {%- if ns.multi_step_tool %}
79
+ {{- raise_exception('No user query found in messages.') }}
80
+ {%- endif %}
81
+ {%- for message in messages %}
82
+ {%- set content = render_content(message.content, true)|trim %}
83
+ {%- if message.role == "system" %}
84
+ {%- if not loop.first %}
85
+ {{- raise_exception('System message must be at the beginning.') }}
86
+ {%- endif %}
87
+ {%- elif message.role == "user" %}
88
+ {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
89
+ {%- elif message.role == "assistant" %}
90
+ {%- set reasoning_content = '' %}
91
+ {%- if message.reasoning_content is string %}
92
+ {%- set reasoning_content = message.reasoning_content %}
93
+ {%- else %}
94
+ {%- if '</think>' in content %}
95
+ {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
96
+ {%- set content = content.split('</think>')[-1].lstrip('\n') %}
97
+ {%- endif %}
98
+ {%- endif %}
99
+ {%- set reasoning_content = reasoning_content|trim %}
100
+ {%- if loop.index0 > ns.last_query_index %}
101
+ {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content + '\n</think>\n\n' + content }}
102
+ {%- else %}
103
+ {{- '<|im_start|>' + message.role + '\n' + content }}
104
+ {%- endif %}
105
+ {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %}
106
+ {%- for tool_call in message.tool_calls %}
107
+ {%- if tool_call.function is defined %}
108
+ {%- set tool_call = tool_call.function %}
109
+ {%- endif %}
110
+ {%- if loop.first %}
111
+ {%- if content|trim %}
112
+ {{- '\n\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
113
+ {%- else %}
114
+ {{- '<tool_call>\n<function=' + tool_call.name + '>\n' }}
115
+ {%- endif %}
116
+ {%- else %}
117
+ {{- '\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
118
+ {%- endif %}
119
+ {%- if tool_call.arguments is defined %}
120
+ {%- for args_name, args_value in tool_call.arguments|items %}
121
+ {{- '<parameter=' + args_name + '>\n' }}
122
+ {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %}
123
+ {{- args_value }}
124
+ {{- '\n</parameter>\n' }}
125
+ {%- endfor %}
126
+ {%- endif %}
127
+ {{- '</function>\n</tool_call>' }}
128
+ {%- endfor %}
129
+ {%- endif %}
130
+ {{- '<|im_end|>\n' }}
131
+ {%- elif message.role == "tool" %}
132
+ {%- if loop.previtem and loop.previtem.role != "tool" %}
133
+ {{- '<|im_start|>user' }}
134
+ {%- endif %}
135
+ {{- '\n<tool_response>\n' }}
136
+ {{- content }}
137
+ {{- '\n</tool_response>' }}
138
+ {%- if not loop.last and loop.nextitem.role != "tool" %}
139
+ {{- '<|im_end|>\n' }}
140
+ {%- elif loop.last %}
141
+ {{- '<|im_end|>\n' }}
142
+ {%- endif %}
143
+ {%- else %}
144
+ {{- raise_exception('Unexpected message role.') }}
145
+ {%- endif %}
146
+ {%- endfor %}
147
+ {%- if add_generation_prompt %}
148
+ {{- '<|im_start|>assistant\n' }}
149
+ {%- if enable_thinking is defined and enable_thinking is true %}
150
+ {{- '<think>\n' }}
151
+ {%- else %}
152
+ {{- '<think>\n\n</think>\n\n' }}
153
+ {%- endif %}
154
+ {%- endif %}
config.json ADDED
@@ -0,0 +1,75 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "Qwen3_5ForCausalLM"
4
+ ],
5
+ "attention_bias": false,
6
+ "attention_dropout": 0.0,
7
+ "attn_output_gate": true,
8
+ "bos_token_id": null,
9
+ "dtype": "bfloat16",
10
+ "eos_token_id": 65503,
11
+ "full_attention_interval": 4,
12
+ "head_dim": 256,
13
+ "hidden_act": "silu",
14
+ "hidden_size": 2048,
15
+ "initializer_range": 0.02,
16
+ "intermediate_size": 5504,
17
+ "layer_types": [
18
+ "linear_attention",
19
+ "linear_attention",
20
+ "linear_attention",
21
+ "full_attention",
22
+ "linear_attention",
23
+ "linear_attention",
24
+ "linear_attention",
25
+ "full_attention",
26
+ "linear_attention",
27
+ "linear_attention",
28
+ "linear_attention",
29
+ "full_attention",
30
+ "linear_attention",
31
+ "linear_attention",
32
+ "linear_attention",
33
+ "full_attention",
34
+ "linear_attention",
35
+ "linear_attention",
36
+ "linear_attention",
37
+ "full_attention",
38
+ "linear_attention",
39
+ "linear_attention",
40
+ "linear_attention",
41
+ "full_attention"
42
+ ],
43
+ "linear_conv_kernel_dim": 4,
44
+ "linear_key_head_dim": 128,
45
+ "linear_num_key_heads": 16,
46
+ "linear_num_value_heads": 16,
47
+ "linear_value_head_dim": 128,
48
+ "mamba_ssm_dtype": "float32",
49
+ "max_position_embeddings": 262144,
50
+ "mlp_only_layers": [],
51
+ "model_type": "qwen3_5_text",
52
+ "mtp_num_hidden_layers": 1,
53
+ "mtp_use_dedicated_embeddings": false,
54
+ "num_attention_heads": 8,
55
+ "num_hidden_layers": 24,
56
+ "num_key_value_heads": 2,
57
+ "pad_token_id": null,
58
+ "partial_rotary_factor": 0.25,
59
+ "rms_norm_eps": 1e-06,
60
+ "rope_parameters": {
61
+ "mrope_interleaved": true,
62
+ "mrope_section": [
63
+ 11,
64
+ 11,
65
+ 10
66
+ ],
67
+ "partial_rotary_factor": 0.25,
68
+ "rope_theta": 10000000,
69
+ "rope_type": "default"
70
+ },
71
+ "tie_word_embeddings": true,
72
+ "transformers_version": "5.18.0",
73
+ "use_cache": false,
74
+ "vocab_size": 65536
75
+ }
generation_config.json ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ {
2
+ "_from_model_config": true,
3
+ "eos_token_id": 65503,
4
+ "transformers_version": "5.18.0",
5
+ "use_cache": true
6
+ }
model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:da49bcd0acbe36d7d1776edd67c1fc43d16041fa13cf54351342dc03d9121005
3
+ size 2826264920
starlight.json ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ {
2
+ "product": "Starlight Pro 1",
3
+ "base": "Qwen/Qwen3.5-2B",
4
+ "step": 2000,
5
+ "arch": "agent-register",
6
+ "slots": 16
7
+ }
starlight_arch.py ADDED
@@ -0,0 +1,70 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Starlight agent memory.
2
+
3
+ Qwen3.5 stays the backbone. After each full-attention layer, a causal register
4
+ mixes a low-rank memory back into the residual stream. The output projection
5
+ starts at zero, so step 0 matches Qwen. The input projection does not, so the
6
+ register gets a gradient and can learn.
7
+ """
8
+ from __future__ import annotations
9
+
10
+ import torch
11
+ from torch import nn
12
+
13
+
14
+ class AgentRegister(nn.Module):
15
+ def __init__(self, hidden: int, slots: int = 16):
16
+ super().__init__()
17
+ self.slots = slots
18
+ self.down = nn.Linear(hidden, slots, bias=False)
19
+ self.up = nn.Linear(slots, hidden, bias=False)
20
+ self.gate = nn.Linear(hidden, slots)
21
+ nn.init.normal_(self.down.weight, std=0.02)
22
+ nn.init.zeros_(self.up.weight)
23
+ nn.init.zeros_(self.gate.weight)
24
+ nn.init.constant_(self.gate.bias, -2.0)
25
+
26
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
27
+ gate = torch.sigmoid(self.gate(hidden_states))
28
+ memory = torch.cumsum(gate * self.down(hidden_states), dim=1)
29
+ return hidden_states + self.up(memory)
30
+
31
+
32
+ def full_attention_indexes(model) -> list[int]:
33
+ config = model.config
34
+ kinds = list(getattr(config, "layer_types", None) or [])
35
+ if not kinds and getattr(config, "text_config", None) is not None:
36
+ kinds = list(getattr(config.text_config, "layer_types", None) or [])
37
+ return [i for i, kind in enumerate(kinds) if kind == "full_attention"]
38
+
39
+
40
+ def attach_registers(model, hidden: int, slots: int = 16) -> nn.ModuleList:
41
+ indexes = full_attention_indexes(model)
42
+ if not indexes:
43
+ raise RuntimeError("Qwen config has no full_attention layers")
44
+ registers = nn.ModuleList(AgentRegister(hidden, slots) for _ in indexes)
45
+ device = next(model.parameters()).device
46
+ dtype = next(model.parameters()).dtype
47
+ registers.to(device=device, dtype=dtype)
48
+
49
+ def make_hook(register: AgentRegister):
50
+ def hook(_module, _inputs, output):
51
+ if isinstance(output, tuple):
52
+ return (register(output[0]),) + output[1:]
53
+ return register(output)
54
+
55
+ return hook
56
+
57
+ layers = model.model.layers
58
+ for index, register in zip(indexes, registers):
59
+ layers[index].register_forward_hook(make_hook(register))
60
+ return registers
61
+
62
+
63
+ def hidden_size_of(model) -> int:
64
+ config = model.config
65
+ size = getattr(config, "hidden_size", None)
66
+ if not size and getattr(config, "text_config", None) is not None:
67
+ size = getattr(config.text_config, "hidden_size", None)
68
+ if not size:
69
+ raise RuntimeError("missing hidden size")
70
+ return int(size)
tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
tokenizer_config.json ADDED
@@ -0,0 +1,33 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "add_prefix_space": false,
3
+ "audio_bos_token": "<|audio_start|>",
4
+ "audio_eos_token": "<|audio_end|>",
5
+ "audio_token": "<|audio_pad|>",
6
+ "backend": "tokenizers",
7
+ "bos_token": null,
8
+ "clean_up_tokenization_spaces": false,
9
+ "eos_token": "<|im_end|>",
10
+ "errors": "replace",
11
+ "fix_mistral_regex": true,
12
+ "image_token": "<|image_pad|>",
13
+ "is_local": true,
14
+ "local_files_only": false,
15
+ "model_max_length": 262144,
16
+ "model_specific_special_tokens": {
17
+ "audio_bos_token": "<|audio_start|>",
18
+ "audio_eos_token": "<|audio_end|>",
19
+ "audio_token": "<|audio_pad|>",
20
+ "image_token": "<|image_pad|>",
21
+ "video_token": "<|video_pad|>",
22
+ "vision_bos_token": "<|vision_start|>",
23
+ "vision_eos_token": "<|vision_end|>"
24
+ },
25
+ "pad_token": "<|endoftext|>",
26
+ "pretokenize_regex": "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?[\\p{L}\\p{M}]+|\\p{N}| ?[^\\s\\p{L}\\p{M}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
27
+ "split_special_tokens": false,
28
+ "tokenizer_class": "Qwen2Tokenizer",
29
+ "unk_token": null,
30
+ "video_token": "<|video_pad|>",
31
+ "vision_bos_token": "<|vision_start|>",
32
+ "vision_eos_token": "<|vision_end|>"
33
+ }