feat: merge LEK into lemrd weights

LEK-2 LoRA merged into Gemma 4 31B Dense attention projections.
Converged at loss 0.0002 in 278 steps via patience-stop (best at 248).
Gradient checkpointing enabled. KV-shared layers restored from base.

Co-Authored-By: Virgil <virgil@lethean.io>

README.md

@@ -1,9 +1,7 @@
 ---
-library_name: mlx
-license: apache-2.0
-license_link: https://ai.google.dev/gemma/docs/gemma_4_license
-pipeline_tag: text-generation
-base_model: google/gemma-4-31b-it
+language: en
 tags:
 - mlx
+pipeline_tag: image-text-to-text
+library_name: mlx
 ---

chat_template.jinja

@@ -11,15 +11,34 @@
                 description:<|"|>{{ value['description'] }}<|"|>
                 {%- set add_comma = true -%}
             {%- endif -%}
+            {%- if value['nullable'] %}
+                {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                nullable: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 == 'OBJECT' -%}
+                ,properties:{
+                {%- if value['properties'] is defined and value['properties'] is mapping -%}
+                    {{- format_parameters(value['properties'], value['required'] | default([])) -}}
+                {%- elif value is mapping -%}
+                    {{- format_parameters(value, value['required'] | default([])) -}}
+                {%- endif -%}
+                }
+                {%- if value['required'] -%}
+                    ,required:[
+                    {%- for item in value['required'] | default([]) -%}
+                        <|"|>{{- item -}}<|"|>
+                        {%- if not loop.last %},{% endif -%}
+                    {%- endfor -%}
+                    ]
+                {%- endif -%}
             {%- elif value['type'] | upper == 'ARRAY' -%}
                 {%- if value['items'] is mapping and value['items'] -%}
-                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
-                    items:{
+                    ,items:{
                     {%- set ns_items = namespace(found_first=false) -%}
                     {%- for item_key, item_value in value['items'] | dictsort -%}
                         {%- if item_value is not none -%}
@@ -52,32 +71,6 @@
                     }
                 {%- 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([])) -}}
-                    }
-                {%- 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 -%}
@@ -157,31 +150,16 @@
     {{- 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 -%}
-
 {%- set ns = namespace(prev_message_type=None) -%}
 {%- set loop_messages = messages -%}
-{{- bos_token -}}
+{{ bos_token }}
 {#- Handle System/Tool Definitions Block -#}
 {%- if (enable_thinking is defined and enable_thinking) or tools or messages[0]['role'] in ['system', 'developer'] -%}
     {{- '<|turn>system\n' -}}
 
     {#- Inject Thinking token at the very top of the FIRST system turn -#}
     {%- if enable_thinking is defined and enable_thinking -%}
-        {{- '<|think|>\n' -}}
+        {{- '<|think|>' -}}
         {%- set ns.prev_message_type = 'think' -%}
     {%- endif -%}
 
@@ -202,41 +180,11 @@
     {{- '<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: suppress duplicate <|turn>model when previous non-tool message was also assistant -#}
-    {%- set prev_nt = namespace(role=None, found=false) -%}
-    {%- if loop.index0 > 0 -%}
-        {%- for j in range(loop.index0 - 1, -1, -1) -%}
-            {%- if not prev_nt.found -%}
-                {%- if loop_messages[j]['role'] != 'tool' -%}
-                    {%- set prev_nt.role = loop_messages[j]['role'] -%}
-                    {%- set prev_nt.found = true -%}
-                {%- endif -%}
-            {%- endif -%}
-        {%- endfor -%}
-    {%- endif -%}
-    {%- set continue_same_model_turn = (role == 'model' and prev_nt.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') -%}
-    {%- if thinking_text and loop.index0 > ns_turn.last_user_idx and message.get('tool_calls') -%}
-        {{- '<|channel>thought\n' + thinking_text + '\n<channel|>' -}}
-    {%- endif -%}
 
             {%- if message['tool_calls'] -%}
                 {%- for tool_call in message['tool_calls'] -%}
@@ -257,49 +205,23 @@
                 {%- 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) -#}
+            {%- if message['tool_responses'] -%}
+                {#- Tool Response handling -#}
                 {%- for tool_response in message['tool_responses'] -%}
-                    {{- format_tool_response_block(tool_response['name'] | default('unknown'), 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') | default('unknown')) -%}
-                        {%- for tc in message['tool_calls'] -%}
-                            {%- if tc.get('id') == follow.get('tool_call_id') -%}
-                                {%- set ns_tname.name = tc['function']['name'] -%}
-                            {%- endif -%}
+                    {{- '<|tool_response>' -}}
+                    {%- if tool_response['response'] is mapping -%}
+                        {{- 'response:' + tool_response['name'] | default('unknown') + '{' -}}
+                        {%- for key, value in tool_response['response'] | dictsort -%}
+                            {{- key -}}:{{- format_argument(value, escape_keys=False) -}}
+                            {%- if not loop.last %},{% 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) -}}
-                        {%- else -%}
-                            {{- format_tool_response_block(ns_tname.name, tool_body) -}}
-                        {%- endif -%}
-                        {%- set ns_tr_out.flag = true -%}
-                        {%- set ns.prev_message_type = 'tool_response' -%}
+                        {{- '}' -}}
+                    {%- else -%}
+                        {{- 'response:' + tool_response['name'] | default('unknown') + '{value:' + format_argument(tool_response['response'], escape_keys=False) + '}' -}}
                     {%- endif -%}
+                    {{- '<tool_response|>' -}}
                 {%- endfor -%}
+                {%- set ns.prev_message_type = 'tool_response' -%}
             {%- endif -%}
 
             {%- if message['content'] is string -%}
@@ -317,31 +239,28 @@
                             {{- item['text'] | trim -}}
                         {%- endif -%}
                     {%- elif item['type'] == 'image' -%}
-                        {{- '<|image|>' -}}
+                        {{- '\n\n<|image|>\n\n' -}}
                         {%- set ns.prev_message_type = 'image' -%}
                     {%- elif item['type'] == 'audio' -%}
                         {{- '<|audio|>' -}}
                         {%- set ns.prev_message_type = 'audio' -%}
                     {%- elif item['type'] == 'video' -%}
-                        {{- '<|video|>' -}}
+                        {{- '\n\n<|video|>\n\n' -}}
                         {%- set ns.prev_message_type = 'video' -%}
                     {%- endif -%}
                 {%- endfor -%}
             {%- endif -%}
 
-        {%- if ns.prev_message_type == 'tool_call' and not ns_tr_out.flag -%}
-            {{- '<|tool_response>' -}}
-        {%- elif not (ns_tr_out.flag and not message.get('content')) -%}
+        {%- if not (message['tool_responses'] and not message['content']) -%}
             {{- '<turn|>\n' -}}
         {%- endif -%}
-    {%- endif -%}
 {%- endfor -%}
 
 {%- if add_generation_prompt -%}
-    {%- if ns.prev_message_type != 'tool_response' and ns.prev_message_type != 'tool_call' -%}
+    {%- if ns.prev_message_type != 'tool_response' -%}
         {{- '<|turn>model\n' -}}
-        {%- if not enable_thinking | default(false) -%}
-            {{- '<|channel>thought\n<channel|>' -}}
-        {%- endif -%}
+    {%- endif -%}
+    {%- if not enable_thinking | default(false) -%}
+        {{- '<|channel>thought\n<channel|>' -}}
     {%- endif -%}
 {%- endif -%}

model-00001-of-00004.safetensors → model-00001-of-00006.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:988e2b1fd41d93b62b8c432f52f632c43b8cb7f86df4b957db36a3cc0dab40ca
-size 5366617512
+oid sha256:6a93ebff49222bf837fec630cca73782ff2662da0db93e74549fa743085d36dd
+size 5326590887

model-00002-of-00004.safetensors → model-00002-of-00006.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a496a96fbd39cd11a9871f91d026013d91069dcac15f89ca861b93976f3857cf
-size 5361642573
+oid sha256:b60c6d7570b9ceb1bcfd22091888121311882505d5bf7758db3fabe769f2d5bb
+size 5351454850

model-00003-of-00004.safetensors → model-00003-of-00006.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:afa555ff0e1bc458c5b08aeef1f4499dce63e2bfb5d3a2aac716e47c0a5672c1
-size 5367276094
+oid sha256:0f889e77a2debd1346a4df0661e3daa0ff8c12fe0332a29f54d1d71e2ebf8a00
+size 5363842431

model-00004-of-00004.safetensors → model-00004-of-00006.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:0bde02793a3d2ed3f29c3d1629ea9fbeb9c25720bc42db0eca9720abe094235a
-size 1173848301
+oid sha256:2b367bddf857b2a78f72d2c079d331e9a4d62354bfb05fc4b83cb20a724c3a5c
+size 5363842477

model-00005-of-00006.safetensors

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

model-00006-of-00006.safetensors

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

processor_config.json

@@ -28,5 +28,15 @@
     }
   },
   "image_seq_length": 280,
-  "processor_class": "Gemma4Processor"
-}
+  "processor_class": "Gemma4Processor",
+  "feature_extractor": {
+    "feature_extractor_type": "Gemma4AudioFeatureExtractor",
+    "sampling_rate": 16000,
+    "num_mel_filters": 128,
+    "fft_length": 512,
+    "hop_length": 160,
+    "chunk_duration": 8.0,
+    "overlap_duration": 1.0
+  },
+  "audio_ms_per_token": 40
+}

tokenizer_config.json

@@ -13,6 +13,9 @@
   "etc_token": "<tool_call|>",
   "etd_token": "<tool|>",
   "etr_token": "<tool_response|>",
+  "extra_special_tokens": [
+    "<|video|>"
+  ],
   "image_token": "<|image|>",
   "is_local": true,
   "mask_token": "<mask>",
@@ -35,8 +38,7 @@
     "stc_token": "<|tool_call>",
     "std_token": "<|tool>",
     "str_token": "<|tool_response>",
-    "think_token": "<|think|>",
-    "video_token": "<|video|>"
+    "think_token": "<|think|>"
   },
   "pad_token": "<pad>",
   "padding_side": "left",
@@ -89,6 +91,5 @@
   "str_token": "<|tool_response>",
   "think_token": "<|think|>",
   "tokenizer_class": "GemmaTokenizer",
-  "unk_token": "<unk>",
-  "video_token": "<|video|>"
+  "unk_token": "<unk>"
 }