Image-Text-to-Text
MLX
Safetensors
inkling_mm_model
Mixture of Experts
multimodal
inkling
thinking-machines
conversational
Instructions to use pipenetwork/Inkling-MLX-4bit with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- MLX
How to use pipenetwork/Inkling-MLX-4bit with MLX:
# Make sure mlx-vlm is installed # pip install --upgrade mlx-vlm from mlx_vlm import load, generate from mlx_vlm.prompt_utils import apply_chat_template from mlx_vlm.utils import load_config # Load the model model, processor = load("pipenetwork/Inkling-MLX-4bit") config = load_config("pipenetwork/Inkling-MLX-4bit") # Prepare input image = ["http://images.cocodataset.org/val2017/000000039769.jpg"] prompt = "Describe this image." # Apply chat template formatted_prompt = apply_chat_template( processor, config, prompt, num_images=1 ) # Generate output output = generate(model, processor, formatted_prompt, image) print(output) - Notebooks
- Google Colab
- Kaggle
- Local Apps Settings
- LM Studio
- Pi
How to use pipenetwork/Inkling-MLX-4bit with Pi:
Start the MLX server
# Install MLX LM: uv tool install mlx-lm # Start a local OpenAI-compatible server: mlx_lm.server --model "pipenetwork/Inkling-MLX-4bit"
Configure the model in Pi
# Install Pi: npm install -g @mariozechner/pi-coding-agent # Add to ~/.pi/agent/models.json: { "providers": { "mlx-lm": { "baseUrl": "http://localhost:8080/v1", "api": "openai-completions", "apiKey": "none", "models": [ { "id": "pipenetwork/Inkling-MLX-4bit" } ] } } }Run Pi
# Start Pi in your project directory: pi
- Hermes Agent new
How to use pipenetwork/Inkling-MLX-4bit with Hermes Agent:
Start the MLX server
# Install MLX LM: uv tool install mlx-lm # Start a local OpenAI-compatible server: mlx_lm.server --model "pipenetwork/Inkling-MLX-4bit"
Configure Hermes
# Install Hermes: curl -fsSL https://hermes-agent.nousresearch.com/install.sh | bash hermes setup # Point Hermes at the local server: hermes config set model.provider custom hermes config set model.base_url http://127.0.0.1:8080/v1 hermes config set model.default pipenetwork/Inkling-MLX-4bit
Run Hermes
hermes
- OpenClaw new
How to use pipenetwork/Inkling-MLX-4bit with OpenClaw:
Start the MLX server
# Install MLX LM: uv tool install mlx-lm # Start a local OpenAI-compatible server: mlx_lm.server --model "pipenetwork/Inkling-MLX-4bit"
Configure OpenClaw
# Install OpenClaw: npm install -g openclaw@latest # Register the local server and set it as the default model: openclaw onboard --non-interactive --mode local \ --auth-choice custom-api-key \ --custom-base-url http://127.0.0.1:8080/v1 \ --custom-model-id "pipenetwork/Inkling-MLX-4bit" \ --custom-provider-id mlx-lm \ --custom-compatibility openai \ --custom-text-input \ --accept-risk \ --skip-health
Run OpenClaw
openclaw agent --local --agent main --message "Hello from Hugging Face"
Add files using upload-large-folder tool
Browse files- README.md +73 -0
- chat_template.jinja +129 -0
- config.json +139 -0
- inkling_mlx/__init__.py +14 -0
- inkling_mlx/attention.py +135 -0
- inkling_mlx/audio.py +34 -0
- inkling_mlx/cache.py +66 -0
- inkling_mlx/common.py +63 -0
- inkling_mlx/config.py +207 -0
- inkling_mlx/convert.py +217 -0
- inkling_mlx/convert_cli.py +44 -0
- inkling_mlx/generate.py +74 -0
- inkling_mlx/layers.py +47 -0
- inkling_mlx/load.py +69 -0
- inkling_mlx/model.py +79 -0
- inkling_mlx/moe.py +117 -0
- inkling_mlx/text.py +50 -0
- inkling_mlx/vision.py +126 -0
- model-00013-of-00097.safetensors +3 -0
- model-00015-of-00097.safetensors +3 -0
- model-00016-of-00097.safetensors +3 -0
- model-00019-of-00097.safetensors +3 -0
- model-00021-of-00097.safetensors +3 -0
- model-00022-of-00097.safetensors +3 -0
- model-00024-of-00097.safetensors +3 -0
- model-00027-of-00097.safetensors +3 -0
- model-00035-of-00097.safetensors +3 -0
- model-00036-of-00097.safetensors +3 -0
- model-00051-of-00097.safetensors +3 -0
- model-00052-of-00097.safetensors +3 -0
- model-00054-of-00097.safetensors +3 -0
- model-00057-of-00097.safetensors +3 -0
- model-00058-of-00097.safetensors +3 -0
- model-00060-of-00097.safetensors +3 -0
- model-00063-of-00097.safetensors +3 -0
- model-00065-of-00097.safetensors +3 -0
- model-00069-of-00097.safetensors +3 -0
- model-00074-of-00097.safetensors +3 -0
- model-00077-of-00097.safetensors +3 -0
- model-00082-of-00097.safetensors +3 -0
- model-00084-of-00097.safetensors +3 -0
- model-00090-of-00097.safetensors +3 -0
- model-00093-of-00097.safetensors +3 -0
- model-00095-of-00097.safetensors +3 -0
- model-00096-of-00097.safetensors +3 -0
- model.safetensors.index.json +0 -0
- processor_config.json +46 -0
- special_tokens_map.json +22 -0
- tiktoken/tokenizer.model +3 -0
- 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 |
+
}
|