Upload folder using huggingface_hub

.gitattributes

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

README.md

@@ -2,4 +2,184 @@
 license: apache-2.0
 base_model:
 - chromadb/context-1
----
+tags:
+- mxfp4
+- quantized
+- gpt-oss
+- moe
+- agentic
+- search
+---
+
+# Chroma Context-1 (MXFP4)
+
+MXFP4-quantized version of [chromadb/context-1](https://huggingface.co/chromadb/context-1),
+a 20B parameter agentic search model fine-tuned from
+[openai/gpt-oss-20b](https://huggingface.co/openai/gpt-oss-20b).
+
+This checkpoint reduces the model size from **39 GB** (BF16) to **~14 GB** (MXFP4)
+with minimal quality degradation, enabling inference on a single GPU with more
+headroom for KV cache.
+
+## Model Details
+
+- **Base model:** chromadb/context-1 (BF16)
+- **Architecture:** GptOssForCausalLM (Mixture of Experts, 24 layers, 32 experts, top-4 routing)
+- **Parameters:** 20B total
+- **Quantization:** MXFP4 (E2M1 weights + E8M0 scales, group size 32)
+- **Quantized layers:** MoE expert weights (`gate_up_proj`, `down_proj`)
+- **Non-quantized layers:** attention, router, embeddings, LM head (remain in BF16)
+- **File size:** ~14 GB (model.safetensors)
+
+## Quantization Details
+
+MXFP4 (Microscaling FP4) is a 4-bit floating-point format using the E2M1
+representation (2-bit exponent, 1-bit mantissa) with shared E8M0 per-group
+scaling factors. Each group of 32 weights shares a single 8-bit scale,
+and each individual weight is stored as a 4-bit FP4 code. Two FP4 codes
+are packed into one uint8 byte (low nibble = even index, high nibble = odd index).
+
+This matches the quantization format used by the official
+[openai/gpt-oss-20b](https://huggingface.co/openai/gpt-oss-20b) MXFP4 checkpoint
+and is natively supported by vLLM's Marlin MXFP4 kernels.
+
+### What is quantized
+
+| Component | Format | Notes |
+|-----------|--------|-------|
+| `mlp.experts.gate_up_proj` | MXFP4 | Stored as `_blocks` (U8) + `_scales` (U8) |
+| `mlp.experts.down_proj` | MXFP4 | Stored as `_blocks` (U8) + `_scales` (U8) |
+| `self_attn.*` | BF16 | Kept at full precision |
+| `mlp.router.*` | BF16 | Kept at full precision |
+| `embed_tokens`, `lm_head` | BF16 | Kept at full precision |
+
+### Tensor layout
+
+Expert weights are transposed from the BF16 checkpoint layout
+`[num_experts, in_features, out_features]` to the vLLM-expected
+`[num_experts, out_features, in_features]` before quantization.
+This is required because vLLM's MXFP4 loader performs a direct copy
+without transposition (unlike the BF16 loader which transposes on load).
+
+## Usage with vLLM
+
+Requires vLLM >= 0.18.0.
+
+```bash
+vllm serve evilfreelancer/context-1-mxfp4 \
+    --served-model-name openai/gpt-oss-20b \
+    --trust-remote-code \
+    --dtype auto \
+    --gpu-memory-utilization 0.9 \
+    --max-model-len 64000 \
+    --max-num-batched-tokens 64000 \
+    --kv-cache-dtype fp8 \
+    --enable-auto-tool-choice \
+    --tool-call-parser openai
+```
+
+Docker Compose example:
+
+```yaml
+services:
+  gpt-oss-20b:
+    image: vllm/vllm-openai:v0.18.0
+    restart: always
+    entrypoint: vllm
+    command: >
+      serve evilfreelancer/context-1-mxfp4
+      --served-model-name openai/gpt-oss-20b
+      --trust-remote-code
+      --dtype auto
+      --gpu-memory-utilization 0.9
+      --max-model-len 64000
+      --max-num-batched-tokens 64000
+      --kv-cache-dtype fp8
+      --enable-auto-tool-choice
+      --tool-call-parser openai
+    ports:
+      - 8081:8000
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              device_ids: ["0"]
+              capabilities: [gpu]
+```
+
+### Note on `generation_config.json`
+
+The `eos_token_id` includes token 200012 (`<|call|>`) in addition to the
+standard `<|return|>` (200002). This is required for tool calling to work
+correctly - without it, the model does not stop generation after emitting
+a tool call, and the Harmony parser fails.
+
+## Conversion Script
+
+The included `convert_mxfp4.py` converts the original BF16
+[chromadb/context-1](https://huggingface.co/chromadb/context-1) weights
+to MXFP4 format. It requires only numpy.
+
+```bash
+pip install numpy
+python convert_mxfp4.py
+```
+
+The script expects the BF16 model in a sibling `context-1/` directory
+and writes the quantized model to `context-1-mxfp4/`.
+
+### What the script does
+
+1. Reads each tensor from the BF16 `model.safetensors`
+2. For MoE expert weights (`gate_up_proj` and `down_proj`):
+   - Transposes from `[E, in, out]` to `[E, out, in]`
+   - Computes per-group E8M0 scales (group size 32)
+   - Quantizes to E2M1 FP4 codes using nearest rounding
+   - Packs two FP4 codes per byte
+   - Saves as `*_blocks` (packed weights) and `*_scales` (shared exponents)
+3. Copies all other tensors (attention, router, embeddings) as-is in BF16
+4. Copies tokenizer and chat template files
+5. Writes `config.json` with `quantization_config` for vLLM
+
+## Key Capabilities
+
+(Inherited from chromadb/context-1)
+
+- **Query decomposition:** Breaks complex multi-constraint questions into
+  targeted subqueries.
+- **Parallel tool calling:** Averages 2.56 tool calls per turn, reducing
+  total turns and end-to-end latency.
+- **Self-editing context:** Selectively prunes irrelevant documents
+  mid-search to sustain retrieval quality over long horizons within a
+  bounded context window (0.94 prune accuracy).
+- **Cross-domain generalization:** Trained on web, legal, and finance
+  tasks; generalizes to held-out domains and public benchmarks
+  (BrowseComp-Plus, SealQA, FRAMES, HLE).
+
+## Important: Agent Harness Required
+
+Context-1 is trained to operate within a specific agent harness that
+manages tool execution, token budgets, context pruning, and deduplication.
+**The harness is not yet public.** Running the model without it will not
+reproduce the results reported in the technical report.
+
+See the [technical report](https://trychroma.com/research/context-1) for
+details on the harness design.
+
+## Citation
+
+```bibtex
+@techreport{bashir2026context1,
+  title = {Chroma Context-1: Training a Self-Editing Search Agent},
+  author = {Bashir, Hammad and Hong, Kelly and Jiang, Patrick and Shi, Zhiyi},
+  year = {2026},
+  month = {March},
+  institution = {Chroma},
+  url = {https://trychroma.com/research/context-1},
+}
+```
+
+## License
+
+Apache 2.0

chat_template.jinja

@@ -0,0 +1,315 @@
+{# Copyright 2025-present Unsloth. Apache 2.0 License. Unsloth chat template fixes. Edited from ggml-org & OpenAI #}
+{#-
+  In addition to the normal inputs of `messages` and `tools`, this template also accepts the
+  following kwargs:
+  - "builtin_tools": A list, can contain "browser" and/or "python".
+  - "model_identity": A string that optionally describes the model identity.
+  - "reasoning_effort": A string that describes the reasoning effort, defaults to "medium".
+ #}
+
+{#- Tool Definition Rendering ============================================== #}
+{%- macro render_typescript_type(param_spec, required_params, is_nullable=false) -%}
+    {%- if param_spec.type == "array" -%}
+        {%- if param_spec['items'] -%}
+            {%- if param_spec['items']['type'] == "string" -%}
+                {{- "string[]" }}
+            {%- elif param_spec['items']['type'] == "number" -%}
+                {{- "number[]" }}
+            {%- elif param_spec['items']['type'] == "integer" -%}
+                {{- "number[]" }}
+            {%- elif param_spec['items']['type'] == "boolean" -%}
+                {{- "boolean[]" }}
+            {%- else -%}
+                {%- set inner_type = render_typescript_type(param_spec['items'], required_params) -%}
+                {%- if inner_type == "object | object" or inner_type|length > 50 -%}
+                    {{- "any[]" }}
+                {%- else -%}
+                    {{- inner_type + "[]" }}
+                {%- endif -%}
+            {%- endif -%}
+            {%- if param_spec.nullable -%}
+                {{- " | null" }}
+            {%- endif -%}
+        {%- else -%}
+            {{- "any[]" }}
+            {%- if param_spec.nullable -%}
+                {{- " | null" }}
+            {%- endif -%}
+        {%- endif -%}
+    {%- elif param_spec.type is defined and param_spec.type is iterable and param_spec.type is not string and param_spec.type is not mapping and param_spec.type[0] is defined -%}
+        {#- Handle array of types like ["object", "object"] from Union[dict, list] #}
+        {%- if param_spec.type | length > 1 -%}
+            {{- param_spec.type | join(" | ") }}
+        {%- else -%}
+            {{- param_spec.type[0] }}
+        {%- endif -%}
+    {%- elif param_spec.oneOf -%}
+        {#- Handle oneOf schemas - check for complex unions and fallback to any #}
+        {%- set has_object_variants = false -%}
+        {%- for variant in param_spec.oneOf -%}
+            {%- if variant.type == "object" -%}
+                {%- set has_object_variants = true -%}
+            {%- endif -%}
+        {%- endfor -%}
+        {%- if has_object_variants and param_spec.oneOf|length > 1 -%}
+            {{- "any" }}
+        {%- else -%}
+            {%- for variant in param_spec.oneOf -%}
+                {{- render_typescript_type(variant, required_params) -}}
+                {%- if variant.description %}
+                    {{- "// " + variant.description }}
+                {%- endif -%}
+                {%- if variant.default is defined %}
+                    {{ "// default: " + variant.default|tojson }}
+                {%- endif -%}
+                {%- if not loop.last %}
+                    {{- " | " }}
+                {% endif -%}
+            {%- endfor -%}
+        {%- endif -%}
+    {%- elif param_spec.type == "string" -%}
+        {%- if param_spec.enum -%}
+            {{- '"' + param_spec.enum|join('" | "') + '"' -}}
+        {%- else -%}
+            {{- "string" }}
+            {%- if param_spec.nullable %}
+                {{- " | null" }}
+            {%- endif -%}
+        {%- endif -%}
+    {%- elif param_spec.type == "number" -%}
+        {{- "number" }}
+    {%- elif param_spec.type == "integer" -%}
+        {{- "number" }}
+    {%- elif param_spec.type == "boolean" -%}
+        {{- "boolean" }}
+
+    {%- elif param_spec.type == "object" -%}
+        {%- if param_spec.properties -%}
+            {{- "{\n" }}
+            {%- for prop_name, prop_spec in param_spec.properties.items() -%}
+                {{- prop_name -}}
+                {%- if prop_name not in (param_spec.required or []) -%}
+                    {{- "?" }}
+                {%- endif -%}
+                {{- ": " }}
+                {{ render_typescript_type(prop_spec, param_spec.required or []) }}
+                {%- if not loop.last -%}
+                    {{-", " }}
+                {%- endif -%}
+            {%- endfor -%}
+            {{- "}" }}
+        {%- else -%}
+            {{- "object" }}
+        {%- endif -%}
+    {%- else -%}
+        {{- "any" }}
+    {%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_tool_namespace(namespace_name, tools) -%}
+    {{- "## " + namespace_name + "\n\n" }}
+    {{- "namespace " + namespace_name + " {\n\n" }}
+    {%- for tool in tools %}
+        {%- set tool = tool.function %}
+        {{- "// " + tool.description + "\n" }}
+        {{- "type "+ tool.name + " = " }}
+        {%- if tool.parameters and tool.parameters.properties -%}
+            {{- "(_: " }}
+            {{- "{\n" }}
+            {%- for param_name, param_spec in tool.parameters.properties.items() %}
+                {{- "// " + param_spec.description + "\n" }}
+                {{- param_name }}
+                {%- if param_name not in (tool.parameters.required or []) -%}
+                    {{- "?" }}
+                {%- endif -%}
+                {{- ": " }}
+                {{- render_typescript_type(param_spec, tool.parameters.required or []) }}
+                {%- if param_spec.default is defined -%}
+                    {%- if param_spec.enum %}
+                        {{- ", // default: " + param_spec.default }}
+                    {%- elif param_spec.oneOf %}
+                        {{- "// default: " + param_spec.default }}
+                    {%- else %}
+                        {{- ", // default: " + param_spec.default|tojson }}
+                    {%- endif -%}
+                {%- endif -%}
+                {%- if not loop.last %}
+                    {{- ",\n" }}
+                {%- else %}
+                    {{- "\n" }}
+                {%- endif -%}
+            {%- endfor %}
+            {{- "}) => any;\n\n" }}
+        {%- else -%}
+            {{- "() => any;\n\n" }}
+        {%- endif -%}
+    {%- endfor %}
+    {{- "} // namespace " + namespace_name }}
+{%- endmacro -%}
+
+{%- macro render_builtin_tools(browser_tool, python_tool) -%}
+    {%- if browser_tool %}
+        {{- "## browser\n\n" }}
+        {{- "// Tool for browsing.\n" }}
+        {{- "// The `cursor` appears in brackets before each browsing display: `[{cursor}]`.\n" }}
+        {{- "// Cite information from the tool using the following format:\n" }}
+        {{- "// `【{cursor}†L{line_start}(-L{line_end})?】`, for example: `【6†L9-L11】` or `【8†L3】`.\n" }}
+        {{- "// Do not quote more than 10 words directly from the tool output.\n" }}
+        {{- "// sources=web (default: web)\n" }}
+        {{- "namespace browser {\n\n" }}
+        {{- "// Searches for information related to `query` and displays `topn` results.\n" }}
+        {{- "type search = (_: {\n" }}
+        {{- "query: string,\n" }}
+        {{- "topn?: number, // default: 10\n" }}
+        {{- "source?: string,\n" }}
+        {{- "}) => any;\n\n" }}
+        {{- "// Opens the link `id` from the page indicated by `cursor` starting at line number `loc`, showing `num_lines` lines.\n" }}
+        {{- "// Valid link ids are displayed with the formatting: `【{id}†.*】`.\n" }}
+        {{- "// If `cursor` is not provided, the most recent page is implied.\n" }}
+        {{- "// If `id` is a string, it is treated as a fully qualified URL associated with `source`.\n" }}
+        {{- "// If `loc` is not provided, the viewport will be positioned at the beginning of the document or centered on the most relevant passage, if available.\n" }}
+        {{- "// Use this function without `id` to scroll to a new location of an opened page.\n" }}
+        {{- "type open = (_: {\n" }}
+        {{- "id?: number | string, // default: -1\n" }}
+        {{- "cursor?: number, // default: -1\n" }}
+        {{- "loc?: number, // default: -1\n" }}
+        {{- "num_lines?: number, // default: -1\n" }}
+        {{- "view_source?: boolean, // default: false\n" }}
+        {{- "source?: string,\n" }}
+        {{- "}) => any;\n\n" }}
+        {{- "// Finds exact matches of `pattern` in the current page, or the page given by `cursor`.\n" }}
+        {{- "type find = (_: {\n" }}
+        {{- "pattern: string,\n" }}
+        {{- "cursor?: number, // default: -1\n" }}
+        {{- "}) => any;\n\n" }}
+        {{- "} // namespace browser\n\n" }}
+    {%- endif -%}
+
+    {%- if python_tool %}
+        {{- "## python\n\n" }}
+        {{- "Use this tool to execute Python code in your chain of thought. The code will not be shown to the user. This tool should be used for internal reasoning, but not for code that is intended to be visible to the user (e.g. when creating plots, tables, or files).\n\n" }}
+        {{- "When you send a message containing Python code to python, it will be executed in a stateful Jupyter notebook environment. python will respond with the output of the execution or time out after 120.0 seconds. The drive at '/mnt/data' can be used to save and persist user files. Internet access for this session is UNKNOWN. Depends on the cluster.\n\n" }}
+    {%- endif -%}
+{%- endmacro -%}
+
+{#- System Message Construction ============================================ #}
+{%- macro build_system_message() -%}
+    {%- if model_identity is not defined %}
+        {{- "You are ChatGPT, a large language model trained by OpenAI.\n" -}}
+    {%- else %}
+        {{- model_identity }}
+    {%- endif %}
+    {{- "Knowledge cutoff: 2024-06\n" }}
+    {{- "Current date: " + strftime_now("%Y-%m-%d") + "\n\n" }}
+    {%- if reasoning_effort is not defined %}
+        {%- set reasoning_effort = "medium" %}
+    {%- endif %}
+    {{- "Reasoning: " + reasoning_effort + "\n\n" }}
+    {%- if builtin_tools is defined %}
+        {{- "# Tools\n\n" }}
+        {%- set available_builtin_tools = namespace(browser=false, python=false) %}
+        {%- for tool in builtin_tools %}
+            {%- if tool == "browser" %}
+                {%- set available_builtin_tools.browser = true %}
+            {%- elif tool == "python" %}
+                {%- set available_builtin_tools.python = true %}
+            {%- endif %}
+        {%- endfor %}
+        {{- render_builtin_tools(available_builtin_tools.browser, available_builtin_tools.python) }}
+    {%- endif -%}
+    {{- "# Valid channels: analysis, commentary, final. Channel must be included for every message." }}
+    {%- if tools is defined -%}
+        {{- "\nCalls to these tools must go to the commentary channel: 'functions'." }}
+    {%- endif -%}
+{%- endmacro -%}
+
+{#- Main Template Logic ================================================= #}
+{#- Set defaults #}
+
+{#- Render system message #}
+{{- "<|start|>system<|message|>" }}
+{{- build_system_message() }}
+{{- "<|end|>" }}
+
+{#- Extract developer message #}
+{%- if messages[0].role == "developer" or messages[0].role == "system" %}
+    {%- set developer_message = messages[0].content %}
+    {%- set loop_messages = messages[1:] %}
+{%- else %}
+    {%- set developer_message = "" %}
+    {%- set loop_messages = messages %}
+{%- endif %}
+
+{#- Render developer message #}
+{%- if developer_message or tools %}
+    {{- "<|start|>developer<|message|>" }}
+    {%- if developer_message %}
+        {{- "# Instructions\n\n" }}
+        {{- developer_message }}
+    {%- endif %}
+    {%- if tools -%}
+        {{- "\n\n" }}
+        {{- "# Tools\n\n" }}
+        {{- render_tool_namespace("functions", tools) }}
+    {%- endif -%}
+    {{- "<|end|>" }}
+{%- endif %}
+
+{#- Render messages #}
+{%- set last_tool_call = namespace(name=none) %}
+{%- for message in loop_messages -%}
+    {#- At this point only assistant/user/tool messages should remain #}
+    {%- if message.role == 'assistant' -%}
+        {%- if "tool_calls" in message %}
+            {#- We assume max 1 tool call per message, and so we infer the tool call name #}
+            {#- in "tool" messages from the most recent assistant tool call name #}
+            {%- set tool_call = message.tool_calls[0] %}
+            {%- if tool_call.function %}
+                {%- set tool_call = tool_call.function %}
+            {%- endif %}
+            {%- if message.content %}
+                {{- "<|start|>assistant<|channel|>analysis<|message|>" + message.content + "<|end|>" }}
+            {%- endif %}
+            {{- "<|start|>assistant to=" }}
+            {{- "functions." + tool_call.name + "<|channel|>commentary json<|message|>" }}
+            {{- tool_call.arguments|tojson }}
+            {{- "<|call|>" }}
+            {%- set last_tool_call.name = tool_call.name %}
+        {%- elif "thinking" in message and loop.last and not add_generation_prompt %}
+            {#- Only render the CoT if the final turn is an assistant turn and add_generation_prompt is false #}
+            {#- This is a situation that should only occur in training, never in inference. #}
+            {{- "<|start|>assistant<|channel|>analysis<|message|>" + message.thinking + "<|end|>" }}
+            {#- <|return|> indicates the end of generation, but <|end|> does not #}
+            {#- <|return|> should never be an input to the model, but we include it as the final token #}
+            {#- when training, so the model learns to emit it. #}
+            {{- "<|start|>assistant<|channel|>final<|message|>" + message.content + "<|return|>" }}
+            {%- set last_tool_call.name = none %}
+        {%- elif "thinking" in message %}
+            {#- CoT is dropped during all previous turns, so we never render it for inference #}
+            {{- "<|start|>assistant<|channel|>final<|message|>" + message.content + "<|end|>" }}
+            {%- set last_tool_call.name = none %}
+        {%- elif loop.last and not add_generation_prompt %}
+            {#- <|return|> indicates the end of generation, but <|end|> does not #}
+            {#- <|return|> should never be an input to the model, but we include it as the final token #}
+            {#- when training, so the model learns to emit it. #}
+            {{- "<|start|>assistant<|message|>" + message.content + "<|return|>" }}
+        {%- else %}
+            {{- "<|start|>assistant<|message|>" + message.content + "<|end|>" }}
+            {%- set last_tool_call.name = none %}
+        {%- endif %}
+    {%- elif message.role == 'tool' -%}
+        {%- if last_tool_call.name is none %}
+            {{- raise_exception("Message has tool role, but there was no previous assistant message with a tool call!") }}
+        {%- endif %}
+        {{- "<|start|>functions." + last_tool_call.name }}
+        {{- " to=assistant<|channel|>commentary<|message|>" + message.content|tojson + "<|end|>" }}
+    {%- else -%}
+        {{- "<|start|>user<|message|>" + message.content + "<|end|>" }}
+    {%- endif -%}
+{%- endfor -%}
+
+{#- Generation prompt #}
+{%- if add_generation_prompt -%}
+<|start|>assistant
+{%- endif -%}
+{# Copyright 2025-present Unsloth. Apache 2.0 License. Unsloth chat template fixes. Edited from ggml-org & OpenAI #}

config.json

@@ -0,0 +1,78 @@
+{
+  "architectures": [
+    "GptOssForCausalLM"
+  ],
+  "attention_bias": true,
+  "attention_dropout": 0.0,
+  "bos_token_id": null,
+  "dtype": "bfloat16",
+  "eos_token_id": 200002,
+  "experts_per_token": 4,
+  "head_dim": 64,
+  "hidden_act": "silu",
+  "hidden_size": 2880,
+  "initial_context_length": 4096,
+  "initializer_range": 0.02,
+  "intermediate_size": 2880,
+  "layer_types": [
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 131072,
+  "model_type": "gpt_oss",
+  "num_attention_heads": 64,
+  "num_experts_per_tok": 4,
+  "num_hidden_layers": 24,
+  "num_key_value_heads": 8,
+  "num_local_experts": 32,
+  "output_router_logits": false,
+  "pad_token_id": 199999,
+  "rms_norm_eps": 1e-05,
+  "rope_parameters": {
+    "beta_fast": 32.0,
+    "beta_slow": 1.0,
+    "factor": 32.0,
+    "original_max_position_embeddings": 4096,
+    "rope_theta": 150000,
+    "rope_type": "yarn",
+    "truncate": false
+  },
+  "router_aux_loss_coef": 0.9,
+  "sliding_window": 128,
+  "swiglu_limit": 7.0,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.3.0",
+  "use_cache": true,
+  "vocab_size": 201088,
+  "quantization_config": {
+    "modules_to_not_convert": [
+      "model.layers.*.self_attn",
+      "model.layers.*.mlp.router",
+      "model.embed_tokens",
+      "lm_head"
+    ],
+    "quant_method": "mxfp4"
+  }
+}

convert_mxfp4.py

@@ -0,0 +1,282 @@
+#!/usr/bin/env python3
+"""Convert chromadb/context-1 BF16 weights to MXFP4 format for vLLM."""
+
+import json
+import math
+import os
+import shutil
+import struct
+import sys
+import time
+
+import numpy as np
+
+MODEL_DIR = os.path.join(os.path.dirname(__file__), "context-1")
+OUTPUT_DIR = os.path.join(os.path.dirname(__file__), "context-1-mxfp4")
+
+GROUP_SIZE = 32
+
+# E2M1 FP4 positive lookup: index -> value
+FP4_VALUES = np.array(
+    [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=np.float32
+)
+FP4_MAX = 6.0
+
+# Midpoints between consecutive FP4 values for nearest rounding
+FP4_BOUNDARIES = np.array(
+    [0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0, np.inf], dtype=np.float32
+)
+
+
+def float_to_e2m1(values: np.ndarray) -> np.ndarray:
+    """Quantize float values to 4-bit E2M1 codes (0..15)."""
+    sign = (values < 0).astype(np.uint8)
+    abs_val = np.abs(values)
+    codes = np.digitize(abs_val, FP4_BOUNDARIES).astype(np.uint8)
+    codes = np.clip(codes, 0, 7)
+    return (sign << 3) | codes
+
+
+def compute_e8m0_scale(group: np.ndarray) -> tuple[np.uint8, float]:
+    """Compute E8M0 shared exponent for a group. Returns (e8m0_byte, scale_float)."""
+    amax = np.max(np.abs(group))
+    if amax == 0:
+        return np.uint8(0), 1.0
+
+    # scale = 2^ceil(log2(amax / FP4_MAX))
+    # ensures amax / scale <= FP4_MAX
+    log2_scale = math.ceil(math.log2(max(amax / FP4_MAX, 2**-127)))
+    log2_scale = max(log2_scale, -127)
+    log2_scale = min(log2_scale, 127)
+    e8m0 = np.uint8(log2_scale + 127)
+    scale = 2.0 ** log2_scale
+    return e8m0, scale
+
+
+def quantize_mxfp4(weight: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+    """
+    Quantize a 2D weight [rows, cols] from float to MXFP4 (vectorized).
+    cols must be divisible by GROUP_SIZE.
+    Returns (packed_uint8 [rows, cols//2], scales_uint8 [rows, cols//GROUP_SIZE]).
+    """
+    rows, cols = weight.shape
+    assert cols % GROUP_SIZE == 0, f"cols={cols} not divisible by {GROUP_SIZE}"
+
+    n_groups = cols // GROUP_SIZE
+    grouped = weight.reshape(rows, n_groups, GROUP_SIZE).astype(np.float32)
+
+    # Vectorized E8M0 scale computation per group
+    amax = np.max(np.abs(grouped), axis=-1)  # [rows, n_groups]
+    amax = np.maximum(amax, 2**-127)
+    log2_scale = np.ceil(np.log2(amax / FP4_MAX)).astype(np.int32)
+    log2_scale = np.clip(log2_scale, -127, 127)
+    scales = (log2_scale + 127).astype(np.uint8)  # [rows, n_groups]
+    scale_float = np.power(2.0, log2_scale.astype(np.float64)).astype(np.float32)
+
+    # Scale each group
+    scaled = grouped / scale_float[:, :, np.newaxis]  # [rows, n_groups, 32]
+
+    # Vectorized E2M1 quantization
+    flat_scaled = scaled.reshape(rows, cols)
+    fp4_codes = float_to_e2m1(flat_scaled)
+
+    # Pack 2 FP4 codes per byte: low nibble = even index, high nibble = odd index
+    even = fp4_codes[:, 0::2]
+    odd = fp4_codes[:, 1::2]
+    packed = ((odd << 4) | even).astype(np.uint8)
+
+    return packed, scales
+
+
+def read_safetensors_header(path: str) -> dict:
+    with open(path, "rb") as f:
+        header_size = struct.unpack("<Q", f.read(8))[0]
+        header_json = f.read(header_size).decode("utf-8")
+    return json.loads(header_json), header_size
+
+
+def read_tensor(path: str, info: dict) -> np.ndarray:
+    dtype_map = {"BF16": (np.uint16, 2), "F32": (np.float32, 4), "U8": (np.uint8, 1)}
+    dtype_str = info["dtype"]
+    np_dtype, elem_size = dtype_map[dtype_str]
+    shape = info["shape"]
+    offsets = info["data_offsets"]
+    start, end = offsets
+
+    with open(path, "rb") as f:
+        header_size_bytes = struct.unpack("<Q", f.read(8))[0]
+        f.seek(8 + header_size_bytes + start)
+        data = f.read(end - start)
+
+    arr = np.frombuffer(data, dtype=np_dtype).reshape(shape)
+    return arr
+
+
+def bf16_to_f32(arr: np.ndarray) -> np.ndarray:
+    """Convert BF16 (stored as uint16) to float32."""
+    f32_bytes = np.zeros(arr.shape, dtype=np.uint32)
+    f32_bytes[:] = arr.astype(np.uint32) << 16
+    return f32_bytes.view(np.float32)
+
+
+def build_safetensors(tensors: dict, metadata: dict | None = None) -> bytes:
+    """Build safetensors binary from dict of {name: (np_array, dtype_str)}."""
+    header = {}
+    if metadata:
+        header["__metadata__"] = metadata
+
+    data_parts = []
+    offset = 0
+    for name, (arr, dtype_str) in sorted(tensors.items()):
+        raw = arr.tobytes()
+        data_parts.append(raw)
+        header[name] = {
+            "dtype": dtype_str,
+            "shape": list(arr.shape),
+            "data_offsets": [offset, offset + len(raw)],
+        }
+        offset += len(raw)
+
+    header_json = json.dumps(header, separators=(",", ":")).encode("utf-8")
+    # Pad header to 8-byte alignment
+    padding = (8 - len(header_json) % 8) % 8
+    header_json += b" " * padding
+
+    result = struct.pack("<Q", len(header_json)) + header_json
+    for part in data_parts:
+        result += part
+    return result
+
+
+def main():
+    st_path = os.path.join(MODEL_DIR, "model.safetensors")
+
+    print("Reading safetensors header...")
+    header, header_size = read_safetensors_header(st_path)
+
+    metadata = header.pop("__metadata__", {"format": "pt"})
+
+    expert_weight_suffixes = (".mlp.experts.gate_up_proj", ".mlp.experts.down_proj")
+
+    output_tensors = {}
+    total = len([k for k in header if k != "__metadata__"])
+    done = 0
+
+    for name, info in sorted(header.items()):
+        if name == "__metadata__":
+            continue
+        done += 1
+        is_expert_weight = any(name.endswith(s) for s in expert_weight_suffixes)
+
+        if is_expert_weight:
+            print(f"[{done}/{total}] Quantizing {name} {info['shape']}...")
+            t0 = time.time()
+            raw = read_tensor(st_path, info)
+            weight_f32 = bf16_to_f32(raw)
+
+            is_gate_up = name.endswith(".gate_up_proj")
+            num_experts = weight_f32.shape[0]
+
+            # BF16 checkpoint stores weights as [E, in_features, out_features].
+            # vLLM MXFP4 expects [E, out_features, in_features // 2] (packed).
+            # Both gate_up_proj and down_proj need transposing to [E, out, in].
+            weight_f32 = np.ascontiguousarray(
+                np.transpose(weight_f32, (0, 2, 1))
+            )
+
+            blocks_list = []
+            scales_list = []
+            for e in range(num_experts):
+                packed, scales = quantize_mxfp4(weight_f32[e])
+                blocks_list.append(packed)
+                scales_list.append(scales)
+
+            blocks = np.stack(blocks_list, axis=0)
+            scales = np.stack(scales_list, axis=0)
+
+            blocks_name = name.replace(".gate_up_proj", ".gate_up_proj_blocks").replace(
+                ".down_proj", ".down_proj_blocks"
+            )
+            scales_name = name.replace(".gate_up_proj", ".gate_up_proj_scales").replace(
+                ".down_proj", ".down_proj_scales"
+            )
+
+            output_tensors[blocks_name] = (blocks, "U8")
+            output_tensors[scales_name] = (scales, "U8")
+
+            dt = time.time() - t0
+            print(f"    -> {blocks_name} {list(blocks.shape)}, "
+                  f"{scales_name} {list(scales.shape)} ({dt:.1f}s)")
+        else:
+            print(f"[{done}/{total}] Copying {name} {info['shape']}...")
+            raw = read_tensor(st_path, info)
+            output_tensors[name] = (raw, info["dtype"])
+
+    os.makedirs(OUTPUT_DIR, exist_ok=True)
+
+    print("\nWriting output safetensors...")
+    out_path = os.path.join(OUTPUT_DIR, "model.safetensors")
+
+    # Write in streaming fashion to avoid huge memory spike
+    # First build header, then write data
+    header_dict = {}
+    if metadata:
+        header_dict["__metadata__"] = metadata
+
+    offset = 0
+    tensor_order = sorted(output_tensors.keys())
+    for tname in tensor_order:
+        arr, dtype_str = output_tensors[tname]
+        raw_size = arr.nbytes
+        header_dict[tname] = {
+            "dtype": dtype_str,
+            "shape": list(arr.shape),
+            "data_offsets": [offset, offset + raw_size],
+        }
+        offset += raw_size
+
+    header_json = json.dumps(header_dict, separators=(",", ":")).encode("utf-8")
+    padding = (8 - len(header_json) % 8) % 8
+    header_json += b" " * padding
+
+    with open(out_path, "wb") as f:
+        f.write(struct.pack("<Q", len(header_json)))
+        f.write(header_json)
+        for tname in tensor_order:
+            arr, _ = output_tensors[tname]
+            f.write(arr.tobytes())
+
+    print(f"Saved {out_path} ({os.path.getsize(out_path) / 1e9:.2f} GB)")
+
+    # Copy config files and add quantization_config
+    for fname in ["generation_config.json", "tokenizer_config.json",
+                  "tokenizer.json", "chat_template.jinja"]:
+        src = os.path.join(MODEL_DIR, fname)
+        if os.path.exists(src):
+            shutil.copy2(src, os.path.join(OUTPUT_DIR, fname))
+            print(f"Copied {fname}")
+
+    # Update config.json with quantization_config
+    with open(os.path.join(MODEL_DIR, "config.json")) as f:
+        config = json.load(f)
+
+    config["quantization_config"] = {
+        "modules_to_not_convert": [
+            "model.layers.*.self_attn",
+            "model.layers.*.mlp.router",
+            "model.embed_tokens",
+            "lm_head",
+        ],
+        "quant_method": "mxfp4",
+    }
+
+    with open(os.path.join(OUTPUT_DIR, "config.json"), "w") as f:
+        json.dump(config, f, indent=2)
+        f.write("\n")
+    print("Wrote config.json with quantization_config")
+
+    print("\nDone!")
+
+
+if __name__ == "__main__":
+    main()

generation_config.json

@@ -0,0 +1,11 @@
+{
+  "bos_token_id": 199998,
+  "do_sample": true,
+  "eos_token_id": [
+    200002,
+    200012,
+    199999
+  ],
+  "pad_token_id": 199999,
+  "transformers_version": "5.3.0"
+}

model.safetensors

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

tokenizer.json

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

tokenizer_config.json

@@ -0,0 +1,183 @@
+{
+  "added_tokens_decoder": {
+    "199998": {
+      "content": "<|startoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "199999": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200000": {
+      "content": "<|reserved_200000|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200001": {
+      "content": "<|reserved_200001|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200002": {
+      "content": "<|return|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200003": {
+      "content": "<|constrain|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200004": {
+      "content": "<|reserved_200004|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200005": {
+      "content": "<|channel|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200006": {
+      "content": "<|start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200007": {
+      "content": "<|end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200008": {
+      "content": "<|message|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200009": {
+      "content": "<|reserved_200009|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200010": {
+      "content": "<|reserved_200010|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200011": {
+      "content": "<|reserved_200011|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200012": {
+      "content": "<|call|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200013": {
+      "content": "<|reserved_200013|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200014": {
+      "content": "<|reserved_200014|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200015": {
+      "content": "<|reserved_200015|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200016": {
+      "content": "<|reserved_200016|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200017": {
+      "content": "<|reserved_200017|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200018": {
+      "content": "<|endofprompt|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|startoftext|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|return|>",
+  "extra_special_tokens": {},
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<|endoftext|>",
+  "tokenizer_class": "PreTrainedTokenizerFast"
+}