Add files using upload-large-folder tool

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,129 @@
+---
+license: gemma
+base_model: google/gemma-4-E2B-it
+pipeline_tag: text-generation
+tags:
+  - gemma4
+  - hybrid
+  - custom_code
+  - custom_generate
+---
+
+# gemma-4-e2b-it-hybrid
+
+`Cactus-Compute/gemma-4-e2b-it-hybrid` is **google/gemma-4-E2B-it plus a handoff probe**: a
+small head (weight prefix `handoff_probe.*`) that scores every generation with
+
+```
+confidence = 1 - p_wrong
+```
+
+so a client can decide when to keep the on-device answer and when to hand off to
+a cloud model. The base weights are byte-identical to the stock checkpoint (same
+keys); the repo adds eleven probe tensors, a remote-code model class and a
+`custom_generate` recipe. **Stock engine commands work unchanged** — you only add
+`--trust-remote-code` / `trust_remote_code=True`.
+
+## Serving with stock transformers
+
+```bash
+transformers serve --trust-remote-code
+# then request model "Cactus-Compute/gemma-4-e2b-it-hybrid" via the OpenAI-compatible API
+```
+
+or interactively:
+
+```bash
+transformers chat Cactus-Compute/gemma-4-e2b-it-hybrid --trust-remote-code
+```
+
+### How the confidence reaches you (in-band trailer)
+
+`transformers serve` cannot add response fields, so the score travels **in-band**:
+the assistant content's final line is
+
+```
+\n[[hybrid:confidence=0.7812]]
+```
+
+always exactly 4 decimals, ASCII, `confidence` in `[0, 1]`. Strip the final
+`[[hybrid:...]]` line before display and parse the float for routing. The
+trailer is emitted in both streaming and non-streaming modes.
+
+### Limitations
+
+- `--continuous-batching`: not supported — the CB scheduler bypasses
+  `generate()`, so no probe runs and **no trailer is emitted**. Serve without
+  `--continuous-batching` to get confidence scores.
+- The trailer (and confidence) is only produced for single-sequence decoding:
+  batch size 1, no beam search, no assisted/speculative decoding. Unsupported
+  modes fall back to stock behavior (no trailer).
+- The probe scores at most the first 1024 generated tokens.
+- The trailer's token ids are appended to the returned sequences, so reported
+  completion token counts include the trailer (a handful of tokens).
+
+## Python usage
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model = AutoModelForCausalLM.from_pretrained(
+    "Cactus-Compute/gemma-4-e2b-it-hybrid", trust_remote_code=True, dtype="auto"
+)
+tokenizer = AutoTokenizer.from_pretrained("Cactus-Compute/gemma-4-e2b-it-hybrid")
+inputs = tokenizer.apply_chat_template(
+    [{"role": "user", "content": "Explain why this contract clause is risky."}],
+    add_generation_prompt=True, return_tensors="pt",
+).to(model.device)
+
+# Structured API: clean sequences + raw float (no in-band trailer).
+sequences, confidence = model.generate_with_confidence(inputs, max_new_tokens=512)
+print(confidence)                     # e.g. 0.7812
+print(model.last_confidence)          # same value
+
+# Stock generate (custom_generate recipe): plain tensor + in-band trailer.
+sequences = model.generate(inputs, max_new_tokens=512)
+
+# Or opt into a rich return with the raw float:
+out = model.generate(inputs, max_new_tokens=512, return_confidence=True)
+out.sequences, out.confidence, out.trailer_text
+
+# Suppress the trailer while keeping stock behavior:
+sequences = model.generate(inputs, max_new_tokens=512, emit_trailer=False)
+```
+
+## Probe contract
+
+- Input: float32 `[T, 1536]` — output of decoder layer index 28
+  (`config.probe_layer`), captured at the position that predicts each generated
+  token: row 0 = last prompt position at prefill, row t = position captured at
+  generation step t. Only the first 1024 rows are scored.
+- Math (float32): `x = LayerNorm(x, eps=1e-5) * norm.weight + norm.bias`;
+  `p = relu(x @ proj.weight.T + proj.bias)`;
+  `s = p @ attn_query / sqrt(32)`; `w = softmax_T(s - max)`; `pooled = w @ p`;
+  `h = relu(head.0 @ pooled + b)`; `h = relu(head.2 @ h + b)`;
+  `logit = head.4 @ h + b`; `p_wrong = sigmoid(logit)`;
+  `confidence = 1 - p_wrong`.
+- Capture uses a forward hook that keeps only one `[1, 1536]` row per decode
+  step — full hidden-state stacks are never materialized.
+
+## MLX (Apple Silicon)
+
+The repo ships a single-file `mlx-lm` model (`gemma_4_e2b_it_hybrid.py`),
+referenced by `"model_file"` in `config.json` (mlx-lm >= 0.30.1). After
+generation, read `model.last_confidence` (or `model.confidence(num_tokens=N)`).
+An mlx-lm `Model` cannot inject tokens into the stream, so there is **no
+in-band trailer** on MLX today — confidence is Python-API only there.
+
+## Repo contents
+
+| File | Purpose |
+|---|---|
+| `configuration_gemma_4_e2b_it_hybrid.py` | `Gemma4E2BItHybridConfig` (stock Gemma-4 text config + probe hyperparams) |
+| `modeling_gemma_4_e2b_it_hybrid.py` | `Gemma4E2BItHybridForCausalLM` (stock `Gemma4ForCausalLM` + `handoff_probe.*`) |
+| `custom_generate/generate.py` | stock decode loop + confidence + in-band trailer |
+| `model*.safetensors` | base weights (identical keys) + `handoff_probe.*` tensors |
+| `gemma_4_e2b_it_hybrid.py` | single-file `mlx-lm` model, wired via config.json's `model_file` |
+
+Requires `transformers>=5,<6`. Gemma model use remains subject to the Gemma
+terms.

chat_template.jinja

@@ -0,0 +1,386 @@
+{#
+  Template: Google Gemma 4 Canonical Chat Template
+  Author: Google Gemma Engineering Team
+  Published: 2026-07-09
+  Context: Fixed tool-calling loops, turn closures, and thinking content-ordering.
+#}
+{%- macro format_parameters(properties, required, filter_keys=false) -%}
+    {%- set standard_keys = ['description', 'type', 'properties', 'required', 'nullable'] -%}
+    {%- set ns = namespace(found_first=false) -%}
+    {%- for key, value in properties | dictsort -%}
+        {%- set add_comma = false -%}
+        {%- if not filter_keys or key not in standard_keys -%}
+            {%- if ns.found_first %},{% endif -%}
+            {%- set ns.found_first = true -%}
+            {{ key }}:{
+            {%- if value['description'] -%}
+                description:<|"|>{{ value['description'] }}<|"|>
+                {%- set add_comma = true -%}
+            {%- endif -%}
+            {%- if value['type'] | upper == 'STRING' -%}
+                {%- if value['enum'] -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    enum:{{ format_argument(value['enum']) }}
+                {%- endif -%}
+            {%- elif value['type'] | upper == 'ARRAY' -%}
+                {%- if value['items'] is mapping and value['items'] -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    items:{
+                    {%- set ns_items = namespace(found_first=false) -%}
+                    {%- for item_key, item_value in value['items'] | dictsort -%}
+                        {%- if item_value is not none -%}
+                            {%- if ns_items.found_first %},{% endif -%}
+                            {%- set ns_items.found_first = true -%}
+                            {%- if item_key == 'properties' -%}
+                                properties:{
+                                {%- if item_value is mapping -%}
+                                    {{- format_parameters(item_value, value['items']['required'] | default([])) -}}
+                                {%- endif -%}
+                                }
+                            {%- elif item_key == 'required' -%}
+                                required:[
+                                {%- for req_item in item_value -%}
+                                    <|"|>{{- req_item -}}<|"|>
+                                    {%- if not loop.last %},{% endif -%}
+                                {%- endfor -%}
+                                ]
+                            {%- elif item_key == 'type' -%}
+                                {%- if item_value is string -%}
+                                    type:{{ format_argument(item_value | upper) }}
+                                {%- else -%}
+                                    type:{{ format_argument(item_value | map('upper') | list) }}
+                                {%- endif -%}
+                            {%- else -%}
+                                {{ item_key }}:{{ format_argument(item_value) }}
+                            {%- endif -%}
+                        {%- endif -%}
+                    {%- endfor -%}
+                    }
+                {%- endif -%}
+            {%- endif -%}
+            {%- if value['nullable'] %}
+                {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                nullable:true
+            {%- endif -%}
+            {%- if value['type'] | upper == 'OBJECT' -%}
+                {%- if value['properties'] is defined and value['properties'] is mapping -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    properties:{
+                    {{- format_parameters(value['properties'], value['required'] | default([])) -}}
+                    }
+                {%- elif value is mapping -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    properties:{
+                    {{- format_parameters(value, value['required'] | default([]), filter_keys=true) -}}
+                    }
+                {%- endif -%}
+                {%- if value['required'] -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    required:[
+                    {%- for item in value['required'] | default([]) -%}
+                        <|"|>{{- item -}}<|"|>
+                        {%- if not loop.last %},{% endif -%}
+                    {%- endfor -%}
+                    ]
+                {%- endif -%}
+            {%- endif -%}
+            {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+            type:<|"|>{{ value['type'] | upper }}<|"|>}
+        {%- endif -%}
+    {%- endfor -%}
+{%- endmacro -%}
+{%- macro format_function_declaration(tool_data) -%}
+    declaration:{{- tool_data['function']['name'] -}}{description:<|"|>{{- tool_data['function']['description'] -}}<|"|>
+    {%- set params = tool_data['function']['parameters'] -%}
+    {%- if params -%}
+        ,parameters:{
+        {%- if params['properties'] -%}
+            properties:{ {{- format_parameters(params['properties'], params['required']) -}} },
+        {%- endif -%}
+        {%- if params['required'] -%}
+            required:[
+            {%- for item in params['required'] -%}
+                <|"|>{{- item -}}<|"|>
+                {{- ',' if not loop.last -}}
+            {%- endfor -%}
+            ],
+        {%- endif -%}
+        {%- if params['type'] -%}
+            type:<|"|>{{- params['type'] | upper -}}<|"|>}
+        {%- endif -%}
+    {%- endif -%}
+    {%- if 'response' in tool_data['function'] -%}
+        {%- set response_declaration = tool_data['function']['response'] -%}
+        ,response:{
+        {%- if response_declaration['description'] -%}
+            description:<|"|>{{- response_declaration['description'] -}}<|"|>,
+        {%- endif -%}
+        {%- if response_declaration['type'] | upper == 'OBJECT' -%}
+            type:<|"|>{{- response_declaration['type'] | upper -}}<|"|>}
+        {%- endif -%}
+    {%- endif -%}
+    }
+{%- endmacro -%}
+{%- macro format_argument(argument, escape_keys=True) -%}
+    {%- if argument is none -%}
+        {{- 'null' -}}
+    {%- elif argument is string -%}
+        {{- '<|"|>' + argument + '<|"|>' -}}
+    {%- elif argument is boolean -%}
+        {{- 'true' if argument else 'false' -}}
+    {%- elif argument is mapping -%}
+        {{- '{' -}}
+        {%- set ns = namespace(found_first=false) -%}
+        {%- for key, value in argument | dictsort -%}
+            {%- if ns.found_first %},{% endif -%}
+            {%- set ns.found_first = true -%}
+            {%- if escape_keys -%}
+                {{- '<|"|>' + key + '<|"|>' -}}
+            {%- else -%}
+                {{- key -}}
+            {%- endif -%}
+            :{{- format_argument(value, escape_keys=escape_keys) -}}
+        {%- endfor -%}
+        {{- '}' -}}
+    {%- elif argument is sequence -%}
+        {{- '[' -}}
+        {%- for item in argument -%}
+            {{- format_argument(item, escape_keys=escape_keys) -}}
+            {%- if not loop.last %},{% endif -%}
+        {%- endfor -%}
+        {{- ']' -}}
+    {%- else -%}
+        {{- argument -}}
+    {%- endif -%}
+{%- endmacro -%}
+{%- macro strip_thinking(text) -%}
+    {%- set ns = namespace(result='') -%}
+    {%- for part in text.split('<channel|>') -%}
+        {%- if '<|channel>' in part -%}
+            {%- set ns.result = ns.result + part.split('<|channel>')[0] -%}
+        {%- else -%}
+            {%- set ns.result = ns.result + part -%}
+        {%- endif -%}
+    {%- endfor -%}
+    {{- ns.result | trim -}}
+{%- endmacro -%}
+
+{%- macro format_tool_response_block(tool_name, response) -%}
+    {{- '<|tool_response>' -}}
+    {%- if response is mapping -%}
+        {{- 'response:' + tool_name + '{' -}}
+        {%- for key, value in response | dictsort -%}
+            {{- key -}}:{{- format_argument(value, escape_keys=False) -}}
+            {%- if not loop.last %},{% endif -%}
+        {%- endfor -%}
+        {{- '}' -}}
+    {%- else -%}
+        {{- 'response:' + tool_name + '{value:' + format_argument(response, escape_keys=False) + '}' -}}
+    {%- endif -%}
+    {{- '<tool_response|>' -}}
+{%- endmacro -%}
+
+{#- ===== SETUP ===== -#}
+{%- set ns = namespace(prev_message_type=None, prev_non_tool_role=None) -%}
+{%- set loop_messages = messages -%}
+{%- set enable_thinking = enable_thinking | default(false) -%}
+{%- set preserve_thinking = preserve_thinking | default(false) -%}
+{{- bos_token -}}
+{#- Handle System/Tool Definitions Block -#}
+{%- if enable_thinking or tools or (messages and messages[0]['role'] in ['system', 'developer']) -%}
+    {{- '<|turn>system\n' -}}
+    {#- Inject Thinking token at the very top of the FIRST system turn -#}
+    {%- if enable_thinking -%}
+        {{- '<|think|>\n' -}}
+        {%- set ns.prev_message_type = 'think' -%}
+    {%- endif -%}
+    {%- if messages and messages[0]['role'] in ['system', 'developer'] -%}
+        {%- if messages[0]['content'] is string -%}
+            {{- messages[0]['content'] | trim -}}
+        {%- elif messages[0]['content'] is sequence -%}
+            {%- for item in messages[0]['content'] -%}
+                {{- item['text'] | trim + ' '-}}
+            {%- endfor -%}
+        {%- endif -%}
+        {%- set loop_messages = messages[1:] -%}
+    {%- endif -%}
+    {%- if tools -%}
+        {%- for tool in tools %}
+            {{- '<|tool>' -}}
+            {{- format_function_declaration(tool) | trim -}}
+            {{- '<tool|>' -}}
+        {%- endfor %}
+        {%- set ns.prev_message_type = 'tool' -%}
+    {%- endif -%}
+    {{- '<turn|>\n' -}}
+{%- endif %}
+
+{#- Pre-scan: find last user message index for reasoning guard -#}
+{%- set ns_turn = namespace(last_user_idx=-1) -%}
+{%- for i in range(loop_messages | length) -%}
+    {%- if loop_messages[i]['role'] == 'user' -%}
+        {%- set ns_turn.last_user_idx = i -%}
+    {%- endif -%}
+{%- endfor -%}
+
+{#- Loop through messages -#}
+{%- for message in loop_messages -%}
+    {%- if message['role'] != 'tool' -%}
+    {%- set ns.prev_message_type = None -%}
+    {%- set role = 'model' if message['role'] == 'assistant' else message['role'] -%}
+    {#- Detect continuation using tracked state — O(1) instead of O(n) backward scan -#}
+    {%- set continue_same_model_turn = (role == 'model' and ns.prev_non_tool_role == 'assistant') -%}
+    {%- if not continue_same_model_turn -%}
+        {{- '<|turn>' + role + '\n' }}
+    {%- endif -%}
+
+    {#- Render reasoning/reasoning_content as thinking channel -#}
+    {%- set thinking_text = message.get('reasoning') or message.get('reasoning_content') -%}
+    {%- set thinking_gate = (loop.index0 > ns_turn.last_user_idx) or (preserve_thinking and message.get('tool_calls')) -%}
+    {%- if thinking_text and thinking_gate -%}
+        {{- '<|channel>thought\n' + thinking_text + '\n<channel|>' -}}
+    {%- endif -%}
+
+            {%- if message.get('tool_calls') -%}
+                {%- for tool_call in message.get('tool_calls') -%}
+                    {%- set function = tool_call['function'] -%}
+                    {{- '<|tool_call>call:' + function['name'] + '{' -}}
+                    {%- if function['arguments'] is mapping -%}
+                        {%- set ns_args = namespace(found_first=false) -%}
+                        {%- for key, value in function['arguments'] | dictsort -%}
+                            {%- if ns_args.found_first %},{% endif -%}
+                            {%- set ns_args.found_first = true -%}
+                            {{- key -}}:{{- format_argument(value, escape_keys=False) -}}
+                        {%- endfor -%}
+                    {%- elif function['arguments'] is none -%}
+                    {%- else -%}
+                        {{- raise_exception(
+                            "chat_template: tool_calls[].function.arguments must be a "
+                            "JSON object (mapping), not a string. Deserialize arguments "
+                            "before passing to the template."
+                        ) -}}
+                    {%- endif -%}
+                    {{- '}<tool_call|>' -}}
+                {%- endfor -%}
+                {%- set ns.prev_message_type = 'tool_call' -%}
+            {%- endif -%}
+
+            {%- set ns_tr_out = namespace(flag=false) -%}
+            {%- if message.get('tool_responses') -%}
+                {#- Legacy: tool_responses embedded on the assistant message (Google/Gemma native) -#}
+                {%- for tool_response in message.get('tool_responses') -%}
+                    {{- format_tool_response_block(tool_response['name'] | default('unknown', true), tool_response['response']) -}}
+                    {%- set ns_tr_out.flag = true -%}
+                    {%- set ns.prev_message_type = 'tool_response' -%}
+                {%- endfor -%}
+            {%- elif message.get('tool_calls') -%}
+                {#- OpenAI Chat Completions: forward-scan consecutive role:tool messages -#}
+                {%- set ns_tool_scan = namespace(stopped=false) -%}
+                {%- for k in range(loop.index0 + 1, loop_messages | length) -%}
+                    {%- if ns_tool_scan.stopped -%}
+                    {%- elif loop_messages[k]['role'] != 'tool' -%}
+                        {%- set ns_tool_scan.stopped = true -%}
+                    {%- else -%}
+                        {%- set follow = loop_messages[k] -%}
+                        {#- Resolve tool_call_id to function name -#}
+                        {%- set ns_tname = namespace(name=follow.get('name') or 'unknown') -%}
+                        {%- for tc in message.get('tool_calls') -%}
+                            {%- if tc.get('id') == follow.get('tool_call_id') -%}
+                                {%- set ns_tname.name = tc['function']['name'] -%}
+                            {%- endif -%}
+                        {%- endfor -%}
+                        {#- Handle content as string or content-parts array -#}
+                        {%- set tool_body = follow.get('content') -%}
+                        {%- if tool_body is string -%}
+                            {{- format_tool_response_block(ns_tname.name, tool_body) -}}
+                        {%- elif tool_body is sequence and tool_body is not string -%}
+                            {%- set ns_txt = namespace(s='') -%}
+                            {%- for part in tool_body -%}
+                                {%- if part.get('type') == 'text' -%}
+                                    {%- set ns_txt.s = ns_txt.s + (part.get('text') | default('')) -%}
+                                {%- endif -%}
+                            {%- endfor -%}
+                            {{- format_tool_response_block(ns_tname.name, ns_txt.s) -}}
+                            {%- for part in tool_body -%}
+                                {%- if part.get('type') in ['image', 'image_url'] -%}
+                                    {{- '<|image|>' -}}
+                                {%- elif part.get('type') in ['audio', 'input_audio'] -%}
+                                    {{- '<|audio|>' -}}
+                                {%- elif part.get('type') == 'video' -%}
+                                    {{- '<|video|>' -}}
+                                {%- endif -%}
+                            {%- endfor -%}
+                        {%- else -%}
+                            {{- format_tool_response_block(ns_tname.name, tool_body) -}}
+                        {%- endif -%}
+                        {%- set ns_tr_out.flag = true -%}
+                        {%- set ns.prev_message_type = 'tool_response' -%}
+                    {%- endif -%}
+                {%- endfor -%}
+            {%- endif -%}
+
+            {%- set captured_content -%}
+            {%- if message.get('content') is string -%}
+                {%- if role == 'model' -%}
+                    {{- strip_thinking(message['content']) -}}
+                {%- else -%}
+                    {{- message['content'] | trim -}}
+                {%- endif -%}
+            {%- elif message.get('content') is sequence -%}
+                {%- for item in message['content'] -%}
+                    {%- if item.get('type') == 'text' -%}
+                        {%- if role == 'model' -%}
+                            {{- strip_thinking(item['text']) -}}
+                        {%- else -%}
+                            {{- item['text'] | trim -}}
+                        {%- endif -%}
+                    {%- elif item.get('type') in ['image', 'image_url'] -%}
+                        {{- '<|image|>' -}}
+                    {%- elif item.get('type') in ['audio', 'input_audio'] -%}
+                        {{- '<|audio|>' -}}
+                    {%- elif item.get('type') == 'video' -%}
+                        {{- '<|video|>' -}}
+                    {%- endif -%}
+                {%- endfor -%}
+            {%- endif -%}
+            {%- endset -%}
+
+            {{- captured_content -}}
+            {%- set has_content = captured_content | trim | length > 0 -%}
+
+        {#- Forward-scan: find next non-tool message role for continuation detection -#}
+        {%- set next_nt = namespace(role=None, found=false) -%}
+        {%- for j in range(loop.index0 + 1, loop_messages | length) -%}
+            {%- if not next_nt.found -%}
+                {%- if loop_messages[j]['role'] != 'tool' -%}
+                    {%- set next_nt.role = loop_messages[j]['role'] -%}
+                    {%- set next_nt.found = true -%}
+                {%- endif -%}
+            {%- endif -%}
+        {%- endfor -%}
+
+        {%- set continues_into_next = (
+            role == 'model'
+            and next_nt.role == 'assistant'
+            and (not message.get('tool_calls') or ns_tr_out.flag)
+        ) -%}
+
+        {%- if ns.prev_message_type == 'tool_call' and not ns_tr_out.flag -%}
+            {{- '<|tool_response>' -}}
+        {%- elif continues_into_next -%}
+        {%- elif not (ns_tr_out.flag and not has_content and not next_nt.found) -%}
+            {{- '<turn|>\n' -}}
+        {%- endif -%}
+
+    {#- Track previous non-tool role for next iteration (avoids O(n) backward scan) -#}
+    {%- set ns.prev_non_tool_role = message['role'] -%}
+    {%- endif -%}
+{%- endfor -%}
+
+{%- if add_generation_prompt -%}
+    {%- if ns.prev_message_type != 'tool_response' and ns.prev_message_type != 'tool_call' -%}
+        {{- '<|turn>model\n' -}}
+    {%- elif ns.prev_message_type == 'tool_response' and enable_thinking -%}
+        {{- '<|channel>thought\n' -}}
+    {%- endif -%}
+{%- endif -%}

config.json

@@ -0,0 +1,96 @@
+{
+  "architectures": [
+    "Gemma4E2BItHybridForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "attention_k_eq_v": false,
+  "auto_map": {
+    "AutoConfig": "configuration_gemma_4_e2b_it_hybrid.Gemma4E2BItHybridConfig",
+    "AutoModelForCausalLM": "modeling_gemma_4_e2b_it_hybrid.Gemma4E2BItHybridForCausalLM"
+  },
+  "bos_token_id": 2,
+  "dtype": "bfloat16",
+  "enable_moe_block": false,
+  "eos_token_id": 1,
+  "expert_intermediate_size": null,
+  "final_logit_softcapping": 30.0,
+  "global_head_dim": 512,
+  "head_dim": 256,
+  "hidden_activation": "gelu_pytorch_tanh",
+  "hidden_size": 1536,
+  "hidden_size_per_layer_input": 256,
+  "initializer_range": 0.02,
+  "intermediate_size": 6144,
+  "layer_types": [
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 131072,
+  "model_file": "gemma_4_e2b_it_hybrid.py",
+  "model_type": "gemma-4-e2b-it-hybrid",
+  "num_attention_heads": 8,
+  "num_experts": null,
+  "num_global_key_value_heads": null,
+  "num_hidden_layers": 35,
+  "num_key_value_heads": 1,
+  "num_kv_shared_layers": 20,
+  "pad_token_id": 0,
+  "probe_feature_size": 1536,
+  "probe_layer": 28,
+  "probe_max_tokens": 1024,
+  "probe_proj_dim": 32,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "full_attention": {
+      "partial_rotary_factor": 0.25,
+      "rope_theta": 1000000.0,
+      "rope_type": "proportional"
+    },
+    "sliding_attention": {
+      "rope_theta": 10000.0,
+      "rope_type": "default"
+    }
+  },
+  "sliding_window": 512,
+  "tie_word_embeddings": true,
+  "top_k_experts": null,
+  "use_bidirectional_attention": null,
+  "use_cache": true,
+  "use_double_wide_mlp": true,
+  "vocab_size": 262144,
+  "vocab_size_per_layer_input": 262144
+}

configuration_gemma_4_e2b_it_hybrid.py

@@ -0,0 +1,46 @@
+"""Configuration for Cactus-Compute/gemma-4-e2b-it-hybrid.
+
+`Gemma4E2BItHybridConfig` is the stock Gemma-4 text config plus the hyperparameters of
+the handoff probe that scores every generation with ``confidence = 1 - p_wrong``.
+Everything the base model needs is inherited unchanged, so the base weights load
+with identical keys.
+"""
+
+from transformers.models.gemma4.configuration_gemma4 import Gemma4TextConfig
+
+
+class Gemma4E2BItHybridConfig(Gemma4TextConfig):
+    """Gemma-4 text config extended with handoff-probe hyperparameters.
+
+    Extra fields (all serialized to ``config.json``):
+
+    - ``probe_layer`` (`int`, defaults to 28): zero-indexed decoder layer whose
+      output feeds the probe. Must be ``< num_hidden_layers``.
+    - ``probe_feature_size`` (`int`, defaults to 1536): width of the captured
+      hidden states; must equal ``hidden_size``.
+    - ``probe_max_tokens`` (`int`, defaults to 1024): the probe scores at most
+      the first ``probe_max_tokens`` generated-token rows.
+    - ``probe_proj_dim`` (`int`, defaults to 32): width of the probe's attention
+      projection (fixed by the released checkpoint).
+    """
+
+    model_type = "gemma-4-e2b-it-hybrid"
+
+    probe_layer: int = 28
+    probe_feature_size: int = 1536
+    probe_max_tokens: int = 1024
+    probe_proj_dim: int = 32
+
+    def __post_init__(self, **kwargs):
+        super().__post_init__(**kwargs)
+        if not 0 <= self.probe_layer < self.num_hidden_layers:
+            raise ValueError(
+                f"probe_layer={self.probe_layer} must be in [0, num_hidden_layers="
+                f"{self.num_hidden_layers})"
+            )
+        # Note: probe_feature_size == hidden_size is enforced by the model, not
+        # here — `to_diff_dict()` default-constructs the config class, so this
+        # class must stay constructible with pure defaults.
+
+
+__all__ = ["Gemma4E2BItHybridConfig"]

custom_generate/generate.py

@@ -0,0 +1,178 @@
+"""Hub ``custom_generate`` override for Cactus-Compute/gemma-4-e2b-it-hybrid.
+
+Runs the stock decode loop, scores the generation with the handoff probe, and
+transports ``confidence = 1 - p_wrong`` in-band: the assistant text gains a
+final line ``\\n[[hybrid:confidence=0.7812]]`` (exactly 4 decimals, ASCII).
+
+Return contract:
+
+- default: a plain ``sequences`` tensor — required by ``transformers serve``,
+  which slices ``sequences[0, input_len:]`` on the raw return. The trailer's
+  token ids are appended to the returned sequences, and, when a ``streamer``
+  is passed, the trailer is also pushed through it, so serve emits the trailer
+  in both streaming and non-streaming modes.
+- ``return_confidence=True``: a ``GenerateWithConfidenceOutput`` with clean
+  sequences (no trailer) and the raw ``confidence`` float.
+- ``emit_trailer=False``: suppress the in-band trailer (Python API opt-out).
+
+Generations the probe contract does not cover (batch > 1, beam search,
+assisted decoding) fall back to the stock behavior: no trailer, confidence
+``None``.
+"""
+
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import torch
+
+TRAILER_TEMPLATE = "\n[[hybrid:confidence={confidence:.4f}]]"
+
+
+@dataclass
+class GenerateWithConfidenceOutput:
+    """Rich result for ``generate(..., return_confidence=True)``."""
+
+    sequences: torch.LongTensor
+    confidence: Optional[float]
+    trailer_text: Optional[str]
+    generate_output: Any  # the raw return of the stock `generate`
+
+
+class _DeferredEndStreamer:
+    """Forwards ``put`` but swallows ``end`` so the trailer can be pushed
+    through the real streamer after the stock decode loop finishes."""
+
+    def __init__(self, inner):
+        self._inner = inner
+
+    def put(self, value):
+        self._inner.put(value)
+
+    def end(self):
+        pass
+
+
+def _resolve_tokenizer(model, kwargs):
+    """Best-effort tokenizer for encoding the trailer. Never raises."""
+    tok = kwargs.get("tokenizer")
+    if tok is None:
+        tok = getattr(model, "_hybrid_trailer_tokenizer", None)
+    if tok is None:
+        try:
+            from transformers import AutoTokenizer
+
+            name_or_path = getattr(model.config, "_name_or_path", "") or ""
+            if name_or_path:
+                tok = AutoTokenizer.from_pretrained(name_or_path)
+                model._hybrid_trailer_tokenizer = tok
+        except Exception:
+            return None
+    # Unwrap processors down to the underlying tokenizer.
+    inner = getattr(tok, "tokenizer", None)
+    if inner is not None and hasattr(inner, "encode"):
+        tok = inner
+    return tok
+
+
+def _encode_trailer(tokenizer, text):
+    """Encode the trailer to token ids. Returns [] when encoding fails."""
+    if tokenizer is None:
+        return []
+    try:
+        ids = tokenizer(text, add_special_tokens=False)["input_ids"]
+    except Exception:
+        try:
+            ids = tokenizer.encode(text, add_special_tokens=False)
+        except Exception:
+            return []
+    if ids and isinstance(ids[0], list):  # batched encoding shape
+        ids = ids[0]
+    return [int(i) for i in ids]
+
+
+def generate(*args, model=None, **kwargs):
+    """Stock generation + handoff-probe confidence with an in-band trailer."""
+    if model is None:
+        raise TypeError("custom generate requires the `model` keyword argument")
+    if args:
+        if len(args) > 1:
+            raise TypeError(
+                "gemma-4-e2b-it-hybrid generate accepts at most one positional argument (`inputs`); "
+                "pass everything else as keyword arguments"
+            )
+        if kwargs.get("inputs") is not None:
+            raise TypeError("got `inputs` both positionally and as a keyword argument")
+        kwargs["inputs"] = args[0]
+
+    return_confidence = bool(kwargs.pop("return_confidence", False))
+    emit_trailer = bool(kwargs.pop("emit_trailer", True))
+    streamer = kwargs.pop("streamer", None)
+
+    # The class attribute is the stock GenerationMixin.generate; the instance
+    # attribute is this function (installed by `from_pretrained`), so calling
+    # through the class avoids recursion.
+    stock_generate = type(model).generate
+
+    scorable = kwargs.get("assistant_model") is None and hasattr(model, "probe_capture")
+    wrapped_streamer = _DeferredEndStreamer(streamer) if streamer is not None else None
+
+    try:
+        if scorable:
+            with model.probe_capture() as rows:
+                output = stock_generate(model, streamer=wrapped_streamer, **kwargs)
+            confidence = model.confidence_from_rows(rows)
+        else:
+            output = stock_generate(model, streamer=wrapped_streamer, **kwargs)
+            confidence = None
+    except BaseException:
+        if streamer is not None:
+            try:
+                streamer.end()
+            except Exception:
+                pass
+        raise
+
+    sequences = output if isinstance(output, torch.Tensor) else getattr(output, "sequences", None)
+
+    if return_confidence:
+        if streamer is not None:
+            streamer.end()
+        trailer_text = (
+            TRAILER_TEMPLATE.format(confidence=confidence) if confidence is not None else None
+        )
+        return GenerateWithConfidenceOutput(
+            sequences=sequences,
+            confidence=confidence,
+            trailer_text=trailer_text,
+            generate_output=output,
+        )
+
+    trailer_ids: list = []
+    if (
+        emit_trailer
+        and confidence is not None
+        and sequences is not None
+        and sequences.ndim == 2
+        and sequences.shape[0] == 1
+    ):
+        trailer_text = TRAILER_TEMPLATE.format(confidence=confidence)
+        trailer_ids = _encode_trailer(_resolve_tokenizer(model, kwargs), trailer_text)
+
+    if trailer_ids:
+        trailer_tensor = torch.tensor(
+            [trailer_ids], dtype=sequences.dtype, device=sequences.device
+        )
+        new_sequences = torch.cat([sequences, trailer_tensor], dim=1)
+        if isinstance(output, torch.Tensor):
+            output = new_sequences
+        else:
+            try:
+                output.sequences = new_sequences
+            except Exception:
+                pass
+        if streamer is not None:
+            streamer.put(trailer_tensor[0].cpu())
+
+    if streamer is not None:
+        streamer.end()
+    return output

custom_generate/requirements.txt

@@ -0,0 +1 @@
+transformers>=5,<6

gemma_4_e2b_it_hybrid.py

@@ -0,0 +1,219 @@
+"""Cactus-Compute/gemma-4-e2b-it-hybrid — single-file mlx-lm model (``model_file`` mechanism).
+
+Drop this file into the converted MLX repo and add ``"model_file":
+"gemma_4_e2b_it_hybrid.py"`` to its ``config.json`` (mlx-lm >= 0.30.1, the mechanism
+from mlx-lm PR #830). ``mlx_lm.load`` then builds ``Model``/``ModelArgs`` from
+this file instead of the built-in ``gemma4_text`` classes.
+
+The model is the stock ``mlx_lm.models.gemma4_text`` Gemma-4 text model plus
+the handoff probe (weight prefix ``handoff_probe.*``). During generation the
+output of decoder layer ``probe_layer`` is captured at the position that
+predicts each generated token; after generation
+
+    ``model.last_confidence``  ->  float in [0, 1] or None
+
+exposes ``confidence = 1 - p_wrong``. All probe math runs in float32.
+
+Capture semantics (tuned to ``mlx_lm.generate_step``):
+
+- multi-token (prefill-chunk) forwards reset the capture buffer and are never
+  kept; ``generate_step`` always feeds the final prompt token through a
+  1-token step, and that forward's last position is row 0 of the contract;
+- ``generate_step`` pipelines one forward ahead, so the buffer ends with one
+  lookahead row for a token that is never emitted; ``last_confidence`` drops
+  that final row. Use ``model.confidence(num_tokens=N)`` to score exactly the
+  first N generated tokens instead.
+
+Limitations:
+
+- Server-side surfacing is Python-API only for MLX today: an mlx-lm ``Model``
+  cannot inject tokens into the stream, so there is no in-band
+  ``[[hybrid:confidence=...]]`` trailer here (unlike the transformers repo).
+- Speculative decoding is not supported (draft verification feeds multiple
+  tokens per forward, which resets the buffer).
+- Prompts of exactly 2 tokens leave one extra leading row in the buffer (the
+  1-token prefill chunk is indistinguishable from a decode step). Chat-template
+  prompts are always far longer.
+"""
+
+import math
+from dataclasses import dataclass
+from typing import Any, List, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from mlx_lm.models import gemma4_text
+
+
+@dataclass
+class ModelArgs(gemma4_text.ModelArgs):
+    model_type: str = "gemma-4-e2b-it-hybrid"
+    probe_layer: int = 28
+    probe_feature_size: int = 1536
+    probe_max_tokens: int = 1024
+    probe_proj_dim: int = 32
+
+    def __post_init__(self):
+        super().__post_init__()
+        if not 0 <= self.probe_layer < self.num_hidden_layers:
+            raise ValueError(
+                f"probe_layer={self.probe_layer} must be in "
+                f"[0, num_hidden_layers={self.num_hidden_layers})"
+            )
+
+
+# Hidden widths of the probe MLP head, fixed by the released checkpoint.
+PROBE_HEAD_DIMS = (128, 64)
+
+
+class HandoffProbe(nn.Module):
+    """The released handoff probe; parameter tree matches the checkpoint keys
+    ``norm.{weight,bias}``, ``proj.{weight,bias}``, ``attn_query``,
+    ``head.{0,2,4}.{weight,bias}``."""
+
+    def __init__(self, feature_size: int, proj_dim: int = 32):
+        super().__init__()
+        h1, h2 = PROBE_HEAD_DIMS
+        self.norm = nn.LayerNorm(feature_size, eps=1e-5)
+        self.proj = nn.Linear(feature_size, proj_dim)
+        self.attn_query = mx.zeros((proj_dim,))
+        self.head = [
+            nn.Linear(proj_dim, h1),
+            nn.ReLU(),
+            nn.Linear(h1, h2),
+            nn.ReLU(),
+            nn.Linear(h2, 1),
+        ]
+
+    def p_wrong(self, hidden_states: mx.array, max_tokens: int = 1024) -> mx.array:
+        """Score ``[T, feature_size]`` rows; float32 math per the contract."""
+        x = hidden_states[:max_tokens].astype(mx.float32)
+        x = mx.fast.layer_norm(
+            x,
+            self.norm.weight.astype(mx.float32),
+            self.norm.bias.astype(mx.float32),
+            1e-5,
+        )
+        projected = mx.maximum(
+            x @ self.proj.weight.astype(mx.float32).T + self.proj.bias.astype(mx.float32),
+            0.0,
+        )
+        scores = projected @ self.attn_query.astype(mx.float32)
+        scores = scores / math.sqrt(projected.shape[-1])
+        weights = mx.softmax(scores - scores.max(), axis=-1)
+        pooled = weights @ projected
+
+        h0, h2_, h4 = self.head[0], self.head[2], self.head[4]
+        h = mx.maximum(
+            pooled @ h0.weight.astype(mx.float32).T + h0.bias.astype(mx.float32), 0.0
+        )
+        h = mx.maximum(
+            h @ h2_.weight.astype(mx.float32).T + h2_.bias.astype(mx.float32), 0.0
+        )
+        logit = h @ h4.weight.astype(mx.float32).T + h4.bias.astype(mx.float32)
+        return mx.sigmoid(logit)[0]
+
+
+class _ProbeState:
+    """Plain-object row buffer, invisible to the MLX module tree."""
+
+    def __init__(self):
+        self.rows: List[mx.array] = []
+
+    def reset(self):
+        self.rows.clear()
+
+
+class _CaptureDecoderLayer(gemma4_text.DecoderLayer):
+    """Stock decoder layer that reports its output to a capture sink.
+
+    Same submodule tree as ``DecoderLayer``, so weight keys are unchanged.
+    """
+
+    def __call__(self, x, *args, **kwargs):
+        out = super().__call__(x, *args, **kwargs)
+        sink = getattr(self, "capture_sink", None)
+        if sink is not None:
+            sink(out[0])
+        return out
+
+
+class Model(gemma4_text.Model):
+    def __init__(self, args: ModelArgs):
+        super().__init__(args)
+        self.args = args
+        self.handoff_probe = HandoffProbe(args.probe_feature_size, args.probe_proj_dim)
+        self._probe_state = _ProbeState()
+
+        capture = _CaptureDecoderLayer(args, layer_idx=args.probe_layer)
+        capture.capture_sink = self._capture_row
+        self.model.layers[args.probe_layer] = capture
+
+    # --- capture -----------------------------------------------------------
+
+    def _capture_row(self, hidden: mx.array) -> None:
+        """Keep the last position of each 1-token probe-layer forward.
+
+        ``generate_step``'s prefill loop always leaves the final prompt token
+        to a 1-token ``_step`` call, so multi-token (prefill-chunk) forwards
+        never produce a contract row — they just reset the buffer. Row 0 is
+        the step that consumes the last prompt token (= last prompt position).
+        """
+        rows = self._probe_state.rows
+        if hidden.shape[1] > 1:
+            rows.clear()
+            return
+        # Keep one row beyond the window so dropping the lookahead row cannot
+        # lose a real one.
+        if len(rows) <= self.args.probe_max_tokens:
+            rows.append(hidden[:, -1, :])
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+        input_embeddings: Optional[mx.array] = None,
+        per_layer_inputs: Optional[mx.array] = None,
+    ):
+        # A fresh cache (or none) means a new generation: reset the buffer.
+        if cache is None or (len(cache) > 0 and getattr(cache[0], "offset", 0) == 0):
+            self._probe_state.reset()
+        return super().__call__(
+            inputs,
+            cache=cache,
+            input_embeddings=input_embeddings,
+            per_layer_inputs=per_layer_inputs,
+        )
+
+    # --- scoring -----------------------------------------------------------
+
+    def reset_probe(self) -> None:
+        """Clear the capture buffer (e.g. between manual forward calls)."""
+        self._probe_state.reset()
+
+    def confidence(self, num_tokens: Optional[int] = None) -> Optional[float]:
+        """``1 - p_wrong`` for the captured generation.
+
+        ``num_tokens`` scores exactly the first N captured rows. Without it,
+        the final row is dropped to discard ``generate_step``'s one-step
+        lookahead forward. Returns ``None`` when nothing (scoreable) was
+        captured.
+        """
+        rows = list(self._probe_state.rows)
+        if not rows or rows[0].shape[0] != 1:
+            return None
+        if num_tokens is not None:
+            rows = rows[:num_tokens]
+        elif len(rows) >= 2:
+            rows = rows[:-1]
+        if not rows:
+            return None
+        stacked = mx.concatenate(rows, axis=0)  # [T, features]
+        p_wrong = self.handoff_probe.p_wrong(stacked, self.args.probe_max_tokens)
+        return float(1.0 - p_wrong.item())
+
+    @property
+    def last_confidence(self) -> Optional[float]:
+        """Confidence of the most recent generation (None before any)."""
+        return self.confidence()

generation_config.json

@@ -0,0 +1,14 @@
+{
+  "bos_token_id": 2,
+  "do_sample": true,
+  "eos_token_id": [
+    1,
+    106,
+    50
+  ],
+  "pad_token_id": 0,
+  "temperature": 1.0,
+  "top_k": 64,
+  "top_p": 0.95,
+  "transformers_version": "5.5.0.dev0"
+}

model-00001-of-00003.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3da4a9ebb8bada980bdfbc213ee28464d6ce39647f3d6dfcb2833b76e592b8d8
+size 805306504

model-00002-of-00003.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:794f283c42b320dee5824838f9d33ebcdd07cc2ce2cf23fca279dd0de362a824
+size 4697620632

model-00003-of-00003.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6b979b9af49171e398743e21a205e14ae562aca42a767c1daeba12d6599a77fa
+size 3792302306

model.safetensors.index.json

@@ -0,0 +1,618 @@
+{
+  "metadata": {
+    "total_size": 9295159114
+  },
+  "weight_map": {
+    "handoff_probe.attn_query": "model-00003-of-00003.safetensors",
+    "handoff_probe.head.0.bias": "model-00003-of-00003.safetensors",
+    "handoff_probe.head.0.weight": "model-00003-of-00003.safetensors",
+    "handoff_probe.head.2.bias": "model-00003-of-00003.safetensors",
+    "handoff_probe.head.2.weight": "model-00003-of-00003.safetensors",
+    "handoff_probe.head.4.bias": "model-00003-of-00003.safetensors",
+    "handoff_probe.head.4.weight": "model-00003-of-00003.safetensors",
+    "handoff_probe.norm.bias": "model-00003-of-00003.safetensors",
+    "handoff_probe.norm.weight": "model-00003-of-00003.safetensors",
+    "handoff_probe.proj.bias": "model-00003-of-00003.safetensors",
+    "handoff_probe.proj.weight": "model-00003-of-00003.safetensors",
+    "model.embed_tokens.weight": "model-00001-of-00003.safetensors",
+    "model.embed_tokens_per_layer.weight": "model-00002-of-00003.safetensors",
+    "model.layers.0.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.0.layer_scalar": "model-00003-of-00003.safetensors",
+    "model.layers.0.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.0.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.0.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.0.per_layer_input_gate.weight": "model-00003-of-00003.safetensors",
+    "model.layers.0.per_layer_projection.weight": "model-00003-of-00003.safetensors",
+    "model.layers.0.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.0.post_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.0.post_per_layer_input_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.0.pre_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.0.self_attn.k_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.0.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.0.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.0.self_attn.q_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.0.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.0.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.1.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.1.layer_scalar": "model-00003-of-00003.safetensors",
+    "model.layers.1.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.1.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.1.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.1.per_layer_input_gate.weight": "model-00003-of-00003.safetensors",
+    "model.layers.1.per_layer_projection.weight": "model-00003-of-00003.safetensors",
+    "model.layers.1.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.1.post_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.1.post_per_layer_input_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.1.pre_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.1.self_attn.k_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.1.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.1.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.1.self_attn.q_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.1.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.1.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.10.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.10.layer_scalar": "model-00003-of-00003.safetensors",
+    "model.layers.10.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.10.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.10.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.10.per_layer_input_gate.weight": "model-00003-of-00003.safetensors",
+    "model.layers.10.per_layer_projection.weight": "model-00003-of-00003.safetensors",
+    "model.layers.10.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.10.post_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.10.post_per_layer_input_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.10.pre_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.10.self_attn.k_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.10.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.10.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.10.self_attn.q_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.10.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.10.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.11.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.11.layer_scalar": "model-00003-of-00003.safetensors",
+    "model.layers.11.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.11.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.11.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.11.per_layer_input_gate.weight": "model-00003-of-00003.safetensors",
+    "model.layers.11.per_layer_projection.weight": "model-00003-of-00003.safetensors",
+    "model.layers.11.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.11.post_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.11.post_per_layer_input_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.11.pre_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.11.self_attn.k_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.11.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.11.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.11.self_attn.q_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.11.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.11.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.12.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.12.layer_scalar": "model-00003-of-00003.safetensors",
+    "model.layers.12.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.12.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.12.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.12.per_layer_input_gate.weight": "model-00003-of-00003.safetensors",
+    "model.layers.12.per_layer_projection.weight": "model-00003-of-00003.safetensors",
+    "model.layers.12.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.12.post_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.12.post_per_layer_input_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.12.pre_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.12.self_attn.k_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.12.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.12.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.12.self_attn.q_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.12.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.12.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.13.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.13.layer_scalar": "model-00003-of-00003.safetensors",
+    "model.layers.13.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.13.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.13.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.13.per_layer_input_gate.weight": "model-00003-of-00003.safetensors",
+    "model.layers.13.per_layer_projection.weight": "model-00003-of-00003.safetensors",
+    "model.layers.13.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.13.post_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.13.post_per_layer_input_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.13.pre_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.13.self_attn.k_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.13.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.13.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.13.self_attn.q_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.13.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.13.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.14.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.14.layer_scalar": "model-00003-of-00003.safetensors",
+    "model.layers.14.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.14.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.14.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.14.per_layer_input_gate.weight": "model-00003-of-00003.safetensors",
+    "model.layers.14.per_layer_projection.weight": "model-00003-of-00003.safetensors",
+    "model.layers.14.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.14.post_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.14.post_per_layer_input_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.14.pre_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.14.self_attn.k_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.14.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.14.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.14.self_attn.q_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.14.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.14.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.15.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.15.layer_scalar": "model-00003-of-00003.safetensors",
+    "model.layers.15.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.15.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.15.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.15.per_layer_input_gate.weight": "model-00003-of-00003.safetensors",
+    "model.layers.15.per_layer_projection.weight": "model-00003-of-00003.safetensors",
+    "model.layers.15.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.15.post_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.15.post_per_layer_input_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.15.pre_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.15.self_attn.k_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.15.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.15.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.15.self_attn.q_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.15.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.15.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.16.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.16.layer_scalar": "model-00003-of-00003.safetensors",
+    "model.layers.16.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.16.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.16.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.16.per_layer_input_gate.weight": "model-00003-of-00003.safetensors",
+    "model.layers.16.per_layer_projection.weight": "model-00003-of-00003.safetensors",
+    "model.layers.16.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.16.post_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.16.post_per_layer_input_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.16.pre_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.16.self_attn.k_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.16.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.16.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.16.self_attn.q_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.16.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.16.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.17.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.17.layer_scalar": "model-00003-of-00003.safetensors",
+    "model.layers.17.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.17.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.17.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.17.per_layer_input_gate.weight": "model-00003-of-00003.safetensors",
+    "model.layers.17.per_layer_projection.weight": "model-00003-of-00003.safetensors",
+    "model.layers.17.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.17.post_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.17.post_per_layer_input_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.17.pre_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.17.self_attn.k_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.17.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.17.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.17.self_attn.q_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.17.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.17.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.18.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.18.layer_scalar": "model-00003-of-00003.safetensors",
+    "model.layers.18.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.18.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.18.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.18.per_layer_input_gate.weight": "model-00003-of-00003.safetensors",
+    "model.layers.18.per_layer_projection.weight": "model-00003-of-00003.safetensors",
+    "model.layers.18.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.18.post_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.18.post_per_layer_input_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.18.pre_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.18.self_attn.k_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.18.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.18.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.18.self_attn.q_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.18.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.18.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.19.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.19.layer_scalar": "model-00003-of-00003.safetensors",
+    "model.layers.19.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.19.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.19.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.19.per_layer_input_gate.weight": "model-00003-of-00003.safetensors",
+    "model.layers.19.per_layer_projection.weight": "model-00003-of-00003.safetensors",
+    "model.layers.19.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.19.post_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.19.post_per_layer_input_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.19.pre_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.19.self_attn.k_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.19.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.19.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.19.self_attn.q_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.19.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.19.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.2.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.2.layer_scalar": "model-00003-of-00003.safetensors",
+    "model.layers.2.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.2.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.2.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.2.per_layer_input_gate.weight": "model-00003-of-00003.safetensors",
+    "model.layers.2.per_layer_projection.weight": "model-00003-of-00003.safetensors",
+    "model.layers.2.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.2.post_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.2.post_per_layer_input_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.2.pre_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.2.self_attn.k_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.2.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.2.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.2.self_attn.q_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.2.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.2.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.20.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.20.layer_scalar": "model-00003-of-00003.safetensors",
+    "model.layers.20.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.20.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.20.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.20.per_layer_input_gate.weight": "model-00003-of-00003.safetensors",
+    "model.layers.20.per_layer_projection.weight": "model-00003-of-00003.safetensors",
+    "model.layers.20.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.20.post_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.20.post_per_layer_input_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.20.pre_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.20.self_attn.k_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.20.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.20.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.20.self_attn.q_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.20.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.20.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.21.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.21.layer_scalar": "model-00003-of-00003.safetensors",
+    "model.layers.21.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.21.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.21.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.21.per_layer_input_gate.weight": "model-00003-of-00003.safetensors",
+    "model.layers.21.per_layer_projection.weight": "model-00003-of-00003.safetensors",
+    "model.layers.21.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.21.post_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.21.post_per_layer_input_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.21.pre_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.21.self_attn.k_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.21.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.21.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.21.self_attn.q_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.21.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.21.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.22.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.22.layer_scalar": "model-00003-of-00003.safetensors",
+    "model.layers.22.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.22.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.22.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.22.per_layer_input_gate.weight": "model-00003-of-00003.safetensors",
+    "model.layers.22.per_layer_projection.weight": "model-00003-of-00003.safetensors",
+    "model.layers.22.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.22.post_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.22.post_per_layer_input_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.22.pre_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.22.self_attn.k_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.22.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.22.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.22.self_attn.q_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.22.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.22.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.23.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.23.layer_scalar": "model-00003-of-00003.safetensors",
+    "model.layers.23.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.23.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.23.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.23.per_layer_input_gate.weight": "model-00003-of-00003.safetensors",
+    "model.layers.23.per_layer_projection.weight": "model-00003-of-00003.safetensors",
+    "model.layers.23.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.23.post_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.23.post_per_layer_input_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.23.pre_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.23.self_attn.k_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.23.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.23.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.23.self_attn.q_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.23.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.23.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.24.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.24.layer_scalar": "model-00003-of-00003.safetensors",
+    "model.layers.24.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.24.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.24.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.24.per_layer_input_gate.weight": "model-00003-of-00003.safetensors",
+    "model.layers.24.per_layer_projection.weight": "model-00003-of-00003.safetensors",
+    "model.layers.24.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.24.post_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.24.post_per_layer_input_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.24.pre_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.24.self_attn.k_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.24.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.24.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.24.self_attn.q_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.24.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.24.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.25.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.25.layer_scalar": "model-00003-of-00003.safetensors",
+    "model.layers.25.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.25.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.25.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.25.per_layer_input_gate.weight": "model-00003-of-00003.safetensors",
+    "model.layers.25.per_layer_projection.weight": "model-00003-of-00003.safetensors",
+    "model.layers.25.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.25.post_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.25.post_per_layer_input_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.25.pre_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.25.self_attn.k_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.25.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.25.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.25.self_attn.q_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.25.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.25.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.26.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.26.layer_scalar": "model-00003-of-00003.safetensors",
+    "model.layers.26.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.26.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.26.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.26.per_layer_input_gate.weight": "model-00003-of-00003.safetensors",
+    "model.layers.26.per_layer_projection.weight": "model-00003-of-00003.safetensors",
+    "model.layers.26.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.26.post_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.26.post_per_layer_input_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.26.pre_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.26.self_attn.k_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.26.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.26.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.26.self_attn.q_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.26.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.26.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.layer_scalar": "model-00003-of-00003.safetensors",
+    "model.layers.27.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.per_layer_input_gate.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.per_layer_projection.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.post_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.post_per_layer_input_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.pre_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.self_attn.k_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.self_attn.q_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.28.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.28.layer_scalar": "model-00003-of-00003.safetensors",
+    "model.layers.28.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.28.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.28.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.28.per_layer_input_gate.weight": "model-00003-of-00003.safetensors",
+    "model.layers.28.per_layer_projection.weight": "model-00003-of-00003.safetensors",
+    "model.layers.28.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.28.post_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.28.post_per_layer_input_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.28.pre_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.28.self_attn.k_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.28.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.28.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.28.self_attn.q_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.28.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.28.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.29.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.29.layer_scalar": "model-00003-of-00003.safetensors",
+    "model.layers.29.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.29.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.29.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.29.per_layer_input_gate.weight": "model-00003-of-00003.safetensors",
+    "model.layers.29.per_layer_projection.weight": "model-00003-of-00003.safetensors",
+    "model.layers.29.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.29.post_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.29.post_per_layer_input_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.29.pre_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.29.self_attn.k_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.29.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.29.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.29.self_attn.q_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.29.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.29.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.3.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.3.layer_scalar": "model-00003-of-00003.safetensors",
+    "model.layers.3.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.3.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.3.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.3.per_layer_input_gate.weight": "model-00003-of-00003.safetensors",
+    "model.layers.3.per_layer_projection.weight": "model-00003-of-00003.safetensors",
+    "model.layers.3.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.3.post_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.3.post_per_layer_input_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.3.pre_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.3.self_attn.k_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.3.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.3.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.3.self_attn.q_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.3.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.3.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.30.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.30.layer_scalar": "model-00003-of-00003.safetensors",
+    "model.layers.30.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.30.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.30.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.30.per_layer_input_gate.weight": "model-00003-of-00003.safetensors",
+    "model.layers.30.per_layer_projection.weight": "model-00003-of-00003.safetensors",
+    "model.layers.30.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.30.post_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.30.post_per_layer_input_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.30.pre_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.30.self_attn.k_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.30.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.30.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.30.self_attn.q_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.30.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.30.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.31.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.31.layer_scalar": "model-00003-of-00003.safetensors",
+    "model.layers.31.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.31.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.31.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.31.per_layer_input_gate.weight": "model-00003-of-00003.safetensors",
+    "model.layers.31.per_layer_projection.weight": "model-00003-of-00003.safetensors",
+    "model.layers.31.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.31.post_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.31.post_per_layer_input_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.31.pre_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.31.self_attn.k_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.31.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.31.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.31.self_attn.q_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.31.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.31.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.32.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.32.layer_scalar": "model-00003-of-00003.safetensors",
+    "model.layers.32.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.32.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.32.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.32.per_layer_input_gate.weight": "model-00003-of-00003.safetensors",
+    "model.layers.32.per_layer_projection.weight": "model-00003-of-00003.safetensors",
+    "model.layers.32.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.32.post_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.32.post_per_layer_input_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.32.pre_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.32.self_attn.k_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.32.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.32.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.32.self_attn.q_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.32.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.32.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.33.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.33.layer_scalar": "model-00003-of-00003.safetensors",
+    "model.layers.33.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.33.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.33.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.33.per_layer_input_gate.weight": "model-00003-of-00003.safetensors",
+    "model.layers.33.per_layer_projection.weight": "model-00003-of-00003.safetensors",
+    "model.layers.33.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.33.post_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.33.post_per_layer_input_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.33.pre_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.33.self_attn.k_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.33.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.33.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.33.self_attn.q_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.33.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.33.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.34.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.34.layer_scalar": "model-00003-of-00003.safetensors",
+    "model.layers.34.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.34.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.34.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.34.per_layer_input_gate.weight": "model-00003-of-00003.safetensors",
+    "model.layers.34.per_layer_projection.weight": "model-00003-of-00003.safetensors",
+    "model.layers.34.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.34.post_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.34.post_per_layer_input_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.34.pre_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.34.self_attn.k_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.34.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.34.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.34.self_attn.q_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.34.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.34.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.4.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.4.layer_scalar": "model-00003-of-00003.safetensors",
+    "model.layers.4.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.4.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.4.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.4.per_layer_input_gate.weight": "model-00003-of-00003.safetensors",
+    "model.layers.4.per_layer_projection.weight": "model-00003-of-00003.safetensors",
+    "model.layers.4.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.4.post_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.4.post_per_layer_input_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.4.pre_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.4.self_attn.k_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.4.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.4.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.4.self_attn.q_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.4.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.4.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.5.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.5.layer_scalar": "model-00003-of-00003.safetensors",
+    "model.layers.5.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.5.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.5.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.5.per_layer_input_gate.weight": "model-00003-of-00003.safetensors",
+    "model.layers.5.per_layer_projection.weight": "model-00003-of-00003.safetensors",
+    "model.layers.5.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.5.post_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.5.post_per_layer_input_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.5.pre_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.5.self_attn.k_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.5.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.5.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.5.self_attn.q_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.5.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.5.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.6.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.6.layer_scalar": "model-00003-of-00003.safetensors",
+    "model.layers.6.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.6.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.6.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.6.per_layer_input_gate.weight": "model-00003-of-00003.safetensors",
+    "model.layers.6.per_layer_projection.weight": "model-00003-of-00003.safetensors",
+    "model.layers.6.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.6.post_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.6.post_per_layer_input_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.6.pre_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.6.self_attn.k_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.6.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.6.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.6.self_attn.q_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.6.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.6.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.7.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.7.layer_scalar": "model-00003-of-00003.safetensors",
+    "model.layers.7.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.7.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.7.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.7.per_layer_input_gate.weight": "model-00003-of-00003.safetensors",
+    "model.layers.7.per_layer_projection.weight": "model-00003-of-00003.safetensors",
+    "model.layers.7.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.7.post_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.7.post_per_layer_input_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.7.pre_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.7.self_attn.k_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.7.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.7.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.7.self_attn.q_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.7.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.7.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.8.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.8.layer_scalar": "model-00003-of-00003.safetensors",
+    "model.layers.8.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.8.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.8.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.8.per_layer_input_gate.weight": "model-00003-of-00003.safetensors",
+    "model.layers.8.per_layer_projection.weight": "model-00003-of-00003.safetensors",
+    "model.layers.8.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.8.post_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.8.post_per_layer_input_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.8.pre_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.8.self_attn.k_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.8.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.8.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.8.self_attn.q_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.8.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.8.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.9.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.9.layer_scalar": "model-00003-of-00003.safetensors",
+    "model.layers.9.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.9.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.9.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.9.per_layer_input_gate.weight": "model-00003-of-00003.safetensors",
+    "model.layers.9.per_layer_projection.weight": "model-00003-of-00003.safetensors",
+    "model.layers.9.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.9.post_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.9.post_per_layer_input_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.9.pre_feedforward_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.9.self_attn.k_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.9.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.9.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.9.self_attn.q_norm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.9.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.9.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.norm.weight": "model-00003-of-00003.safetensors",
+    "model.per_layer_model_projection.weight": "model-00003-of-00003.safetensors",
+    "model.per_layer_projection_norm.weight": "model-00003-of-00003.safetensors"
+  }
+}

modeling_gemma_4_e2b_it_hybrid.py

@@ -0,0 +1,165 @@
+"""Cactus-Compute/gemma-4-e2b-it-hybrid — Gemma-4 causal LM with a handoff probe.
+
+``Gemma4E2BItHybridForCausalLM`` is the stock ``Gemma4ForCausalLM`` plus a small
+"handoff probe" head (weight prefix ``handoff_probe.*``) that scores each
+generation with ``confidence = 1 - p_wrong``. Base weights keep identical keys,
+so the checkpoint is the stock checkpoint with eleven extra probe tensors.
+
+Probe contract (checkpoint layer 28, float32 math):
+
+- input: ``[T, 1536]`` — output of decoder layer index ``config.probe_layer``
+  at the position that predicts each generated token (row 0 = last prompt
+  position at prefill, row t = position captured at generation step t). Only
+  the first ``config.probe_max_tokens`` rows are used.
+- ``x = LayerNorm(x, eps=1e-5) * norm.weight + norm.bias``
+- ``p = relu(x @ proj.weight.T + proj.bias)``
+- ``s = p @ attn_query / sqrt(probe_proj_dim)``; ``w = softmax_T(s)``
+- ``pooled = w @ p``
+- ``h = relu(head.0 @ pooled); h = relu(head.2 @ h); logit = head.4 @ h``
+- ``p_wrong = sigmoid(logit)``; ``confidence = 1 - p_wrong``
+
+Layer capture uses a forward hook that keeps only the last position of each
+decode step (a ``[batch, hidden]`` row), never the full hidden-state stack.
+"""
+
+import math
+from contextlib import contextmanager
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from transformers.generation.utils import GenerationMixin
+from transformers.models.gemma4.modeling_gemma4 import Gemma4ForCausalLM
+
+try:
+    from .configuration_gemma_4_e2b_it_hybrid import Gemma4E2BItHybridConfig
+except ImportError:  # direct (non-package) execution
+    from configuration_gemma_4_e2b_it_hybrid import Gemma4E2BItHybridConfig
+
+# Hidden widths of the probe MLP head, fixed by the released checkpoint.
+PROBE_HEAD_DIMS = (128, 64)
+
+
+class HandoffProbe(nn.Module):
+    """The released handoff probe. Weight keys match the probe checkpoint:
+
+    ``norm.{weight,bias}``, ``proj.{weight,bias}``, ``attn_query``,
+    ``head.{0,2,4}.{weight,bias}``.
+    """
+
+    def __init__(self, feature_size: int, proj_dim: int = 32) -> None:
+        super().__init__()
+        h1, h2 = PROBE_HEAD_DIMS
+        self.norm = nn.LayerNorm(feature_size, eps=1e-5)
+        self.proj = nn.Linear(feature_size, proj_dim)
+        self.attn_query = nn.Parameter(torch.zeros(proj_dim))
+        self.head = nn.Sequential(
+            nn.Linear(proj_dim, h1),
+            nn.ReLU(),
+            nn.Linear(h1, h2),
+            nn.ReLU(),
+            nn.Linear(h2, 1),
+        )
+
+    @torch.no_grad()
+    def p_wrong(self, hidden_states: torch.Tensor, max_tokens: int = 1024) -> float:
+        """Score ``[T, feature_size]`` generated-token hidden states.
+
+        All math runs in float32 on the probe's own device, regardless of the
+        dtype the module weights were loaded in; the input may arrive on any
+        device (capture hooks store rows on CPU).
+        """
+        if hidden_states.ndim != 2:
+            raise ValueError(f"expected [tokens, features], got {tuple(hidden_states.shape)}")
+        if hidden_states.shape[0] == 0:
+            raise ValueError("cannot score an empty generation")
+        x = hidden_states[:max_tokens].to(self.norm.weight.device, torch.float32)
+
+        x = F.layer_norm(
+            x, (x.shape[-1],), self.norm.weight.float(), self.norm.bias.float(), self.norm.eps
+        )
+        projected = F.relu(F.linear(x, self.proj.weight.float(), self.proj.bias.float()))
+        scores = projected @ self.attn_query.float() / math.sqrt(projected.shape[-1])
+        weights = torch.softmax(scores, dim=0)
+        pooled = weights @ projected
+
+        h = F.relu(F.linear(pooled, self.head[0].weight.float(), self.head[0].bias.float()))
+        h = F.relu(F.linear(h, self.head[2].weight.float(), self.head[2].bias.float()))
+        logit = F.linear(h, self.head[4].weight.float(), self.head[4].bias.float())
+        return float(torch.sigmoid(logit)[0].item())
+
+
+class Gemma4E2BItHybridForCausalLM(Gemma4ForCausalLM):
+    """Stock Gemma-4 causal LM plus the ``handoff_probe.*`` scoring head."""
+
+    config_class = Gemma4E2BItHybridConfig
+    config: Gemma4E2BItHybridConfig
+
+    def __init__(self, config: Gemma4E2BItHybridConfig) -> None:
+        if config.probe_feature_size != config.hidden_size:
+            raise ValueError(
+                f"probe_feature_size={config.probe_feature_size} must equal "
+                f"hidden_size={config.hidden_size}"
+            )
+        super().__init__(config)
+        self.handoff_probe = HandoffProbe(
+            config.probe_feature_size, getattr(config, "probe_proj_dim", 32)
+        )
+        #: Confidence of the most recent scored generation (``None`` before any).
+        self.last_confidence: float | None = None
+
+    @contextmanager
+    def probe_capture(self):
+        """Capture probe-layer rows during generation.
+
+        Yields a list that fills with one ``[batch, hidden]`` float32 CPU tensor
+        per decode step (row 0 comes from the prefill forward at the last prompt
+        position). At most ``config.probe_max_tokens`` rows are kept.
+        """
+        rows: list[torch.Tensor] = []
+        layer = self.model.layers[self.config.probe_layer]
+        max_rows = self.config.probe_max_tokens
+
+        def hook(module, args, output):
+            if len(rows) < max_rows:
+                hidden = output[0] if isinstance(output, tuple) else output
+                rows.append(hidden[:, -1, :].detach().to(torch.float32).cpu())
+
+        handle = layer.register_forward_hook(hook)
+        try:
+            yield rows
+        finally:
+            handle.remove()
+
+    def confidence_from_rows(self, rows: list[torch.Tensor]) -> float | None:
+        """Turn captured probe rows into ``confidence = 1 - p_wrong``.
+
+        Returns ``None`` when the capture is empty or not a single-sequence
+        generation (batch size > 1, beam search, ...), for which the probe
+        contract is undefined.
+        """
+        if not rows or any(row.shape[0] != 1 for row in rows):
+            return None
+        states = torch.cat(rows, dim=0)  # [T, hidden]
+        p_wrong = self.handoff_probe.p_wrong(states, self.config.probe_max_tokens)
+        self.last_confidence = 1.0 - p_wrong
+        return self.last_confidence
+
+    def generate_with_confidence(self, *args, **kwargs):
+        """Run the stock ``generate`` and score it with the handoff probe.
+
+        Returns ``(sequences, confidence)`` where ``sequences`` is exactly what
+        the stock ``generate`` returns for the given arguments (no in-band
+        trailer is appended) and ``confidence`` is a float in ``[0, 1]``, or
+        ``None`` when the generation cannot be scored (batch > 1, beams, or
+        assisted decoding).
+        """
+        kwargs.pop("return_confidence", None)
+        with self.probe_capture() as rows:
+            sequences = GenerationMixin.generate(self, *args, **kwargs)
+        confidence = self.confidence_from_rows(rows)
+        return sequences, confidence
+
+
+__all__ = ["Gemma4E2BItHybridForCausalLM", "HandoffProbe"]

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cc8d3a0ce36466ccc1278bf987df5f71db1719b9ca6b4118264f45cb627bfe0f
+size 32169626

tokenizer_config.json

@@ -0,0 +1,120 @@
+{
+  "audio_token": "<|audio|>",
+  "backend": "tokenizers",
+  "boa_token": "<|audio>",
+  "boi_token": "<|image>",
+  "bos_token": "<bos>",
+  "eoa_token": "<audio|>",
+  "eoc_token": "<channel|>",
+  "eoi_token": "<image|>",
+  "eos_token": "<eos>",
+  "eot_token": "<turn|>",
+  "escape_token": "<|\"|>",
+  "etc_token": "<tool_call|>",
+  "etd_token": "<tool|>",
+  "etr_token": "<tool_response|>",
+  "extra_special_tokens": [
+    "<|video|>"
+  ],
+  "image_token": "<|image|>",
+  "mask_token": "<mask>",
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<pad>",
+  "padding_side": "left",
+  "processor_class": "Gemma4Processor",
+  "response_schema": {
+    "type": "object",
+    "properties": {
+      "role": {
+        "const": "assistant"
+      },
+      "thinking": {
+        "type": "string"
+      },
+      "content": {
+        "type": "string"
+      },
+      "tool_calls": {
+        "x-regex-iterator": "<\\|tool_call>(.*?)<tool_call\\|>",
+        "type": "array",
+        "items": {
+          "type": "object",
+          "properties": {
+            "type": {
+              "const": "function"
+            },
+            "function": {
+              "type": "object",
+              "x-regex": "call\\:(?P<name>\\w+)(?P<arguments>\\{.*\\})",
+              "properties": {
+                "name": {
+                  "type": "string"
+                },
+                "arguments": {
+                  "type": "object",
+                  "x-parser": "gemma4-tool-call",
+                  "additionalProperties": {}
+                }
+              }
+            }
+          }
+        }
+      }
+    },
+    "x-regex": "(\\<\\|channel\\>thought\\n(?P<thinking>.*?)\\<channel\\|\\>)?(?P<tool_calls>\\<\\|tool_call\\>.*\\<tool_call\\|\\>)?(?P<content>(?:(?!\\<turn\\|\\>)(?!\\<\\|tool_response\\>).)+)?(?:\\<turn\\|\\>|\\<\\|tool_response\\>)?"
+  },
+  "response_template": {
+    "defaults": {
+      "role": "assistant"
+    },
+    "fields": {
+      "content": {
+        "close": [
+          "<turn|>",
+          "<|tool_response>",
+          "<eos>"
+        ],
+        "content": "text"
+      },
+      "thinking": {
+        "close": "<channel|>",
+        "content": "text",
+        "open": "<|channel>thought\n"
+      },
+      "tool_calls": {
+        "close": "<tool_call|>",
+        "content": "json",
+        "content_args": {
+          "string_delims": [
+            [
+              "<|\"|>",
+              "<|\"|>"
+            ]
+          ],
+          "unquoted_keys": true
+        },
+        "open_pattern": "<\\|tool_call>call:(?P<name>\\w+)",
+        "repeats": true,
+        "transform": {
+          "function": {
+            "arguments": "{content}",
+            "name": "{name}"
+          },
+          "type": "function"
+        }
+      }
+    },
+    "start_anchor": [
+      "<|turn>model\n",
+      "<tool_response|>"
+    ]
+  },
+  "soc_token": "<|channel>",
+  "sot_token": "<|turn>",
+  "stc_token": "<|tool_call>",
+  "std_token": "<|tool>",
+  "str_token": "<|tool_response>",
+  "think_token": "<|think|>",
+  "tokenizer_class": "GemmaTokenizer",
+  "unk_token": "<unk>"
+}