pipenetwork commited on
Commit
ecbc3d2
·
verified ·
1 Parent(s): 7882f28

Add files using upload-large-folder tool

Browse files
Files changed (50) hide show
  1. README.md +73 -0
  2. chat_template.jinja +129 -0
  3. config.json +139 -0
  4. inkling_mlx/__init__.py +14 -0
  5. inkling_mlx/attention.py +135 -0
  6. inkling_mlx/audio.py +34 -0
  7. inkling_mlx/cache.py +66 -0
  8. inkling_mlx/common.py +63 -0
  9. inkling_mlx/config.py +207 -0
  10. inkling_mlx/convert.py +217 -0
  11. inkling_mlx/convert_cli.py +44 -0
  12. inkling_mlx/generate.py +74 -0
  13. inkling_mlx/layers.py +47 -0
  14. inkling_mlx/load.py +69 -0
  15. inkling_mlx/model.py +79 -0
  16. inkling_mlx/moe.py +117 -0
  17. inkling_mlx/text.py +50 -0
  18. inkling_mlx/vision.py +126 -0
  19. model-00013-of-00097.safetensors +3 -0
  20. model-00015-of-00097.safetensors +3 -0
  21. model-00016-of-00097.safetensors +3 -0
  22. model-00019-of-00097.safetensors +3 -0
  23. model-00021-of-00097.safetensors +3 -0
  24. model-00022-of-00097.safetensors +3 -0
  25. model-00024-of-00097.safetensors +3 -0
  26. model-00027-of-00097.safetensors +3 -0
  27. model-00035-of-00097.safetensors +3 -0
  28. model-00036-of-00097.safetensors +3 -0
  29. model-00051-of-00097.safetensors +3 -0
  30. model-00052-of-00097.safetensors +3 -0
  31. model-00054-of-00097.safetensors +3 -0
  32. model-00057-of-00097.safetensors +3 -0
  33. model-00058-of-00097.safetensors +3 -0
  34. model-00060-of-00097.safetensors +3 -0
  35. model-00063-of-00097.safetensors +3 -0
  36. model-00065-of-00097.safetensors +3 -0
  37. model-00069-of-00097.safetensors +3 -0
  38. model-00074-of-00097.safetensors +3 -0
  39. model-00077-of-00097.safetensors +3 -0
  40. model-00082-of-00097.safetensors +3 -0
  41. model-00084-of-00097.safetensors +3 -0
  42. model-00090-of-00097.safetensors +3 -0
  43. model-00093-of-00097.safetensors +3 -0
  44. model-00095-of-00097.safetensors +3 -0
  45. model-00096-of-00097.safetensors +3 -0
  46. model.safetensors.index.json +0 -0
  47. processor_config.json +46 -0
  48. special_tokens_map.json +22 -0
  49. tiktoken/tokenizer.model +3 -0
  50. tokenizer_config.json +508 -0
README.md ADDED
@@ -0,0 +1,73 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: apache-2.0
3
+ base_model: thinkingmachines/Inkling
4
+ base_model_relation: quantized
5
+ pipeline_tag: image-text-to-text
6
+ library_name: mlx
7
+ tags:
8
+ - mlx
9
+ - moe
10
+ - multimodal
11
+ - inkling
12
+ - thinking-machines
13
+ ---
14
+
15
+ # Inkling-MLX-4bit
16
+
17
+ **Built with Inkling (Thinking Machines Lab).**
18
+
19
+ MLX (Apple Silicon) conversion of
20
+ [thinkingmachines/Inkling](https://huggingface.co/thinkingmachines/Inkling),
21
+ quantized to **4-bit** (affine group quant, group size 64).
22
+
23
+ **Code / loader:** [github.com/PipeNetwork/inkling-mlx](https://github.com/PipeNetwork/inkling-mlx)
24
+
25
+ Inkling is a **975B-total / 41B-active** sparse-MoE, natively multimodal model
26
+ (text + image/video + audio → text). This is the **full multimodal** conversion:
27
+ all three towers (text backbone, HMLP vision, dMel audio) are ported; the
28
+ multi-token-prediction head is dropped (inference-irrelevant).
29
+
30
+ ## Quantizations
31
+
32
+ | Variant | Size | Notes |
33
+ |---|---|---|
34
+ | [8bit](https://huggingface.co/pipenetwork/Inkling-MLX-8bit) | ~937 GB | near-lossless |
35
+ | [6bit](https://huggingface.co/pipenetwork/Inkling-MLX-6bit) | ~717 GB | high quality |
36
+ | [4bit](https://huggingface.co/pipenetwork/Inkling-MLX-4bit) | ~490 GB | balanced default |
37
+
38
+ ## ⚠️ Loading requires the bundled `inkling_mlx` loader
39
+
40
+ The `inkling_mm_model` architecture is **not** in stock `mlx-lm` / `mlx-vlm`, so this
41
+ repo bundles a minimal, numerically-validated MLX implementation under `inkling_mlx/`.
42
+
43
+ ```bash
44
+ pip install mlx mlx-lm transformers
45
+ ```
46
+ ```python
47
+ from inkling_mlx.load import load
48
+ from inkling_mlx.generate import greedy_generate
49
+ from transformers import AutoTokenizer
50
+
51
+ model, config = load("/path/to/this/repo")
52
+ tok = AutoTokenizer.from_pretrained("/path/to/this/repo", trust_remote_code=True)
53
+ ids = tok("The capital of France is")["input_ids"]
54
+ print(tok.decode(greedy_generate(model, config, ids, max_new_tokens=64)))
55
+ ```
56
+
57
+ Needs an Apple-Silicon Mac with enough unified memory to hold the weights (≈ the
58
+ size above).
59
+
60
+ ## Status & caveats
61
+
62
+ - **Text generation** works end-to-end via an incremental KV + short-convolution cache.
63
+ - The **vision and audio towers are ported and numerically validated** (fp32 parity
64
+ vs the PyTorch reference, max |Δ| ~1e-5), but currently take **pre-featurized**
65
+ inputs (`pixel_values` / `audio_input_ids`); the image/audio *preprocessing* pipeline
66
+ is not yet bundled.
67
+ - Quantized: attention / MLP / expert projections, token embed+unembed, and the
68
+ vision/audio matmuls. Kept in higher precision: the MoE router, RMSNorms, the four
69
+ short-convolutions per layer, and the relative-position bias.
70
+
71
+ Conversion is streaming (tensor-by-tensor; the ~1.9 TB bf16 model never fully loads
72
+ into RAM) and was validated with fp32 numerical parity against transformers PR #47347.
73
+ License: Apache-2.0 (inherits the base model).
chat_template.jinja ADDED
@@ -0,0 +1,129 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {%- set effort_map = {"none": 0.0, "minimal": 0.1, "low": 0.2, "medium": 0.7, "high": 0.9, "max": 0.99} -%}
2
+ {%- set role_token = {"user": "<|message_user|>", "assistant": "<|message_model|>", "system": "<|message_system|>", "tool": "<|message_tool|>"} -%}
3
+
4
+ {%- macro emit_thinking_effort() -%}
5
+ {%- set eff = reasoning_effort if reasoning_effort is defined and reasoning_effort is not none else 0.9 -%}
6
+ {%- if eff is string -%}
7
+ {%- set key = eff | trim -%}
8
+ {%- if key not in effort_map -%}
9
+ {{- raise_exception("Unknown reasoning_effort: " ~ eff) -}}
10
+ {%- endif -%}
11
+ {%- set num = effort_map[key] -%}
12
+ {%- else -%}
13
+ {%- set num = eff | float -%}
14
+ {%- endif -%}
15
+ {%- if num < 0.0 or num > 0.99 -%}
16
+ {{- raise_exception("reasoning_effort must be in [0.0, 0.99]") -}}
17
+ {%- endif -%}
18
+ {{- "<|message_system|><|content_text|>Thinking effort level: " -}}
19
+ {%- if num == 0.0 -%}0{%- else -%}{{ num }}{%- endif -%}
20
+ {{- "<|end_message|>" -}}
21
+ {%- endmacro -%}
22
+
23
+ {%- if tools -%}
24
+ {%- set tool_state = namespace(specs=[]) -%}
25
+ {%- for tool in tools -%}
26
+ {%- set fn = tool.function if tool.function is defined else tool -%}
27
+ {%- set spec = {
28
+ "description": (fn.description if fn.description is defined and fn.description else ""),
29
+ "name": fn.name,
30
+ "parameters": (fn.parameters if fn.parameters is defined and fn.parameters else {}),
31
+ "type": (tool.type if tool.type is defined and tool.type else "function"),
32
+ } -%}
33
+ {%- set tool_state.specs = tool_state.specs + [spec] -%}
34
+ {%- endfor -%}
35
+ {{- "<|message_system|>tool_declare<|content_xml|>" -}}
36
+ {{- tool_state.specs | tojson(sort_keys=true, separators=(",", ":")) -}}
37
+ {{- "<|end_message|>" -}}
38
+ {%- endif -%}
39
+
40
+ {%- set state = namespace(effort_emitted=false) -%}
41
+ {%- for message in messages -%}
42
+ {%- if message.role not in role_token -%}
43
+ {{- raise_exception("Unknown message role: " ~ message.role) -}}
44
+ {%- endif -%}
45
+ {%- if not state.effort_emitted and message.role != "system" -%}
46
+ {{- emit_thinking_effort() -}}
47
+ {%- set state.effort_emitted = true -%}
48
+ {%- endif -%}
49
+
50
+ {%- set rtok = role_token[message.role] -%}
51
+
52
+ {%- if message.role == "tool" -%}
53
+ {%- set tool_name_state = namespace(name="") -%}
54
+ {%- if message.name is defined and message.name -%}
55
+ {%- set tool_name_state.name = message.name -%}
56
+ {%- elif message.tool_call_id is defined and message.tool_call_id -%}
57
+ {%- for prev in messages -%}
58
+ {%- if prev.role == "assistant" and prev.tool_calls -%}
59
+ {%- for tc in prev.tool_calls -%}
60
+ {%- if tc.id is defined and tc.id == message.tool_call_id and tc.function.name is defined -%}
61
+ {%- set tool_name_state.name = tc.function.name -%}
62
+ {%- endif -%}
63
+ {%- endfor -%}
64
+ {%- endif -%}
65
+ {%- endfor -%}
66
+ {%- endif -%}
67
+ {{- rtok -}}
68
+ {%- if tool_name_state.name -%}{{- tool_name_state.name -}}{%- endif -%}
69
+ {{- "<|content_text|>" -}}
70
+ {%- if message.content is string -%}{{- message.content -}}{%- endif -%}
71
+ {{- "<|end_message|>" -}}
72
+
73
+ {%- else -%}
74
+ {%- if message.role == "assistant" and message.reasoning_content is defined and message.reasoning_content -%}
75
+ {{- "<|message_model|><|content_thinking|>" ~ message.reasoning_content ~ "<|end_message|>" -}}
76
+ {%- endif -%}
77
+
78
+ {%- if message.content is string -%}
79
+ {{- rtok ~ "<|content_text|>" ~ message.content ~ "<|end_message|>" -}}
80
+ {%- elif message.content -%}
81
+ {%- for part in message.content -%}
82
+ {%- if part is string -%}
83
+ {{- rtok ~ "<|content_text|>" ~ part ~ "<|end_message|>" -}}
84
+ {%- elif part.type is not defined or part.type in ("text", "input_text") -%}
85
+ {%- set text_part = (part.text if part.text is defined and part.text is string else "") -%}
86
+ {{- rtok ~ "<|content_text|>" ~ text_part ~ "<|end_message|>" -}}
87
+ {%- elif part.type in ("image", "input_image", "image_url") -%}
88
+ {{- rtok ~ "<|content_image|><|unused_200054|><|end_message|>" -}}
89
+ {%- elif part.type in ("audio", "input_audio", "audio_url") -%}
90
+ {{- rtok ~ "<|content_audio_input|><|unused_200053|><|audio_end|><|end_message|>" -}}
91
+ {%- else -%}
92
+ {{- raise_exception("Unsupported content part type: " ~ part.type) -}}
93
+ {%- endif -%}
94
+ {%- endfor -%}
95
+ {%- endif -%}
96
+
97
+ {%- if message.role == "assistant" and message.tool_calls -%}
98
+ {%- for tc in message.tool_calls -%}
99
+ {%- set fn = tc.function -%}
100
+ {%- if fn.name is not defined or fn.name is not string -%}
101
+ {{- raise_exception("tool call function name must be a string") -}}
102
+ {%- endif -%}
103
+ {%- set args = fn.arguments if fn.arguments is defined and fn.arguments else {} -%}
104
+ {%- if args is string -%}
105
+ {{- raise_exception("tool call arguments must be a parsed object, not a JSON string; canonicalize upstream") -}}
106
+ {%- endif -%}
107
+ {%- if args is not mapping -%}
108
+ {{- raise_exception("tool call arguments must be an object") -}}
109
+ {%- endif -%}
110
+ {{- "<|message_model|>" ~ fn.name ~ "<|content_invoke_tool_json|>" -}}
111
+ {{- '{"name":' ~ (fn.name | tojson(sort_keys=true, separators=(",", ":"))) ~ ',"args":' -}}
112
+ {{- (args | tojson(sort_keys=true, separators=(",", ":"))) -}}
113
+ {{- "}<|end_message|>" -}}
114
+ {%- endfor -%}
115
+ {%- endif -%}
116
+
117
+ {%- if message.role == "assistant" -%}
118
+ {{- "<|content_model_end_sampling|>" -}}
119
+ {%- endif -%}
120
+ {%- endif -%}
121
+ {%- endfor -%}
122
+
123
+ {%- if not state.effort_emitted -%}
124
+ {{- emit_thinking_effort() -}}
125
+ {%- endif -%}
126
+
127
+ {%- if add_generation_prompt -%}
128
+ {{- "<|message_model|>" -}}
129
+ {%- endif -%}
config.json ADDED
@@ -0,0 +1,139 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "InklingForConditionalGeneration"
4
+ ],
5
+ "model_type": "inkling_mm_model",
6
+ "eos_token_id": 200006,
7
+ "text_config": {
8
+ "model_max_length": 1048576,
9
+ "torch_dtype": "bfloat16",
10
+ "hidden_size": 6144,
11
+ "num_hidden_layers": 66,
12
+ "vocab_size": 201024,
13
+ "num_attention_heads": 64,
14
+ "num_key_value_heads": 8,
15
+ "head_dim": 128,
16
+ "d_rel": 16,
17
+ "rel_extent": 1024,
18
+ "q_bias": false,
19
+ "o_bias": false,
20
+ "log_scaling_n_floor": 128000,
21
+ "log_scaling_alpha": 0.1,
22
+ "rms_norm_eps": 1e-06,
23
+ "use_embed_norm": true,
24
+ "local_layer_ids": [
25
+ 0,
26
+ 1,
27
+ 2,
28
+ 3,
29
+ 4,
30
+ 6,
31
+ 7,
32
+ 8,
33
+ 9,
34
+ 10,
35
+ 12,
36
+ 13,
37
+ 14,
38
+ 15,
39
+ 16,
40
+ 18,
41
+ 19,
42
+ 20,
43
+ 21,
44
+ 22,
45
+ 24,
46
+ 25,
47
+ 26,
48
+ 27,
49
+ 28,
50
+ 30,
51
+ 31,
52
+ 32,
53
+ 33,
54
+ 34,
55
+ 36,
56
+ 37,
57
+ 38,
58
+ 39,
59
+ 40,
60
+ 42,
61
+ 43,
62
+ 44,
63
+ 45,
64
+ 46,
65
+ 48,
66
+ 49,
67
+ 50,
68
+ 51,
69
+ 52,
70
+ 54,
71
+ 55,
72
+ 56,
73
+ 57,
74
+ 58,
75
+ 60,
76
+ 61,
77
+ 62,
78
+ 63,
79
+ 64
80
+ ],
81
+ "dense_mlp_idx": 2,
82
+ "use_sconv": true,
83
+ "sconv_kernel_size": 4,
84
+ "unpadded_vocab_size": 200058,
85
+ "logits_mup_width_multiplier": 24.0,
86
+ "final_logit_softcapping": null,
87
+ "swa_head_dim": 128,
88
+ "swa_num_attention_heads": 64,
89
+ "swa_num_key_value_heads": 16,
90
+ "sliding_window_size": 512,
91
+ "n_routed_experts": 256,
92
+ "num_experts_per_tok": 6,
93
+ "n_shared_experts": 2,
94
+ "shared_expert_sink": true,
95
+ "dense_intermediate_size": 24576,
96
+ "intermediate_size": 3072,
97
+ "route_scale": 8.0,
98
+ "use_gate_bias": true,
99
+ "gate_activation": "sigmoid",
100
+ "norm_after_topk": true,
101
+ "use_global_scale": true
102
+ },
103
+ "audio_config": {
104
+ "decoder_dmodel": 6144,
105
+ "n_mel_bins": 80,
106
+ "mel_vocab_size": 16,
107
+ "bias": false,
108
+ "dmel_min_value": -7.0,
109
+ "dmel_max_value": 2.0,
110
+ "use_audio_norm": true,
111
+ "audio_mode": "dmel"
112
+ },
113
+ "vision_config": {
114
+ "vision_encoder_type": "hmlp",
115
+ "decoder_dmodel": 6144,
116
+ "patch_size": 40,
117
+ "temporal_patch_size": 2,
118
+ "n_channels": 3,
119
+ "n_layers": 4,
120
+ "use_vision_norm": true
121
+ },
122
+ "mtp_config": {
123
+ "num_nextn_predict_layers": 8,
124
+ "chain_hidden_post_norm": false,
125
+ "local_layer_ids": [
126
+ 0,
127
+ 2,
128
+ 4,
129
+ 5,
130
+ 6,
131
+ 7
132
+ ]
133
+ },
134
+ "quantization": {
135
+ "group_size": 64,
136
+ "bits": 4,
137
+ "recipe": "uniform"
138
+ }
139
+ }
inkling_mlx/__init__.py ADDED
@@ -0,0 +1,14 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """MLX port of thinkingmachines/Inkling (975B MoE, natively multimodal)."""
2
+
3
+ from .config import AudioConfig, InklingConfig, TextConfig, VisionConfig
4
+ from .model import InklingForConditionalGeneration
5
+ from .text import TextModel
6
+
7
+ __all__ = [
8
+ "InklingConfig",
9
+ "TextConfig",
10
+ "VisionConfig",
11
+ "AudioConfig",
12
+ "InklingForConditionalGeneration",
13
+ "TextModel",
14
+ ]
inkling_mlx/attention.py ADDED
@@ -0,0 +1,135 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Inkling attention: hybrid local/global, per-head q/k RMSNorm, relative-position
2
+ logits bias, optional log-scaling, and short-convolution on k/v.
3
+
4
+ Mirrors ``InklingAttention`` + ``InklingRelativeLogits`` from transformers PR #47347.
5
+ This implementation is prefill-oriented (full-sequence, no KV cache); an incremental
6
+ cache (including the 4 per-layer conv states) can be layered on top later.
7
+ """
8
+
9
+ from __future__ import annotations
10
+
11
+ import math
12
+
13
+ import mlx.core as mx
14
+ import mlx.nn as nn
15
+
16
+ from .common import RMSNorm, ShortConvolution
17
+ from .config import TextConfig
18
+
19
+ NEG_INF = -1e30
20
+
21
+
22
+ class RelativeLogits(nn.Module):
23
+ """Hidden-state-conditioned relative position bias.
24
+
25
+ ``proj`` is a bank of bias-vs-distance profiles ``[d_rel, rel_extent]``. Each
26
+ query's ``d_rel`` relative-state vector mixes them into one bias value per
27
+ backward distance; the bias is zero outside ``0 <= distance < rel_extent``.
28
+ """
29
+
30
+ def __init__(self, d_rel: int, rel_extent: int):
31
+ super().__init__()
32
+ self.rel_extent = rel_extent
33
+ self.proj = mx.zeros((d_rel, rel_extent))
34
+
35
+ def __call__(self, relative_states, q_pos, kv_pos):
36
+ # relative_states: [B, Lq, heads, d_rel]
37
+ # rel_logits: [B, Lq, heads, rel_extent] -> [B, heads, Lq, rel_extent]
38
+ rel_logits = mx.swapaxes(relative_states @ self.proj, 1, 2)
39
+ B, H, Lq, _ = rel_logits.shape
40
+ distance = q_pos[:, None] - kv_pos[None, :] # [Lq, Lkv]
41
+ gather = mx.clip(distance, 0, self.rel_extent - 1) # [Lq, Lkv]
42
+ gather = mx.broadcast_to(gather[None, None], (B, H, Lq, gather.shape[-1]))
43
+ bias = mx.take_along_axis(rel_logits, gather, axis=-1) # [B, H, Lq, Lkv]
44
+ valid = (distance >= 0) & (distance < self.rel_extent) # [Lq, Lkv]
45
+ return mx.where(valid[None, None], bias, 0.0)
46
+
47
+
48
+ class Attention(nn.Module):
49
+ def __init__(self, config: TextConfig, layer_idx: int):
50
+ super().__init__()
51
+ self.config = config
52
+ self.layer_idx = layer_idx
53
+ self.is_sliding = config.layer_types[layer_idx] == "hybrid_sliding"
54
+
55
+ self.head_dim = config.swa_head_dim if self.is_sliding else config.head_dim
56
+ self.num_heads = config.swa_num_attention_heads if self.is_sliding else config.num_attention_heads
57
+ self.num_kv_heads = config.swa_num_key_value_heads if self.is_sliding else config.num_key_value_heads
58
+ self.n_rep = self.num_heads // self.num_kv_heads
59
+ self.sliding_window = config.sliding_window_size if self.is_sliding else None
60
+ self.rel_extent = config.sliding_window_size if self.is_sliding else config.rel_extent
61
+ self.d_rel = config.d_rel
62
+ # q/k are per-head RMS-normalized, hence 1/d rather than 1/sqrt(d)
63
+ self.scaling = 1.0 / self.head_dim
64
+
65
+ h = config.hidden_size
66
+ self.wq_du = nn.Linear(h, self.num_heads * self.head_dim, bias=False)
67
+ self.wk_dv = nn.Linear(h, self.num_kv_heads * self.head_dim, bias=False)
68
+ self.wv_dv = nn.Linear(h, self.num_kv_heads * self.head_dim, bias=False)
69
+ self.wr_du = nn.Linear(h, self.num_heads * self.d_rel, bias=False)
70
+ self.wo_ud = nn.Linear(self.num_heads * self.head_dim, h, bias=False)
71
+
72
+ self.k_sconv = ShortConvolution(self.num_kv_heads * self.head_dim, config.sconv_kernel_size)
73
+ self.v_sconv = ShortConvolution(self.num_kv_heads * self.head_dim, config.sconv_kernel_size)
74
+ self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
75
+ self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
76
+ self.rel_logits_proj = RelativeLogits(self.d_rel, self.rel_extent)
77
+
78
+ def __call__(self, hidden_states, start_pos=0, kv_cache=None,
79
+ k_conv=None, v_conv=None, conv_mask=None):
80
+ B, L, _ = hidden_states.shape
81
+
82
+ q = self.wq_du(hidden_states)
83
+ k = self.k_sconv(self.wk_dv(hidden_states), mask=conv_mask, cache=k_conv)
84
+ v = self.v_sconv(self.wv_dv(hidden_states), mask=conv_mask, cache=v_conv)
85
+ rel = self.wr_du(hidden_states)
86
+
87
+ q = self.q_norm(q.reshape(B, L, self.num_heads, self.head_dim))
88
+ k = self.k_norm(k.reshape(B, L, self.num_kv_heads, self.head_dim))
89
+ v = v.reshape(B, L, self.num_kv_heads, self.head_dim)
90
+
91
+ # -> [B, heads, L, head_dim]
92
+ q = q.transpose(0, 2, 1, 3)
93
+ k = k.transpose(0, 2, 1, 3)
94
+ v = v.transpose(0, 2, 1, 3)
95
+
96
+ q_pos = mx.arange(L) + start_pos
97
+ if kv_cache is not None:
98
+ k, v = kv_cache.update(k, v) # full history
99
+ kv_pos = mx.arange(k.shape[2])
100
+
101
+ rel = rel.reshape(B, L, self.num_heads, self.d_rel)
102
+ position_bias = self.rel_logits_proj(rel, q_pos, kv_pos) # [B, heads, Lq, Lkv]
103
+
104
+ # log-scaling (global layers only; no-op for context <= n_floor)
105
+ if not self.is_sliding and self.config.log_scaling_n_floor is not None:
106
+ n_floor = self.config.log_scaling_n_floor
107
+ eff_n = (q_pos + 1).astype(mx.float32)
108
+ tau = 1.0 + self.config.log_scaling_alpha * mx.log(
109
+ mx.maximum(eff_n / n_floor, 1.0)
110
+ )
111
+ tau_q = tau.reshape(1, 1, -1, 1)
112
+ q = (q.astype(mx.float32) * tau_q).astype(q.dtype)
113
+ position_bias = (position_bias.astype(mx.float32) * tau_q).astype(position_bias.dtype)
114
+
115
+ # GQA expand
116
+ if self.n_rep > 1:
117
+ k = mx.repeat(k, self.n_rep, axis=1)
118
+ v = mx.repeat(v, self.n_rep, axis=1)
119
+
120
+ scores = (q.astype(mx.float32) @ mx.swapaxes(k, 2, 3).astype(mx.float32)) * self.scaling
121
+ scores = scores + position_bias.astype(mx.float32)
122
+ scores = scores + self._causal_mask(q_pos, kv_pos)
123
+ weights = mx.softmax(scores, axis=-1)
124
+ out = weights.astype(v.dtype) @ v # [B, heads, Lq, head_dim]
125
+
126
+ out = out.transpose(0, 2, 1, 3).reshape(B, L, self.num_heads * self.head_dim)
127
+ return self.wo_ud(out)
128
+
129
+ def _causal_mask(self, q_pos, kv_pos):
130
+ distance = q_pos[:, None] - kv_pos[None, :] # [Lq, Lkv]
131
+ allowed = distance >= 0
132
+ if self.sliding_window is not None:
133
+ allowed = allowed & (distance < self.sliding_window)
134
+ mask = mx.where(allowed, 0.0, NEG_INF)
135
+ return mask[None, None].astype(mx.float32)
inkling_mlx/audio.py ADDED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Inkling audio tower: discrete dMel-token embedding + norm.
2
+
3
+ Each audio frame is ``n_mel_bins`` discretized bins (values in ``[0, mel_vocab_size)``);
4
+ each bin is embedded from its own slice of a shared table (offset ``bin * mel_vocab_size``)
5
+ and the per-bin embeddings are summed. Mirrors ``InklingAudioModel`` /
6
+ ``InklingAudioModelEmbeddings``. Checkpoint keys: ``audio.encoder.weight`` (the
7
+ ``[n_mel_bins*mel_vocab_size, hidden]`` table) and ``audio.final_norm.weight``.
8
+ """
9
+
10
+ from __future__ import annotations
11
+
12
+ import mlx.core as mx
13
+ import mlx.nn as nn
14
+
15
+ from .common import RMSNorm
16
+ from .config import AudioConfig
17
+
18
+
19
+ class AudioModel(nn.Module):
20
+ def __init__(self, config: AudioConfig):
21
+ super().__init__()
22
+ self.config = config
23
+ self.encoder = nn.Embedding(
24
+ config.n_mel_bins * config.mel_vocab_size, config.text_hidden_size
25
+ )
26
+ self.final_norm = RMSNorm(config.text_hidden_size, eps=config.rms_norm_eps)
27
+ # non-persistent: arange(n_mel_bins) * mel_vocab_size
28
+ self._offsets = mx.arange(config.n_mel_bins) * config.mel_vocab_size
29
+
30
+ def __call__(self, audio_input_ids: mx.array) -> mx.array:
31
+ # audio_input_ids: [..., n_mel_bins] with values in [0, mel_vocab_size)
32
+ embeds = self.encoder(audio_input_ids + self._offsets) # [..., n_mel_bins, hidden]
33
+ embeds = embeds.sum(axis=-2) # [..., hidden]
34
+ return self.final_norm(embeds)
inkling_mlx/cache.py ADDED
@@ -0,0 +1,66 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Incremental caches for Inkling generation.
2
+
3
+ Two kinds of per-layer state must persist across decode steps:
4
+
5
+ * ``KVCache`` — the appended key/value tensors for attention.
6
+ * ``ConvCache`` — the last ``kernel-1`` inputs of each depthwise short-convolution
7
+ (there are 4 per layer: k, v, post-attn, post-mlp).
8
+
9
+ A ``LayerCache`` bundles one KVCache + the 4 ConvCaches; ``make_cache`` builds one
10
+ per decoder layer. Absolute key positions are always ``arange(kv_len)`` because the
11
+ cache holds every key from position 0 (KVCache keeps the full history — correct for
12
+ both global and sliding layers, since the sliding-window constraint is enforced by
13
+ the attention mask).
14
+ """
15
+
16
+ from __future__ import annotations
17
+
18
+ import mlx.core as mx
19
+
20
+
21
+ class ConvCache:
22
+ """Holds the last ``kernel-1`` inputs of a short convolution."""
23
+
24
+ __slots__ = ("state",)
25
+
26
+ def __init__(self):
27
+ self.state = None # [B, kernel-1, C] or None
28
+
29
+
30
+ class KVCache:
31
+ """Appends keys/values along the sequence axis (full history)."""
32
+
33
+ __slots__ = ("keys", "values")
34
+
35
+ def __init__(self):
36
+ self.keys = None # [B, heads, T, d]
37
+ self.values = None
38
+
39
+ @property
40
+ def offset(self) -> int:
41
+ return 0 if self.keys is None else self.keys.shape[2]
42
+
43
+ def update(self, k: mx.array, v: mx.array):
44
+ if self.keys is None:
45
+ self.keys, self.values = k, v
46
+ else:
47
+ self.keys = mx.concatenate([self.keys, k], axis=2)
48
+ self.values = mx.concatenate([self.values, v], axis=2)
49
+ return self.keys, self.values
50
+
51
+
52
+ class LayerCache:
53
+ __slots__ = ("kv", "k_conv", "v_conv", "attn_conv", "mlp_conv")
54
+
55
+ def __init__(self):
56
+ self.kv = KVCache()
57
+ self.k_conv = ConvCache()
58
+ self.v_conv = ConvCache()
59
+ self.attn_conv = ConvCache()
60
+ self.mlp_conv = ConvCache()
61
+
62
+
63
+ def make_cache(model) -> list[LayerCache]:
64
+ """One LayerCache per text decoder layer."""
65
+ n = len(model.model.llm.layers) if hasattr(model, "model") else len(model.layers)
66
+ return [LayerCache() for _ in range(n)]
inkling_mlx/common.py ADDED
@@ -0,0 +1,63 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Shared low-level modules for the Inkling MLX port."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import mlx.core as mx
6
+ import mlx.nn as nn
7
+
8
+
9
+ class RMSNorm(nn.Module):
10
+ """Llama-style RMSNorm (compute in fp32, weight is a gain).
11
+
12
+ Matches ``LlamaRMSNorm``: ``x_fp32 * rsqrt(mean(x^2) + eps) * weight``.
13
+ """
14
+
15
+ def __init__(self, dims: int, eps: float = 1e-6):
16
+ super().__init__()
17
+ self.weight = mx.ones((dims,))
18
+ self.eps = eps
19
+
20
+ def __call__(self, x: mx.array) -> mx.array:
21
+ return mx.fast.rms_norm(x, self.weight, self.eps)
22
+
23
+
24
+ class ShortConvolution(nn.Module):
25
+ """Depthwise causal 1-D convolution with a residual add, computed in fp32.
26
+
27
+ Mirrors ``InklingShortConvolution``: a per-channel (groups == channels) causal
28
+ conv1d of ``kernel_size`` taps, no bias, no activation, then ``out + input``.
29
+ The reference keeps this module in fp32 regardless of the model dtype
30
+ (``_keep_in_fp32_modules_strict``), so we upcast here too.
31
+
32
+ Weight layout (MLX ``conv1d``): ``[channels, kernel_size, 1]``.
33
+ """
34
+
35
+ def __init__(self, channels: int, kernel_size: int):
36
+ super().__init__()
37
+ self.channels = channels
38
+ self.kernel_size = kernel_size
39
+ # [C_out, K, C_in // groups] with groups == channels -> [C, K, 1]
40
+ self.weight = mx.zeros((channels, kernel_size, 1))
41
+
42
+ def __call__(self, x: mx.array, mask: mx.array | None = None, cache=None) -> mx.array:
43
+ # x: [batch, seq, channels]
44
+ in_dtype = x.dtype
45
+ xf = x.astype(mx.float32)
46
+ residual = xf
47
+ if mask is not None:
48
+ xf = xf * mask.astype(mx.float32)
49
+ k = self.kernel_size
50
+ B, seq, C = xf.shape
51
+ w = self.weight.astype(mx.float32)
52
+ if cache is not None:
53
+ # left-context = cached last (k-1) inputs (zeros on the first call);
54
+ # a "valid" conv over [left, xf] yields exactly `seq` causal outputs.
55
+ left = cache.state if cache.state is not None else mx.zeros((B, k - 1, C), dtype=mx.float32)
56
+ x_in = mx.concatenate([left, xf], axis=1)
57
+ out = mx.conv1d(x_in, w, padding=0, groups=self.channels)
58
+ cache.state = x_in[:, -(k - 1):, :]
59
+ else:
60
+ # causal: left-pad by (k-1), keep first `seq` outputs (== zero left-context)
61
+ out = mx.conv1d(xf, w, padding=k - 1, groups=self.channels)[:, :seq, :]
62
+ out = out + residual
63
+ return out.astype(in_dtype)
inkling_mlx/config.py ADDED
@@ -0,0 +1,207 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Configuration for the Inkling multimodal model (MLX port).
2
+
3
+ Mirrors ``thinkingmachines/Inkling`` ``config.json`` and the transformers PR #47347
4
+ reference (``InklingConfig`` / ``InklingTextConfig`` / ``InklingVisionConfig`` /
5
+ ``InklingAudioConfig``). We parse the *checkpoint* config layout (top-level
6
+ ``text_config`` / ``vision_config`` / ``audio_config`` / ``mtp_config``), not the
7
+ flattened transformers layout.
8
+ """
9
+
10
+ from __future__ import annotations
11
+
12
+ from dataclasses import dataclass, field
13
+ from typing import Any
14
+
15
+
16
+ def _get(d: dict, *names, default=None):
17
+ for n in names:
18
+ if n in d and d[n] is not None:
19
+ return d[n]
20
+ return default
21
+
22
+
23
+ @dataclass
24
+ class TextConfig:
25
+ hidden_size: int = 6144
26
+ num_hidden_layers: int = 66
27
+ vocab_size: int = 201024
28
+ unpadded_vocab_size: int | None = 200058
29
+
30
+ # global (full) attention
31
+ num_attention_heads: int = 64
32
+ num_key_value_heads: int = 8
33
+ head_dim: int = 128
34
+ # sliding-window attention
35
+ swa_num_attention_heads: int = 64
36
+ swa_num_key_value_heads: int = 16
37
+ swa_head_dim: int = 128
38
+ sliding_window_size: int = 512
39
+
40
+ # relative-position logits
41
+ d_rel: int = 16
42
+ rel_extent: int = 1024
43
+ log_scaling_n_floor: int | None = 128000
44
+ log_scaling_alpha: float = 0.1
45
+
46
+ rms_norm_eps: float = 1e-6
47
+ use_embed_norm: bool = True
48
+
49
+ # short convolution
50
+ sconv_kernel_size: int = 4
51
+
52
+ # dense vs MoE MLP
53
+ dense_mlp_idx: int = 2
54
+ dense_intermediate_size: int = 24576 # dense MLP intermediate
55
+ moe_intermediate_size: int = 3072 # per-expert intermediate
56
+
57
+ # MoE routing
58
+ n_routed_experts: int = 256
59
+ num_experts_per_tok: int = 6
60
+ n_shared_experts: int = 2
61
+ shared_expert_sink: bool = True
62
+ route_scale: float = 8.0
63
+ use_gate_bias: bool = True
64
+ norm_after_topk: bool = True
65
+ use_global_scale: bool = True
66
+
67
+ logits_mup_width_multiplier: float = 24.0
68
+ hidden_act: str = "silu"
69
+
70
+ max_position_embeddings: int = 1048576
71
+
72
+ # which layer indices use sliding-window ("local") attention
73
+ local_layer_ids: list[int] = field(default_factory=list)
74
+
75
+ # MTP head (dropped for inference)
76
+ num_mtp_layers: int | None = None
77
+
78
+ @property
79
+ def layer_types(self) -> list[str]:
80
+ local = set(self.local_layer_ids)
81
+ return [
82
+ "hybrid_sliding" if i in local else "hybrid"
83
+ for i in range(self.num_hidden_layers)
84
+ ]
85
+
86
+ @property
87
+ def mlp_layer_types(self) -> list[str]:
88
+ return [
89
+ "dense" if i < self.dense_mlp_idx else "sparse"
90
+ for i in range(self.num_hidden_layers)
91
+ ]
92
+
93
+ @classmethod
94
+ def from_dict(cls, tc: dict) -> "TextConfig":
95
+ return cls(
96
+ hidden_size=_get(tc, "hidden_size", default=6144),
97
+ num_hidden_layers=_get(tc, "num_hidden_layers", default=66),
98
+ vocab_size=_get(tc, "vocab_size", default=201024),
99
+ unpadded_vocab_size=_get(tc, "unpadded_vocab_size"),
100
+ num_attention_heads=_get(tc, "num_attention_heads", default=64),
101
+ num_key_value_heads=_get(tc, "num_key_value_heads", default=8),
102
+ head_dim=_get(tc, "head_dim", default=128),
103
+ swa_num_attention_heads=_get(tc, "swa_num_attention_heads", default=64),
104
+ swa_num_key_value_heads=_get(tc, "swa_num_key_value_heads", default=16),
105
+ swa_head_dim=_get(tc, "swa_head_dim", default=128),
106
+ sliding_window_size=_get(tc, "sliding_window_size", default=512),
107
+ d_rel=_get(tc, "d_rel", default=16),
108
+ rel_extent=_get(tc, "rel_extent", default=1024),
109
+ log_scaling_n_floor=_get(tc, "log_scaling_n_floor"),
110
+ log_scaling_alpha=_get(tc, "log_scaling_alpha", default=0.1),
111
+ rms_norm_eps=_get(tc, "rms_norm_eps", default=1e-6),
112
+ use_embed_norm=_get(tc, "use_embed_norm", default=True),
113
+ sconv_kernel_size=_get(tc, "sconv_kernel_size", default=4),
114
+ dense_mlp_idx=_get(tc, "dense_mlp_idx", default=2),
115
+ dense_intermediate_size=_get(tc, "dense_intermediate_size", default=24576),
116
+ # checkpoint labels the *MoE* intermediate as `intermediate_size`
117
+ moe_intermediate_size=_get(tc, "intermediate_size", default=3072),
118
+ n_routed_experts=_get(tc, "n_routed_experts", default=256),
119
+ num_experts_per_tok=_get(tc, "num_experts_per_tok", default=6),
120
+ n_shared_experts=_get(tc, "n_shared_experts", default=2),
121
+ shared_expert_sink=_get(tc, "shared_expert_sink", default=True),
122
+ route_scale=_get(tc, "route_scale", default=8.0),
123
+ use_gate_bias=_get(tc, "use_gate_bias", default=True),
124
+ norm_after_topk=_get(tc, "norm_after_topk", default=True),
125
+ use_global_scale=_get(tc, "use_global_scale", default=True),
126
+ logits_mup_width_multiplier=_get(tc, "logits_mup_width_multiplier", default=24.0),
127
+ max_position_embeddings=_get(tc, "model_max_length", "max_position_embeddings", default=1048576),
128
+ local_layer_ids=list(_get(tc, "local_layer_ids", default=[]) or []),
129
+ )
130
+
131
+
132
+ @dataclass
133
+ class VisionConfig:
134
+ text_hidden_size: int = 6144
135
+ patch_size: int = 40
136
+ temporal_patch_size: int = 2
137
+ num_channels: int = 3
138
+ n_layers: int = 4
139
+ rms_norm_eps: float = 1e-6
140
+ use_vision_norm: bool = True
141
+
142
+ @classmethod
143
+ def from_dict(cls, vc: dict, text_hidden: int) -> "VisionConfig":
144
+ return cls(
145
+ text_hidden_size=text_hidden,
146
+ patch_size=_get(vc, "patch_size", default=40),
147
+ temporal_patch_size=_get(vc, "temporal_patch_size", default=2),
148
+ num_channels=_get(vc, "n_channels", "num_channels", default=3),
149
+ n_layers=_get(vc, "n_layers", "num_hidden_layers", default=4),
150
+ rms_norm_eps=_get(vc, "rms_norm_eps", default=1e-6),
151
+ use_vision_norm=_get(vc, "use_vision_norm", default=True),
152
+ )
153
+
154
+
155
+ @dataclass
156
+ class AudioConfig:
157
+ text_hidden_size: int = 6144
158
+ n_mel_bins: int = 80
159
+ mel_vocab_size: int = 16
160
+ rms_norm_eps: float = 1e-6
161
+
162
+ @classmethod
163
+ def from_dict(cls, ac: dict, text_hidden: int) -> "AudioConfig":
164
+ return cls(
165
+ text_hidden_size=text_hidden,
166
+ n_mel_bins=_get(ac, "n_mel_bins", default=80),
167
+ mel_vocab_size=_get(ac, "mel_vocab_size", default=16),
168
+ rms_norm_eps=_get(ac, "rms_norm_eps", default=1e-6),
169
+ )
170
+
171
+
172
+ @dataclass
173
+ class InklingConfig:
174
+ text: TextConfig
175
+ vision: VisionConfig
176
+ audio: AudioConfig
177
+ image_token_id: int = 200054
178
+ audio_token_id: int = 200053
179
+ image_bos_token_id: int = 200005
180
+ audio_bos_token_id: int = 200020
181
+ eos_token_id: int = 200006
182
+ model_type: str = "inkling_mm_model"
183
+
184
+ @classmethod
185
+ def from_dict(cls, cfg: dict) -> "InklingConfig":
186
+ tc = dict(cfg.get("text_config", {}))
187
+ mtp = cfg.get("mtp_config") or {}
188
+ if mtp.get("num_nextn_predict_layers") is not None:
189
+ tc.setdefault("num_mtp_layers", mtp.get("num_nextn_predict_layers"))
190
+ text = TextConfig.from_dict(tc)
191
+ vision = VisionConfig.from_dict(cfg.get("vision_config", {}) or {}, text.hidden_size)
192
+ audio = AudioConfig.from_dict(cfg.get("audio_config", {}) or {}, text.hidden_size)
193
+ return cls(
194
+ text=text,
195
+ vision=vision,
196
+ audio=audio,
197
+ image_token_id=_get(cfg, "image_token_id", default=200054),
198
+ audio_token_id=_get(cfg, "audio_token_id", default=200053),
199
+ image_bos_token_id=_get(cfg, "image_bos_token_id", default=200005),
200
+ audio_bos_token_id=_get(cfg, "audio_bos_token_id", default=200020),
201
+ eos_token_id=_get(cfg, "eos_token_id", default=200006),
202
+ model_type=_get(cfg, "model_type", default="inkling_mm_model"),
203
+ )
204
+
205
+ @property
206
+ def raw(self) -> dict[str, Any]:
207
+ return {"model_type": self.model_type}
inkling_mlx/convert.py ADDED
@@ -0,0 +1,217 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Streaming HF -> MLX conversion + quantization for Inkling.
2
+
3
+ The model is far too large (~1.9 TB bf16) to instantiate in RAM, so we convert
4
+ tensor-by-tensor: read each source shard (mmap), remap the name, apply the layout
5
+ transform, optionally affine-quantize the weight, and write output shards. Affine
6
+ quantization has no cross-tensor dependency, so per-tensor streaming is exactly
7
+ equivalent to ``nn.quantize(model)``.
8
+
9
+ Name/layout transforms vs. the checkpoint:
10
+ * ``*_sconv.weight`` [C,1,K] -> [C,K,1] (MLX conv1d layout)
11
+ * ``mlp.w13_dn`` [2I,H] -> gate_proj/up_proj (split dense fused gate+up)
12
+ * ``experts.w13_weight`` [E,2I,H] -> gate_proj/up_proj (split)
13
+ * ``experts.w2_weight`` [E,H,I] -> down_proj (identity)
14
+ * ``shared_experts.shared_w13`` [2,2I,H] -> gate_proj/up_proj (split)
15
+ * ``model.mtp.*`` dropped (inference-irrelevant)
16
+ * everything else: identity
17
+ """
18
+
19
+ from __future__ import annotations
20
+
21
+ import glob
22
+ import json
23
+ import os
24
+ import shutil
25
+
26
+ import mlx.core as mx
27
+
28
+
29
+ def map_name(name: str):
30
+ """HF checkpoint tensor name -> list of (out_name, kind) for the MLX model."""
31
+ if name.startswith("model.mtp."):
32
+ return [] # drop MTP head
33
+
34
+ if name.endswith(("k_sconv.weight", "v_sconv.weight", "attn_sconv.weight", "mlp_sconv.weight")):
35
+ return [(name, "sconv")]
36
+
37
+ # dense MLP fused gate+up / down
38
+ if name.endswith("mlp.w13_dn.weight"):
39
+ base = name[: -len("w13_dn.weight")]
40
+ return [(base + "gate_proj.weight", "w13_gate"), (base + "up_proj.weight", "w13_up")]
41
+ if name.endswith("mlp.w2_md.weight"):
42
+ return [(name[: -len("w2_md.weight")] + "down_proj.weight", "identity")]
43
+
44
+ # routed experts fused
45
+ if name.endswith("experts.w13_weight"):
46
+ base = name[: -len("w13_weight")]
47
+ return [(base + "gate_proj.weight", "w13_gate"), (base + "up_proj.weight", "w13_up")]
48
+ if name.endswith("experts.w2_weight"):
49
+ return [(name[: -len("w2_weight")] + "down_proj.weight", "identity")]
50
+
51
+ # shared experts fused
52
+ if name.endswith("shared_experts.shared_w13_weight"):
53
+ base = name[: -len("shared_w13_weight")]
54
+ return [(base + "gate_proj.weight", "w13_gate"), (base + "up_proj.weight", "w13_up")]
55
+ if name.endswith("shared_experts.shared_w2_weight"):
56
+ return [(name[: -len("shared_w2_weight")] + "down_proj.weight", "identity")]
57
+
58
+ return [(name, "identity")]
59
+
60
+
61
+ def transform(w: mx.array, kind: str) -> mx.array:
62
+ if kind == "identity":
63
+ return w
64
+ if kind == "sconv":
65
+ # [C, 1, K] -> [C, K, 1]
66
+ return mx.swapaxes(w, 1, 2)
67
+ if kind in ("w13_gate", "w13_up"):
68
+ # The checkpoint stores gate/up INTERLEAVED row-wise: [g0, u0, g1, u1, ...]
69
+ # (SGLang `deinterleave_w13`). De-interleave: gate = rows 0::2, up = rows 1::2.
70
+ # A contiguous [:half]/[half:] split scrambles gate<->up in every MLP.
71
+ n = w.shape[-2] // 2
72
+ g = w.reshape(*w.shape[:-2], n, 2, w.shape[-1])
73
+ return g[..., 0, :] if kind == "w13_gate" else g[..., 1, :]
74
+ raise ValueError(kind)
75
+
76
+
77
+ # ---- quantization target predicate (must be identical in convert and load) ----
78
+
79
+ # Quant "recipes" — which module leaves get affine-quantized.
80
+ # uniform : everything (attention, MLP/experts, embed/unembed, audio, vision)
81
+ # experts_only : ONLY the MLP/expert matmuls (+ audio/vision); attention and
82
+ # embed/unembed stay bf16. Inkling attention dominates 4-bit error
83
+ # (~58% per layer vs ~15% for experts), so this keeps a 4-bit-sized
84
+ # build coherent while the ~927 B experts still fit in 512 GB.
85
+ _RECIPES = {
86
+ "uniform": {"wq_du", "wk_dv", "wv_dv", "wr_du", "wo_ud",
87
+ "gate_proj", "up_proj", "down_proj", "embed", "unembed", "encoder"},
88
+ "experts_only": {"gate_proj", "up_proj", "down_proj", "encoder"},
89
+ }
90
+
91
+
92
+ def is_quant_target(out_name: str, quant_axis_size: int, group_size: int, recipe: str = "uniform") -> bool:
93
+ """Whether ``out_name`` (a converted param path) should be affine-quantized."""
94
+ if not out_name.endswith(".weight"):
95
+ return False
96
+ leaf = out_name[: -len(".weight")].rsplit(".", 1)[-1]
97
+ leaves = _RECIPES[recipe]
98
+ # vision projection layers (linear_0 .. linear_3) — quantized in both recipes
99
+ is_vision_linear = leaf.startswith("linear_") and ".visual." in out_name
100
+ if leaf not in leaves and not is_vision_linear:
101
+ return False
102
+ # router gate stays fp (leaf == "gate", excluded above); norms/sconv excluded by leaf
103
+ # can only group-quantize when the input dim is a multiple of group_size
104
+ return quant_axis_size % group_size == 0
105
+
106
+
107
+ # ------------------------------ streaming driver ------------------------------
108
+
109
+ _SHARD_CAP_BYTES = 5_000_000_000 # ~5 GB per output shard
110
+
111
+
112
+ def _process_tensor(name, w, bits, group_size, out_dtype, recipe="uniform"):
113
+ """Yield (out_name, array) pairs for one source tensor."""
114
+ for out_name, kind in map_name(name):
115
+ wt = transform(w, kind)
116
+ quantize = bits is not None and is_quant_target(out_name, wt.shape[-1], group_size, recipe)
117
+ if quantize:
118
+ qw, scales, biases = mx.quantize(wt, group_size=group_size, bits=bits)
119
+ base = out_name[: -len(".weight")]
120
+ yield out_name, qw
121
+ yield base + ".scales", scales
122
+ yield base + ".biases", biases
123
+ else:
124
+ # keep norms/router/sconv/rel-proj in fp32-safe dtype; matmul weights in out_dtype
125
+ keep_hi = wt.dtype == mx.float32 and (".global_scale" in out_name or ".bias" in out_name
126
+ or out_name.endswith(("_norm.weight", "norm.weight")))
127
+ yield out_name, wt.astype(mx.float32 if keep_hi else out_dtype)
128
+
129
+
130
+ def convert_model(src: str, dst: str, bits=None, group_size: int = 64, out_dtype=mx.bfloat16,
131
+ recipe: str = "uniform"):
132
+ """Stream-convert an Inkling checkpoint from ``src`` to ``dst``.
133
+
134
+ ``bits=None`` -> plain dtype cast (bf16). ``bits in {4,6,8}`` -> affine quant.
135
+ ``recipe`` selects which modules are quantized (see ``_RECIPES``).
136
+ Processes one source shard at a time; never holds the whole model in RAM.
137
+ """
138
+ os.makedirs(dst, exist_ok=True)
139
+ index = json.load(open(os.path.join(src, "model.safetensors.index.json")))
140
+ weight_map = index["weight_map"]
141
+
142
+ shard_to_names: dict[str, list[str]] = {}
143
+ for n, s in weight_map.items():
144
+ shard_to_names.setdefault(s, []).append(n)
145
+
146
+ out_index: dict[str, str] = {}
147
+ buffer: dict[str, mx.array] = {}
148
+ buffer_bytes = 0
149
+ out_shard_id = 0
150
+ total_out_shards_placeholder = "{:05d}"
151
+
152
+ def flush(final=False):
153
+ nonlocal buffer, buffer_bytes, out_shard_id
154
+ if not buffer:
155
+ return
156
+ out_shard_id += 1
157
+ fname = f"model-{total_out_shards_placeholder.format(out_shard_id)}.safetensors"
158
+ mx.save_safetensors(os.path.join(dst, fname), buffer, metadata={"format": "mlx"})
159
+ for k in buffer:
160
+ out_index[k] = fname
161
+ buffer = {}
162
+ buffer_bytes = 0
163
+
164
+ for shard in sorted(shard_to_names):
165
+ path = os.path.join(src, shard)
166
+ tensors = mx.load(path) # mmap
167
+ for name in shard_to_names[shard]:
168
+ w = tensors[name]
169
+ for out_name, arr in _process_tensor(name, w, bits, group_size, out_dtype, recipe):
170
+ mx.eval(arr)
171
+ buffer[out_name] = arr
172
+ buffer_bytes += arr.nbytes
173
+ if buffer_bytes >= _SHARD_CAP_BYTES:
174
+ flush()
175
+ del tensors
176
+ flush(final=True)
177
+
178
+ # rename shards with correct total, build index.json
179
+ _finalize_index(dst, out_index, out_shard_id)
180
+ _write_config(src, dst, bits, group_size, recipe)
181
+ _copy_aux(src, dst)
182
+ return dst
183
+
184
+
185
+ def _finalize_index(dst, out_index, n_shards):
186
+ # rewrite shard filenames to model-XXXXX-of-YYYYY.safetensors
187
+ remap = {}
188
+ for i in range(1, n_shards + 1):
189
+ old = f"model-{i:05d}.safetensors"
190
+ new = f"model-{i:05d}-of-{n_shards:05d}.safetensors"
191
+ if old != new and os.path.exists(os.path.join(dst, old)):
192
+ os.rename(os.path.join(dst, old), os.path.join(dst, new))
193
+ remap[old] = new
194
+ weight_map = {k: remap[v] for k, v in out_index.items()}
195
+ total = sum(os.path.getsize(os.path.join(dst, f)) for f in set(weight_map.values()))
196
+ with open(os.path.join(dst, "model.safetensors.index.json"), "w") as f:
197
+ json.dump({"metadata": {"total_size": total}, "weight_map": weight_map}, f, indent=2)
198
+
199
+
200
+ def _write_config(src, dst, bits, group_size, recipe="uniform"):
201
+ cfg = json.load(open(os.path.join(src, "config.json")))
202
+ if bits is not None:
203
+ cfg["quantization"] = {"group_size": group_size, "bits": bits, "recipe": recipe}
204
+ with open(os.path.join(dst, "config.json"), "w") as f:
205
+ json.dump(cfg, f, indent=2)
206
+
207
+
208
+ def _copy_aux(src, dst):
209
+ for pat in ("tokenizer*", "special_tokens_map.json", "*.tiktoken", "tiktoken",
210
+ "chat_template.jinja", "processor_config.json"):
211
+ for p in glob.glob(os.path.join(src, pat)):
212
+ base = os.path.basename(p)
213
+ target = os.path.join(dst, base)
214
+ if os.path.isdir(p):
215
+ shutil.copytree(p, target, dirs_exist_ok=True)
216
+ else:
217
+ shutil.copy2(p, target)
inkling_mlx/convert_cli.py ADDED
@@ -0,0 +1,44 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """CLI: convert/quantize an Inkling checkpoint to MLX.
2
+
3
+ Examples:
4
+ python -m inkling_mlx.convert_cli --src /path/Inkling-src --dst out-bf16
5
+ python -m inkling_mlx.convert_cli --src /path/Inkling-src --dst out-4bit --bits 4
6
+ """
7
+
8
+ from __future__ import annotations
9
+
10
+ import argparse
11
+ import time
12
+
13
+ import mlx.core as mx
14
+
15
+ from .convert import convert_model
16
+
17
+
18
+ def main():
19
+ ap = argparse.ArgumentParser()
20
+ ap.add_argument("--src", required=True, help="Inkling bf16 source dir (HF layout)")
21
+ ap.add_argument("--dst", required=True, help="output dir")
22
+ ap.add_argument("--bits", type=int, default=None, choices=[2, 3, 4, 5, 6, 8],
23
+ help="quantization bits; omit for bf16 passthrough")
24
+ ap.add_argument("--group-size", type=int, default=64)
25
+ ap.add_argument("--dtype", default="bfloat16", choices=["bfloat16", "float16"])
26
+ ap.add_argument("--device", default="gpu", choices=["gpu", "cpu"],
27
+ help="cpu avoids the Metal GPU-timeout watchdog on huge tensors (slower, robust)")
28
+ ap.add_argument("--recipe", default="uniform", choices=["uniform", "experts_only"],
29
+ help="experts_only keeps attention + embed/unembed at bf16 (coherent 4-bit-sized build)")
30
+ args = ap.parse_args()
31
+
32
+ if args.device == "cpu":
33
+ mx.set_default_device(mx.cpu)
34
+ print("[convert] using CPU device (avoids Metal command-buffer timeout)")
35
+
36
+ dtype = {"bfloat16": mx.bfloat16, "float16": mx.float16}[args.dtype]
37
+ t0 = time.time()
38
+ print(f"[convert] {args.src} -> {args.dst} bits={args.bits} group_size={args.group_size} dtype={args.dtype} recipe={args.recipe}")
39
+ convert_model(args.src, args.dst, bits=args.bits, group_size=args.group_size, out_dtype=dtype, recipe=args.recipe)
40
+ print(f"[convert] done in {time.time()-t0:.0f}s -> {args.dst}")
41
+
42
+
43
+ if __name__ == "__main__":
44
+ main()
inkling_mlx/generate.py ADDED
@@ -0,0 +1,74 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Greedy generation for an Inkling MLX model, using an incremental KV + conv-state
2
+ cache: the prompt is prefilled once, then each new token is a single-position step.
3
+ """
4
+
5
+ from __future__ import annotations
6
+
7
+ import argparse
8
+ import time
9
+
10
+ import mlx.core as mx
11
+
12
+ from .cache import make_cache
13
+ from .load import load
14
+
15
+
16
+ def load_tokenizer(path: str):
17
+ try:
18
+ from transformers import AutoTokenizer
19
+ return AutoTokenizer.from_pretrained(path, trust_remote_code=True)
20
+ except Exception:
21
+ from transformers import PreTrainedTokenizerFast
22
+ import os
23
+ return PreTrainedTokenizerFast(tokenizer_file=os.path.join(path, "tokenizer.json"))
24
+
25
+
26
+ def greedy_generate(model, config, input_ids, max_new_tokens=32, eos_id=None):
27
+ eos_id = eos_id if eos_id is not None else config.eos_token_id
28
+ caches = make_cache(model)
29
+ prompt = list(input_ids)
30
+
31
+ # prefill the whole prompt in one pass; only need the last position's logits
32
+ logits = model(mx.array([prompt]), caches=caches, start_pos=0, last_logit_only=True)
33
+ next_id = int(mx.argmax(logits[0, -1]).item())
34
+ out = [next_id]
35
+ pos = len(prompt)
36
+
37
+ for _ in range(max_new_tokens - 1):
38
+ if next_id == eos_id:
39
+ break
40
+ logits = model(mx.array([[next_id]]), caches=caches, start_pos=pos, last_logit_only=True)
41
+ next_id = int(mx.argmax(logits[0, -1]).item())
42
+ out.append(next_id)
43
+ pos += 1
44
+
45
+ return prompt + out
46
+
47
+
48
+ def main():
49
+ ap = argparse.ArgumentParser()
50
+ ap.add_argument("--model", required=True, help="converted MLX model dir")
51
+ ap.add_argument("--prompt", default="The capital of France is")
52
+ ap.add_argument("--max-new-tokens", type=int, default=32)
53
+ args = ap.parse_args()
54
+
55
+ print(f"[load] {args.model}")
56
+ t0 = time.time()
57
+ model, config = load(args.model)
58
+ print(f"[load] ready in {time.time()-t0:.0f}s")
59
+
60
+ tok = load_tokenizer(args.model)
61
+ input_ids = tok(args.prompt)["input_ids"]
62
+ print(f"[prompt] {args.prompt!r} -> {len(input_ids)} tokens")
63
+
64
+ t0 = time.time()
65
+ out_ids = greedy_generate(model, config, input_ids, args.max_new_tokens)
66
+ dt = time.time() - t0
67
+ text = tok.decode(out_ids)
68
+ n_new = len(out_ids) - len(input_ids)
69
+ print(f"\n{text}\n")
70
+ print(f"[gen] {n_new} tokens in {dt:.1f}s ({n_new/dt:.2f} tok/s)")
71
+
72
+
73
+ if __name__ == "__main__":
74
+ main()
inkling_mlx/layers.py ADDED
@@ -0,0 +1,47 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Inkling decoder layer: attention + MLP, each wrapped by a pre-norm and a
2
+ trailing short-convolution, with residual adds. Mirrors ``InklingDecoderLayer``.
3
+ """
4
+
5
+ from __future__ import annotations
6
+
7
+ import mlx.core as mx
8
+ import mlx.nn as nn
9
+
10
+ from .attention import Attention
11
+ from .common import RMSNorm, ShortConvolution
12
+ from .config import TextConfig
13
+ from .moe import DenseMLP, MoE
14
+
15
+
16
+ class DecoderLayer(nn.Module):
17
+ def __init__(self, config: TextConfig, layer_idx: int):
18
+ super().__init__()
19
+ self.attn = Attention(config, layer_idx)
20
+ self.attn_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
21
+ self.mlp_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
22
+ if config.mlp_layer_types[layer_idx] == "sparse":
23
+ self.mlp = MoE(config)
24
+ else:
25
+ self.mlp = DenseMLP(config)
26
+ self.attn_sconv = ShortConvolution(config.hidden_size, config.sconv_kernel_size)
27
+ self.mlp_sconv = ShortConvolution(config.hidden_size, config.sconv_kernel_size)
28
+
29
+ def __call__(self, x, start_pos=0, cache=None, conv_mask=None):
30
+ kv = cache.kv if cache is not None else None
31
+ residual = x
32
+ h = self.attn_norm(x)
33
+ h = self.attn(
34
+ h, start_pos=start_pos, kv_cache=kv,
35
+ k_conv=cache.k_conv if cache is not None else None,
36
+ v_conv=cache.v_conv if cache is not None else None,
37
+ conv_mask=conv_mask,
38
+ )
39
+ h = self.attn_sconv(h, mask=conv_mask, cache=cache.attn_conv if cache is not None else None)
40
+ x = residual + h
41
+
42
+ residual = x
43
+ h = self.mlp_norm(x)
44
+ h = self.mlp(h)
45
+ h = self.mlp_sconv(h, mask=conv_mask, cache=cache.mlp_conv if cache is not None else None)
46
+ x = residual + h
47
+ return x
inkling_mlx/load.py ADDED
@@ -0,0 +1,69 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Load a (possibly quantized) Inkling MLX model produced by ``convert_model``."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import glob
6
+ import json
7
+ import os
8
+
9
+ import mlx.core as mx
10
+ import mlx.nn as nn
11
+ from mlx.utils import tree_flatten
12
+
13
+ from .config import InklingConfig
14
+ from .model import InklingForConditionalGeneration
15
+
16
+
17
+ def quant_predicate(group_size: int, recipe: str = "uniform"):
18
+ """Quantize exactly the modules the converter did, by delegating to
19
+ ``convert.is_quant_target`` with the same ``recipe``. Guarantees the loaded
20
+ module set matches the checkpoint (e.g. under ``experts_only``, attention and
21
+ embed/unembed stay bf16 and must NOT be re-quantized here)."""
22
+ from .convert import is_quant_target
23
+
24
+ def pred(path, module):
25
+ if not hasattr(module, "to_quantized"):
26
+ return False
27
+ w = getattr(module, "weight", None)
28
+ if w is None:
29
+ return False
30
+ return is_quant_target(path + ".weight", w.shape[-1], group_size, recipe)
31
+
32
+ return pred
33
+
34
+
35
+ def load(path: str, lazy: bool = False):
36
+ cfg_dict = json.load(open(os.path.join(path, "config.json")))
37
+ config = InklingConfig.from_dict(cfg_dict)
38
+ model = InklingForConditionalGeneration(config)
39
+
40
+ q = cfg_dict.get("quantization")
41
+ if q:
42
+ nn.quantize(model, group_size=q["group_size"], bits=q["bits"],
43
+ class_predicate=quant_predicate(q["group_size"], q.get("recipe", "uniform")))
44
+
45
+ # Stream shards: assign each, then release its handle. We do NOT eagerly
46
+ # mx.eval() the whole parameter tree — for a ~500 GB model that builds one
47
+ # enormous eval graph and trips a Metal resource limit. Weights stay lazy
48
+ # (mmap-backed) and materialize on demand during the forward pass, exactly
49
+ # like mlx-lm loads large models.
50
+ loaded = set()
51
+ shards = sorted(glob.glob(os.path.join(path, "*.safetensors")))
52
+ for shard in shards:
53
+ w = mx.load(shard)
54
+ model.load_weights(list(w.items()), strict=False)
55
+ if not lazy:
56
+ # materialize THIS shard's tensors now (bounded graph) and keep them
57
+ # resident. Avoids one enormous eval over all ~500 GB of params, which
58
+ # trips a Metal resource limit; also prevents per-token disk paging.
59
+ mx.eval(list(w.values()))
60
+ loaded.update(w.keys())
61
+ del w
62
+
63
+ expected = {k for k, _ in tree_flatten(model.parameters())}
64
+ missing = expected - loaded
65
+ if missing:
66
+ raise ValueError(f"{len(missing)} params not found in checkpoint, e.g. {sorted(missing)[:3]}")
67
+
68
+ model.eval()
69
+ return model, config
inkling_mlx/model.py ADDED
@@ -0,0 +1,79 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Top-level Inkling multimodal model.
2
+
3
+ Checkpoint layout: ``model.llm.*`` (text backbone + untied unembed), ``model.visual.*``
4
+ (HMLP vision tower), ``model.audio.*`` (dMel audio tower). Image/audio features are
5
+ scattered into the token-embedding stream at their placeholder-token positions, then
6
+ the text backbone runs and the untied unembed head produces (muP-scaled) logits.
7
+ The MTP head (``model.mtp.*``) is intentionally not loaded (inference-irrelevant).
8
+ """
9
+
10
+ from __future__ import annotations
11
+
12
+ import mlx.core as mx
13
+ import mlx.nn as nn
14
+ import numpy as np
15
+
16
+ from .audio import AudioModel
17
+ from .config import InklingConfig
18
+ from .text import TextModel
19
+ from .vision import VisionModel
20
+
21
+
22
+ def _scatter_features(embeds, input_ids, token_id, features):
23
+ """Replace ``embeds`` rows where ``input_ids == token_id`` with ``features``
24
+ (in sequence order). ``input_ids`` is host-known so we resolve positions on CPU."""
25
+ B, L, H = embeds.shape
26
+ ids = np.array(input_ids).reshape(-1)
27
+ pos = np.nonzero(ids == token_id)[0]
28
+ if pos.size == 0:
29
+ return embeds
30
+ flat = embeds.reshape(B * L, H)
31
+ flat[mx.array(pos)] = features.astype(flat.dtype)
32
+ return flat.reshape(B, L, H)
33
+
34
+
35
+ class InnerModel(nn.Module):
36
+ """The ``model.`` level holding the three towers."""
37
+
38
+ def __init__(self, config: InklingConfig):
39
+ super().__init__()
40
+ self.llm = TextModel(config.text)
41
+ self.visual = VisionModel(config.vision)
42
+ self.audio = AudioModel(config.audio)
43
+
44
+
45
+ class InklingForConditionalGeneration(nn.Module):
46
+ def __init__(self, config: InklingConfig):
47
+ super().__init__()
48
+ self.config = config
49
+ self.model = InnerModel(config)
50
+
51
+ # --- convenience accessors ---
52
+ @property
53
+ def llm(self) -> TextModel:
54
+ return self.model.llm
55
+
56
+ def __call__(
57
+ self,
58
+ input_ids: mx.array,
59
+ pixel_values: mx.array | None = None,
60
+ audio_input_ids: mx.array | None = None,
61
+ conv_mask=None,
62
+ caches=None,
63
+ start_pos: int = 0,
64
+ last_logit_only: bool = False,
65
+ ) -> mx.array:
66
+ embeds = self.model.llm.embed_tokens(input_ids)
67
+
68
+ if pixel_values is not None:
69
+ img = self.model.visual(pixel_values)
70
+ embeds = _scatter_features(embeds, input_ids, self.config.image_token_id, img)
71
+
72
+ if audio_input_ids is not None:
73
+ aud = self.model.audio(audio_input_ids)
74
+ embeds = _scatter_features(embeds, input_ids, self.config.audio_token_id, aud)
75
+
76
+ hidden = self.model.llm.backbone(embeds, conv_mask=conv_mask, caches=caches, start_pos=start_pos)
77
+ if last_logit_only:
78
+ hidden = hidden[:, -1:, :]
79
+ return self.model.llm.logits(hidden)
inkling_mlx/moe.py ADDED
@@ -0,0 +1,117 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Inkling MLP variants: dense SwiGLU (with a learned output scale) and the
2
+ sparse MoE (sigmoid router with correction bias, softmax-over-selected weights,
3
+ route/global scaling, and 2 always-on shared experts forming a routing "sink").
4
+
5
+ Mirrors ``InklingMLP`` / ``InklingTopkRouter`` / ``InklingExperts`` /
6
+ ``InklingSharedExperts`` / ``InklingMoE`` from transformers PR #47347.
7
+ """
8
+
9
+ from __future__ import annotations
10
+
11
+ import mlx.core as mx
12
+ import mlx.nn as nn
13
+
14
+ from mlx_lm.models.switch_layers import SwitchGLU
15
+
16
+ from .config import TextConfig
17
+
18
+
19
+ class DenseMLP(nn.Module):
20
+ """SwiGLU MLP with a learned scalar output gain (``global_scale``).
21
+
22
+ The checkpoint fuses gate+up into ``w13_dn``; the converter splits it into
23
+ ``gate_proj``/``up_proj`` so the standard MLX quantizer sees plain ``nn.Linear``s.
24
+ """
25
+
26
+ def __init__(self, config: TextConfig):
27
+ super().__init__()
28
+ h = config.hidden_size
29
+ inter = config.dense_intermediate_size
30
+ self.gate_proj = nn.Linear(h, inter, bias=False)
31
+ self.up_proj = nn.Linear(h, inter, bias=False)
32
+ self.down_proj = nn.Linear(inter, h, bias=False)
33
+ self.global_scale = mx.ones((1,))
34
+
35
+ def __call__(self, x: mx.array) -> mx.array:
36
+ y = self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
37
+ return y * self.global_scale
38
+
39
+
40
+ class Router(nn.Module):
41
+ """Sigmoid top-k router with a correction bias and a shared-expert sink.
42
+
43
+ Kept in full precision (tiny). Returns per-token routed weights/indices plus
44
+ the two shared-expert gammas produced by the same softmax (the "sink").
45
+ """
46
+
47
+ def __init__(self, config: TextConfig):
48
+ super().__init__()
49
+ self.num_experts = config.n_routed_experts
50
+ self.n_shared = config.n_shared_experts
51
+ self.n_total = self.num_experts + self.n_shared
52
+ self.top_k = config.num_experts_per_tok
53
+ self.route_scale = config.route_scale
54
+ self.hidden = config.hidden_size
55
+ self.weight = mx.zeros((self.n_total, config.hidden_size))
56
+ self.bias = mx.zeros((self.num_experts,)) # e_score_correction_bias
57
+ self.global_scale = mx.ones((1,))
58
+
59
+ def __call__(self, x: mx.array):
60
+ # Routing (esp. the top-k selection) is precision-sensitive: in bf16 the
61
+ # rounding of near-tied expert scores flips which experts fire, and a wrong
62
+ # choice compounds over 64 MoE layers into incoherent output. Compute the
63
+ # whole router in fp32.
64
+ flat = x.reshape(-1, self.hidden).astype(mx.float32)
65
+ router_logits = flat @ self.weight.T.astype(mx.float32) # [T, n_total]
66
+ scores = mx.sigmoid(router_logits)
67
+ routed_scores = scores[:, : self.num_experts]
68
+ scores_for_choice = routed_scores + self.bias
69
+
70
+ # top-k experts (order within the top-k is irrelevant downstream)
71
+ topk_idx = mx.argpartition(-scores_for_choice, kth=self.top_k - 1, axis=-1)[:, : self.top_k]
72
+
73
+ routed_logits = router_logits[:, : self.num_experts]
74
+ shared_logits = router_logits[:, self.num_experts :] # [T, n_shared]
75
+ gathered = mx.take_along_axis(routed_logits, topk_idx, axis=-1) # [T, top_k]
76
+ topk_logits = mx.concatenate([gathered, shared_logits], axis=-1) # [T, top_k+n_shared]
77
+
78
+ # softmax over the selected (+shared) logits, computed in the log domain
79
+ log_probs = -mx.logaddexp(mx.zeros_like(topk_logits), -topk_logits) # logsigmoid
80
+ weights = mx.softmax(log_probs, axis=-1)
81
+ weights = weights * self.route_scale * self.global_scale
82
+
83
+ shared_gammas = weights[:, self.top_k :].astype(x.dtype) # [T, n_shared]
84
+ topk_weights = weights[:, : self.top_k].astype(x.dtype) # [T, top_k]
85
+ return topk_weights, topk_idx, shared_gammas
86
+
87
+
88
+ class MoE(nn.Module):
89
+ def __init__(self, config: TextConfig):
90
+ super().__init__()
91
+ self.config = config
92
+ self.n_shared = config.n_shared_experts
93
+ self.gate = Router(config)
94
+ self.experts = SwitchGLU(
95
+ config.hidden_size, config.moe_intermediate_size, config.n_routed_experts, bias=False
96
+ )
97
+ self.shared_experts = SwitchGLU(
98
+ config.hidden_size, config.moe_intermediate_size, config.n_shared_experts, bias=False
99
+ )
100
+
101
+ def __call__(self, x: mx.array) -> mx.array:
102
+ B, L, H = x.shape
103
+ topk_weights, topk_idx, shared_gammas = self.gate(x)
104
+ xf = x.reshape(-1, H) # [T, H]
105
+ T = xf.shape[0]
106
+
107
+ routed = self.experts(xf, topk_idx) # [T, top_k, H]
108
+ routed = (routed * topk_weights[..., None]).sum(axis=1)
109
+
110
+ shared_idx = mx.broadcast_to(
111
+ mx.arange(self.n_shared)[None], (T, self.n_shared)
112
+ )
113
+ shared = self.shared_experts(xf, shared_idx) # [T, n_shared, H]
114
+ shared = (shared.astype(mx.float32) * shared_gammas[..., None].astype(mx.float32)).sum(axis=1)
115
+ shared = shared.astype(routed.dtype)
116
+
117
+ return (routed + shared).reshape(B, L, H)
inkling_mlx/text.py ADDED
@@ -0,0 +1,50 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Inkling text backbone (``model.llm.*``): token embedding + embed-norm,
2
+ 66 decoder layers, final norm, and the (untied) unembed head.
3
+
4
+ Mirrors ``InklingTextModel`` + the unembed / muP-logit scaling from
5
+ ``InklingForConditionalGeneration``.
6
+ """
7
+
8
+ from __future__ import annotations
9
+
10
+ import mlx.core as mx
11
+ import mlx.nn as nn
12
+
13
+ from .common import RMSNorm
14
+ from .config import TextConfig
15
+ from .layers import DecoderLayer
16
+
17
+
18
+ class TextModel(nn.Module):
19
+ def __init__(self, config: TextConfig):
20
+ super().__init__()
21
+ self.config = config
22
+ self.embed = nn.Embedding(config.vocab_size, config.hidden_size)
23
+ self.embed_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
24
+ self.layers = [DecoderLayer(config, i) for i in range(config.num_hidden_layers)]
25
+ self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
26
+ self.unembed = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
27
+
28
+ def embed_tokens(self, input_ids: mx.array) -> mx.array:
29
+ return self.embed_norm(self.embed(input_ids))
30
+
31
+ def backbone(self, inputs_embeds: mx.array, conv_mask=None, caches=None, start_pos=0) -> mx.array:
32
+ h = inputs_embeds
33
+ for i, layer in enumerate(self.layers):
34
+ h = layer(h, start_pos=start_pos,
35
+ cache=caches[i] if caches is not None else None,
36
+ conv_mask=conv_mask)
37
+ return self.norm(h)
38
+
39
+ def logits(self, hidden: mx.array) -> mx.array:
40
+ hidden = hidden / self.config.logits_mup_width_multiplier
41
+ logits = self.unembed(hidden)
42
+ uv = self.config.unpadded_vocab_size
43
+ if uv is not None and uv < logits.shape[-1]:
44
+ logits = logits[..., :uv]
45
+ return logits
46
+
47
+ def __call__(self, input_ids: mx.array, conv_mask=None) -> mx.array:
48
+ h = self.embed_tokens(input_ids)
49
+ h = self.backbone(h, conv_mask=conv_mask)
50
+ return self.logits(h)
inkling_mlx/vision.py ADDED
@@ -0,0 +1,126 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Inkling vision tower: an HMLP (hierarchical MLP) patch encoder.
2
+
3
+ No attention — each layer folds space/time into the channel dim then projects
4
+ (Linear -> RMSNorm -> GELU), progressively growing channels up to the text hidden
5
+ size. Mirrors ``InklingVisionModel`` / ``InklingVisionEncoderLayer`` /
6
+ ``plan_out_scales``. Checkpoint keys are flat: ``visual.layers.linear_{i}`` and
7
+ ``visual.layers.norm_{i}`` plus ``visual.final_norm``.
8
+ """
9
+
10
+ from __future__ import annotations
11
+
12
+ import math
13
+
14
+ import mlx.core as mx
15
+ import mlx.nn as nn
16
+
17
+ from .common import RMSNorm
18
+ from .config import VisionConfig
19
+
20
+
21
+ def _prime_factors(n: int) -> list[int]:
22
+ factors = []
23
+ while n % 2 == 0:
24
+ factors.append(2)
25
+ n //= 2
26
+ p = 3
27
+ while p * p <= n:
28
+ while n % p == 0:
29
+ factors.append(p)
30
+ n //= p
31
+ p += 2
32
+ if n > 1:
33
+ factors.append(n)
34
+ return factors
35
+
36
+
37
+ def plan_out_scales(temporal_patch_size: int, patch_size: int, n_layers: int, n_channels: int):
38
+ """Port of the reference ``plan_out_scales`` (returns an ``(n_layers+1, 4)``
39
+ array of (t, h, w, c) grid sizes). Uses numpy + scipy for the assignment."""
40
+ import numpy as np
41
+ from scipy.optimize import linear_sum_assignment
42
+
43
+ h = np.cumprod(np.array(_prime_factors(patch_size)[::-1]))
44
+ t = np.cumprod(np.array(_prime_factors(temporal_patch_size)[::-1]))
45
+
46
+ h_ch = np.ceil(h**2 * n_channels / 64).astype(np.int64) * 64
47
+ t_ch = (np.ceil(h[-1] ** 2 * n_channels * t)).astype(np.int64) * 64
48
+
49
+ base = np.array([[1, 1, 1, n_channels]], dtype=np.int64)
50
+ spatial = np.stack([np.ones_like(h), h, h, h_ch], axis=1)
51
+ temporal = np.stack([t, np.full_like(t, h[-1]), np.full_like(t, h[-1]), t_ch], axis=1)
52
+ scales = np.concatenate([base, spatial, temporal], axis=0).astype(np.int64)
53
+
54
+ size_reduction = np.prod(scales[:, :-1], axis=1).astype(np.float64)
55
+ total_elements = patch_size * patch_size * temporal_patch_size * n_channels
56
+ log_ideal = np.linspace(0.0, math.log(total_elements), n_layers + 1)
57
+ cost = np.abs(log_ideal[:, None] - np.log(size_reduction)[None, :])
58
+
59
+ if n_layers >= scales.shape[0]:
60
+ idxs = np.argmin(cost, axis=1)
61
+ else:
62
+ _, idxs = linear_sum_assignment(cost)
63
+ idxs = np.array(idxs)
64
+ idxs[0] = 0
65
+ idxs[-1] = scales.shape[0] - 1
66
+ return scales[idxs]
67
+
68
+
69
+ def _fold_timespace_to_depth(x, t_fold, hw_fold):
70
+ # x: [B, T, H, W, C] -> [B, T//t, H//hw, W//hw, C*t*hw*hw]
71
+ B, T, H, W, C = x.shape
72
+ t_new, h_new, w_new = T // t_fold, H // hw_fold, W // hw_fold
73
+ x = x.reshape(B, t_new, t_fold, h_new, hw_fold, w_new, hw_fold, C)
74
+ x = x.transpose(0, 1, 3, 5, 2, 4, 6, 7)
75
+ x = x.reshape(B, t_new, h_new, w_new, t_fold * hw_fold * hw_fold * C)
76
+ return x
77
+
78
+
79
+ class _VisionLayers(nn.Module):
80
+ """Holds ``linear_{i}`` / ``norm_{i}`` to match checkpoint keys."""
81
+
82
+ def __init__(self, config: VisionConfig):
83
+ super().__init__()
84
+ scales = plan_out_scales(
85
+ config.temporal_patch_size, config.patch_size, config.n_layers, config.num_channels
86
+ )
87
+ self.n_layers = config.n_layers
88
+ self.folds = [] # (t_fold, hw_fold, add_norm)
89
+ for i in range(config.n_layers):
90
+ start, end = scales[i], scales[i + 1]
91
+ shuffle = (
92
+ (end[0] // start[0]) * (end[1] // start[1]) * (end[2] // start[2])
93
+ )
94
+ hw_fold = int(end[1] // start[1])
95
+ t_fold = int(end[0] // start[0])
96
+ in_dim = int(start[3]) * int(shuffle)
97
+ add_norm = i != config.n_layers - 1
98
+ out_dim = config.text_hidden_size if i == config.n_layers - 1 else int(end[3])
99
+ setattr(self, f"linear_{i}", nn.Linear(in_dim, out_dim, bias=False))
100
+ if add_norm:
101
+ setattr(self, f"norm_{i}", RMSNorm(out_dim, eps=config.rms_norm_eps))
102
+ self.folds.append((t_fold, hw_fold, add_norm))
103
+
104
+ def __call__(self, x):
105
+ for i, (t_fold, hw_fold, add_norm) in enumerate(self.folds):
106
+ if hw_fold > 1 or t_fold > 1:
107
+ x = _fold_timespace_to_depth(x, t_fold, hw_fold)
108
+ x = getattr(self, f"linear_{i}")(x)
109
+ if add_norm:
110
+ x = getattr(self, f"norm_{i}")(x)
111
+ x = nn.gelu(x)
112
+ return x
113
+
114
+
115
+ class VisionModel(nn.Module):
116
+ def __init__(self, config: VisionConfig):
117
+ super().__init__()
118
+ self.config = config
119
+ self.layers = _VisionLayers(config)
120
+ self.final_norm = RMSNorm(config.text_hidden_size, eps=config.rms_norm_eps)
121
+
122
+ def __call__(self, pixel_values: mx.array) -> mx.array:
123
+ num_patches = pixel_values.shape[0]
124
+ h = self.layers(pixel_values)
125
+ h = self.final_norm(h)
126
+ return h.reshape(num_patches, -1)
model-00013-of-00097.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:347bd0e9a78592de1d8b284c8ed730e7df889c6774dfefb3279b999ba1593b17
3
+ size 5435818784
model-00015-of-00097.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:3bd2902330c2cb765b0b099b03b422bf9772ef8ada1a88dd9e509a0bbcbb4005
3
+ size 5435847671
model-00016-of-00097.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:2f44c9b25d12e42fe706bea4b509a47260a2746f1f0a110aa8b8f09085215536
3
+ size 5570459254
model-00019-of-00097.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:18d1963fa26a1370c94da1d9f8a99f69e2666b03849c73bb96547ae029f80387
3
+ size 5552268696
model-00021-of-00097.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:ea283c381a817ac92bbd3cfac3a2c0bac073d2722822ecf1b002dc9ddfd704b0
3
+ size 5439436602
model-00022-of-00097.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:55fd2826eec30ba0504695fa5c3bd558ee497596d1ced2c38cdb5ff3c7ecb26b
3
+ size 5481892886
model-00024-of-00097.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:bdd9110a691f77120d80661a45f56b010f97cd873fe1a782df50853746065225
3
+ size 5555889618
model-00027-of-00097.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:ff814a774c46f077ae1d52a6a479d2ae7141e86c3e4e09ef7209dd24b7bedfe6
3
+ size 5435835282
model-00035-of-00097.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:74f166e9ef09da8cd1a1de091c48cd324f38b01d2828b485cb0424bf3c197542
3
+ size 5478500966
model-00036-of-00097.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:7413fe9ef56f0d6056d428ac56b7cd5259b3d0ffce3507d70b883303906ea95f
3
+ size 5435828491
model-00051-of-00097.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:fa53f3839b888da3993e042447a1ed8ae955696ade3547618b74d007e0c48b49
3
+ size 5510226779
model-00052-of-00097.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:10d628557229be2340453f26644e33f223a342714f78399e1764c8bc75b283fb
3
+ size 5669450472
model-00054-of-00097.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:a39fb3ed64642e61729ac8a3486349aa05fa04661fba2c97c96cdf259e267ab5
3
+ size 5474428577
model-00057-of-00097.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:6d3ad88cc0c3c7097ba4d9c9c0d28f7c091b0f901cfe2eb30a4d48ddb8dfa343
3
+ size 5470856882
model-00058-of-00097.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:eb6cc23978a5a517d93a5e7eb516d09ba0b316c1fa067b95fea3effc4cdf3ac1
3
+ size 5464130796
model-00060-of-00097.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:fb6b9e197466046f11fd92fa22a0bc10c0c2bc13710595c3331f8d7f42547de4
3
+ size 5481474960
model-00063-of-00097.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:529ace358426464006df9a322f0a5ff08a67d10d439ff31ef5ec7c534dee0d01
3
+ size 5574017292
model-00065-of-00097.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:2ba5447632352096bd4479ff52b9549942d758c084392a5b53291157d11b27b4
3
+ size 5474815047
model-00069-of-00097.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:3edb6f874628d7d042b55a9936aaae7f40c2ea5684de45d2f352bbb011ef759a
3
+ size 5602111647
model-00074-of-00097.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:d867d3a5e88dd3acf2d382d90b1fbf9d87bb876b9ef0fa1266e3c9b9b1c21182
3
+ size 5492505269
model-00077-of-00097.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:5aa5ee38350ae4ee8d1fbc8140489973fad2c483cd80a31ef70887bd4bc3f90a
3
+ size 5474462387
model-00082-of-00097.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:3009eb89302944b671fdc045b8613e8786ea90f0753ed93681108a70c42c8fed
3
+ size 5573474243
model-00084-of-00097.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:1fc752a1868f0655b468806329fe95aabb2e1869fa95ae9d34ff15ca92df6b3d
3
+ size 5591306886
model-00090-of-00097.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:508879fcdd31f023915301928b2b47c59c20f9c43dc54b83d35b0a11adb0f066
3
+ size 5435831912
model-00093-of-00097.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:ce1888d0339f7409619d799e335112c3552a20f16edb93d7533286ab5cac8047
3
+ size 5523236292
model-00095-of-00097.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:c2f2b7b4b21d09e46d6e975e8eb7dd25d8cb434a98c47980ba6e9ae7dda15860
3
+ size 5439374554
model-00096-of-00097.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:4c40dbf5ce79741cdf82325013ac23295267a8e53384bc4ff25a74501d78e8aa
3
+ size 5492442567
model.safetensors.index.json ADDED
The diff for this file is too large to render. See raw diff
 
processor_config.json ADDED
@@ -0,0 +1,46 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "audio_token": "<|unused_200053|>",
3
+ "audio_bos_token": "<|content_audio_input|>",
4
+ "dmel_max_value": 2.0,
5
+ "dmel_min_value": -7.0,
6
+ "feature_extractor": {
7
+ "audio_token_duration_s": 0.05,
8
+ "feature_extractor_type": "InklingFeatureExtractor",
9
+ "feature_size": 80,
10
+ "hop_length": 800,
11
+ "n_fft": 1600,
12
+ "padding_side": "right",
13
+ "padding_value": 0.0,
14
+ "return_attention_mask": true,
15
+ "sampling_rate": 16000,
16
+ "window_size": 1600,
17
+ "window_size_multiplier": 2.0
18
+ },
19
+ "image_processor": {
20
+ "do_convert_rgb": true,
21
+ "do_normalize": true,
22
+ "do_rescale": true,
23
+ "do_resize": true,
24
+ "image_mean": [
25
+ 0.48145466,
26
+ 0.4578275,
27
+ 0.40821073
28
+ ],
29
+ "image_processor_type": "InklingImageProcessor",
30
+ "image_std": [
31
+ 0.26862954,
32
+ 0.26130258,
33
+ 0.27577711
34
+ ],
35
+ "resample": 3,
36
+ "rescale_factor": 0.00392156862745098,
37
+ "size": {
38
+ "height": 40,
39
+ "width": 40
40
+ }
41
+ },
42
+ "image_token": "<|unused_200054|>",
43
+ "image_bos_token": "<|content_image|>",
44
+ "num_dmel_bins": 16,
45
+ "processor_class": "InklingProcessor"
46
+ }
special_tokens_map.json ADDED
@@ -0,0 +1,22 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "additional_special_tokens": [
3
+ "<|endoftext|>",
4
+ "<|message_user|>",
5
+ "<|message_model|>",
6
+ "<|message_system|>",
7
+ "<|message_tool|>",
8
+ "<|content_text|>",
9
+ "<|content_image|>",
10
+ "<|content_model_end_sampling|>",
11
+ "<|content_thinking|>",
12
+ "<|end_message|>",
13
+ "<|content_audio_input|>",
14
+ "<|content_tool_error|>",
15
+ "<|audio|>",
16
+ "<|content_xml|>",
17
+ "<|begin_of_text|>",
18
+ "<|audio_end|>",
19
+ "<|content_invoke_tool_json|>",
20
+ "<|content_invoke_tool_text|>"
21
+ ]
22
+ }
tiktoken/tokenizer.model ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:bc253fd2b702f7a6da7105eaa8f3463b2f1247e83614f23e5323b921088bed2a
3
+ size 3615874
tokenizer_config.json ADDED
@@ -0,0 +1,508 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "added_tokens_decoder": {
3
+ "199998": {
4
+ "content": "<|unused|>",
5
+ "single_word": false,
6
+ "lstrip": false,
7
+ "rstrip": false,
8
+ "normalized": false,
9
+ "special": true
10
+ },
11
+ "199999": {
12
+ "content": "<|endoftext|>",
13
+ "single_word": false,
14
+ "lstrip": false,
15
+ "rstrip": false,
16
+ "normalized": false,
17
+ "special": true
18
+ },
19
+ "200000": {
20
+ "content": "<|message_user|>",
21
+ "single_word": false,
22
+ "lstrip": false,
23
+ "rstrip": false,
24
+ "normalized": false,
25
+ "special": true
26
+ },
27
+ "200001": {
28
+ "content": "<|message_model|>",
29
+ "single_word": false,
30
+ "lstrip": false,
31
+ "rstrip": false,
32
+ "normalized": false,
33
+ "special": true
34
+ },
35
+ "200002": {
36
+ "content": "<|message_system|>",
37
+ "single_word": false,
38
+ "lstrip": false,
39
+ "rstrip": false,
40
+ "normalized": false,
41
+ "special": true
42
+ },
43
+ "200003": {
44
+ "content": "<|message_tool|>",
45
+ "single_word": false,
46
+ "lstrip": false,
47
+ "rstrip": false,
48
+ "normalized": false,
49
+ "special": true
50
+ },
51
+ "200004": {
52
+ "content": "<|content_text|>",
53
+ "single_word": false,
54
+ "lstrip": false,
55
+ "rstrip": false,
56
+ "normalized": false,
57
+ "special": true
58
+ },
59
+ "200005": {
60
+ "content": "<|content_image|>",
61
+ "single_word": false,
62
+ "lstrip": false,
63
+ "rstrip": false,
64
+ "normalized": false,
65
+ "special": true
66
+ },
67
+ "200006": {
68
+ "content": "<|content_model_end_sampling|>",
69
+ "single_word": false,
70
+ "lstrip": false,
71
+ "rstrip": false,
72
+ "normalized": false,
73
+ "special": true
74
+ },
75
+ "200007": {
76
+ "content": "<|unused_200007|>",
77
+ "single_word": false,
78
+ "lstrip": false,
79
+ "rstrip": false,
80
+ "normalized": false,
81
+ "special": true
82
+ },
83
+ "200008": {
84
+ "content": "<|content_thinking|>",
85
+ "single_word": false,
86
+ "lstrip": false,
87
+ "rstrip": false,
88
+ "normalized": false,
89
+ "special": true
90
+ },
91
+ "200009": {
92
+ "content": "<|unused_200009|>",
93
+ "single_word": false,
94
+ "lstrip": false,
95
+ "rstrip": false,
96
+ "normalized": false,
97
+ "special": true
98
+ },
99
+ "200010": {
100
+ "content": "<|end_message|>",
101
+ "single_word": false,
102
+ "lstrip": false,
103
+ "rstrip": false,
104
+ "normalized": false,
105
+ "special": true
106
+ },
107
+ "200011": {
108
+ "content": "<|unused_200011|>",
109
+ "single_word": false,
110
+ "lstrip": false,
111
+ "rstrip": false,
112
+ "normalized": false,
113
+ "special": true
114
+ },
115
+ "200012": {
116
+ "content": "<|unused_200012|>",
117
+ "single_word": false,
118
+ "lstrip": false,
119
+ "rstrip": false,
120
+ "normalized": false,
121
+ "special": true
122
+ },
123
+ "200013": {
124
+ "content": "<|unused_200013|>",
125
+ "single_word": false,
126
+ "lstrip": false,
127
+ "rstrip": false,
128
+ "normalized": false,
129
+ "special": true
130
+ },
131
+ "200014": {
132
+ "content": "<|unused_200014|>",
133
+ "single_word": false,
134
+ "lstrip": false,
135
+ "rstrip": false,
136
+ "normalized": false,
137
+ "special": true
138
+ },
139
+ "200015": {
140
+ "content": "<|unused_200015|>",
141
+ "single_word": false,
142
+ "lstrip": false,
143
+ "rstrip": false,
144
+ "normalized": false,
145
+ "special": true
146
+ },
147
+ "200016": {
148
+ "content": "<|unused_200016|>",
149
+ "single_word": false,
150
+ "lstrip": false,
151
+ "rstrip": false,
152
+ "normalized": false,
153
+ "special": true
154
+ },
155
+ "200017": {
156
+ "content": "<|unused_200017|>",
157
+ "single_word": false,
158
+ "lstrip": false,
159
+ "rstrip": false,
160
+ "normalized": false,
161
+ "special": true
162
+ },
163
+ "200018": {
164
+ "content": "<|unused_200018|>",
165
+ "single_word": false,
166
+ "lstrip": false,
167
+ "rstrip": false,
168
+ "normalized": false,
169
+ "special": true
170
+ },
171
+ "200019": {
172
+ "content": "<|unused_200019|>",
173
+ "single_word": false,
174
+ "lstrip": false,
175
+ "rstrip": false,
176
+ "normalized": false,
177
+ "special": true
178
+ },
179
+ "200020": {
180
+ "content": "<|content_audio_input|>",
181
+ "single_word": false,
182
+ "lstrip": false,
183
+ "rstrip": false,
184
+ "normalized": false,
185
+ "special": true
186
+ },
187
+ "200021": {
188
+ "content": "<|unused_200021|>",
189
+ "single_word": false,
190
+ "lstrip": false,
191
+ "rstrip": false,
192
+ "normalized": false,
193
+ "special": true
194
+ },
195
+ "200022": {
196
+ "content": "<|content_tool_error|>",
197
+ "single_word": false,
198
+ "lstrip": false,
199
+ "rstrip": false,
200
+ "normalized": false,
201
+ "special": true
202
+ },
203
+ "200023": {
204
+ "content": "<|audio|>",
205
+ "single_word": false,
206
+ "lstrip": false,
207
+ "rstrip": false,
208
+ "normalized": false,
209
+ "special": true
210
+ },
211
+ "200024": {
212
+ "content": "<|content_xml|>",
213
+ "single_word": false,
214
+ "lstrip": false,
215
+ "rstrip": false,
216
+ "normalized": false,
217
+ "special": true
218
+ },
219
+ "200025": {
220
+ "content": "<|unused_200025|>",
221
+ "single_word": false,
222
+ "lstrip": false,
223
+ "rstrip": false,
224
+ "normalized": false,
225
+ "special": true
226
+ },
227
+ "200026": {
228
+ "content": "<|unused_200026|>",
229
+ "single_word": false,
230
+ "lstrip": false,
231
+ "rstrip": false,
232
+ "normalized": false,
233
+ "special": true
234
+ },
235
+ "200027": {
236
+ "content": "<|unused_200027|>",
237
+ "single_word": false,
238
+ "lstrip": false,
239
+ "rstrip": false,
240
+ "normalized": false,
241
+ "special": true
242
+ },
243
+ "200028": {
244
+ "content": "<|begin_of_text|>",
245
+ "single_word": false,
246
+ "lstrip": false,
247
+ "rstrip": false,
248
+ "normalized": false,
249
+ "special": true
250
+ },
251
+ "200029": {
252
+ "content": "<|unused_200029|>",
253
+ "single_word": false,
254
+ "lstrip": false,
255
+ "rstrip": false,
256
+ "normalized": false,
257
+ "special": true
258
+ },
259
+ "200030": {
260
+ "content": "<|unused_200030|>",
261
+ "single_word": false,
262
+ "lstrip": false,
263
+ "rstrip": false,
264
+ "normalized": false,
265
+ "special": true
266
+ },
267
+ "200031": {
268
+ "content": "<|unused_200031|>",
269
+ "single_word": false,
270
+ "lstrip": false,
271
+ "rstrip": false,
272
+ "normalized": false,
273
+ "special": true
274
+ },
275
+ "200032": {
276
+ "content": "<|unused_200032|>",
277
+ "single_word": false,
278
+ "lstrip": false,
279
+ "rstrip": false,
280
+ "normalized": false,
281
+ "special": true
282
+ },
283
+ "200033": {
284
+ "content": "<|unused_200033|>",
285
+ "single_word": false,
286
+ "lstrip": false,
287
+ "rstrip": false,
288
+ "normalized": false,
289
+ "special": true
290
+ },
291
+ "200034": {
292
+ "content": "<|unused_200034|>",
293
+ "single_word": false,
294
+ "lstrip": false,
295
+ "rstrip": false,
296
+ "normalized": false,
297
+ "special": true
298
+ },
299
+ "200035": {
300
+ "content": "<|unused_200035|>",
301
+ "single_word": false,
302
+ "lstrip": false,
303
+ "rstrip": false,
304
+ "normalized": false,
305
+ "special": true
306
+ },
307
+ "200036": {
308
+ "content": "<|unused_200036|>",
309
+ "single_word": false,
310
+ "lstrip": false,
311
+ "rstrip": false,
312
+ "normalized": false,
313
+ "special": true
314
+ },
315
+ "200037": {
316
+ "content": "<|unused_200037|>",
317
+ "single_word": false,
318
+ "lstrip": false,
319
+ "rstrip": false,
320
+ "normalized": false,
321
+ "special": true
322
+ },
323
+ "200038": {
324
+ "content": "<|unused_200038|>",
325
+ "single_word": false,
326
+ "lstrip": false,
327
+ "rstrip": false,
328
+ "normalized": false,
329
+ "special": true
330
+ },
331
+ "200039": {
332
+ "content": "<|unused_200039|>",
333
+ "single_word": false,
334
+ "lstrip": false,
335
+ "rstrip": false,
336
+ "normalized": false,
337
+ "special": true
338
+ },
339
+ "200040": {
340
+ "content": "<|unused_200040|>",
341
+ "single_word": false,
342
+ "lstrip": false,
343
+ "rstrip": false,
344
+ "normalized": false,
345
+ "special": true
346
+ },
347
+ "200041": {
348
+ "content": "<|unused_200041|>",
349
+ "single_word": false,
350
+ "lstrip": false,
351
+ "rstrip": false,
352
+ "normalized": false,
353
+ "special": true
354
+ },
355
+ "200042": {
356
+ "content": "<|unused_200042|>",
357
+ "single_word": false,
358
+ "lstrip": false,
359
+ "rstrip": false,
360
+ "normalized": false,
361
+ "special": true
362
+ },
363
+ "200043": {
364
+ "content": "<|audio_end|>",
365
+ "single_word": false,
366
+ "lstrip": false,
367
+ "rstrip": false,
368
+ "normalized": false,
369
+ "special": true
370
+ },
371
+ "200044": {
372
+ "content": "<|unused_200044|>",
373
+ "single_word": false,
374
+ "lstrip": false,
375
+ "rstrip": false,
376
+ "normalized": false,
377
+ "special": true
378
+ },
379
+ "200045": {
380
+ "content": "<|unused_200045|>",
381
+ "single_word": false,
382
+ "lstrip": false,
383
+ "rstrip": false,
384
+ "normalized": false,
385
+ "special": true
386
+ },
387
+ "200046": {
388
+ "content": "<|unused_200046|>",
389
+ "single_word": false,
390
+ "lstrip": false,
391
+ "rstrip": false,
392
+ "normalized": false,
393
+ "special": true
394
+ },
395
+ "200047": {
396
+ "content": "<|unused_200047|>",
397
+ "single_word": false,
398
+ "lstrip": false,
399
+ "rstrip": false,
400
+ "normalized": false,
401
+ "special": true
402
+ },
403
+ "200048": {
404
+ "content": "<|unused_200048|>",
405
+ "single_word": false,
406
+ "lstrip": false,
407
+ "rstrip": false,
408
+ "normalized": false,
409
+ "special": true
410
+ },
411
+ "200049": {
412
+ "content": "<|content_invoke_tool_json|>",
413
+ "single_word": false,
414
+ "lstrip": false,
415
+ "rstrip": false,
416
+ "normalized": false,
417
+ "special": true
418
+ },
419
+ "200050": {
420
+ "content": "<|unused_200050|>",
421
+ "single_word": false,
422
+ "lstrip": false,
423
+ "rstrip": false,
424
+ "normalized": false,
425
+ "special": true
426
+ },
427
+ "200051": {
428
+ "content": "<|unused_200051|>",
429
+ "single_word": false,
430
+ "lstrip": false,
431
+ "rstrip": false,
432
+ "normalized": false,
433
+ "special": true
434
+ },
435
+ "200052": {
436
+ "content": "<|unused_200052|>",
437
+ "single_word": false,
438
+ "lstrip": false,
439
+ "rstrip": false,
440
+ "normalized": false,
441
+ "special": true
442
+ },
443
+ "200053": {
444
+ "content": "<|unused_200053|>",
445
+ "single_word": false,
446
+ "lstrip": false,
447
+ "rstrip": false,
448
+ "normalized": false,
449
+ "special": true
450
+ },
451
+ "200054": {
452
+ "content": "<|unused_200054|>",
453
+ "single_word": false,
454
+ "lstrip": false,
455
+ "rstrip": false,
456
+ "normalized": false,
457
+ "special": true
458
+ },
459
+ "200055": {
460
+ "content": "<|unused_200055|>",
461
+ "single_word": false,
462
+ "lstrip": false,
463
+ "rstrip": false,
464
+ "normalized": false,
465
+ "special": true
466
+ },
467
+ "200056": {
468
+ "content": "<|unused_200056|>",
469
+ "single_word": false,
470
+ "lstrip": false,
471
+ "rstrip": false,
472
+ "normalized": false,
473
+ "special": true
474
+ },
475
+ "200057": {
476
+ "content": "<|content_invoke_tool_text|>",
477
+ "single_word": false,
478
+ "lstrip": false,
479
+ "rstrip": false,
480
+ "normalized": false,
481
+ "special": true
482
+ }
483
+ },
484
+ "clean_up_tokenization_spaces": false,
485
+ "fix_mistral_regex": false,
486
+ "extra_special_tokens": {
487
+ "endoftext": "<|endoftext|>",
488
+ "message_user": "<|message_user|>",
489
+ "message_model": "<|message_model|>",
490
+ "message_system": "<|message_system|>",
491
+ "message_tool": "<|message_tool|>",
492
+ "content_text": "<|content_text|>",
493
+ "content_image": "<|content_image|>",
494
+ "content_model_end_sampling": "<|content_model_end_sampling|>",
495
+ "content_thinking": "<|content_thinking|>",
496
+ "end_message": "<|end_message|>",
497
+ "content_audio_input": "<|content_audio_input|>",
498
+ "content_tool_error": "<|content_tool_error|>",
499
+ "audio": "<|audio|>",
500
+ "content_xml": "<|content_xml|>",
501
+ "begin_of_text": "<|begin_of_text|>",
502
+ "audio_end": "<|audio_end|>",
503
+ "content_invoke_tool_json": "<|content_invoke_tool_json|>",
504
+ "content_invoke_tool_text": "<|content_invoke_tool_text|>"
505
+ },
506
+ "model_max_length": 1000000000000000019884624838656,
507
+ "tokenizer_class": "PreTrainedTokenizerFast"
508
+ }