dataopsnick/Qwen3-4B-Instruct-2507-zip-rc
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

初始化项目，由ModelHub XC社区提供模型

Browse Source 
Model: dataopsnick/Qwen3-4B-Instruct-2507-zip-rc
Source: Original Platform
This commit is contained in:
ModelHub XC
2026-06-07 09:56:25 +08:00
commit 6b9f76449b
16 changed files with 153355 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

36
.gitattributes vendored Normal file
View File

@@ -0,0 +1,36 @@
*.7z filter=lfs diff=lfs merge=lfs -text
*.arrow filter=lfs diff=lfs merge=lfs -text
*.bin filter=lfs diff=lfs merge=lfs -text
*.bz2 filter=lfs diff=lfs merge=lfs -text
*.ckpt filter=lfs diff=lfs merge=lfs -text
*.ftz filter=lfs diff=lfs merge=lfs -text
*.gz filter=lfs diff=lfs merge=lfs -text
*.h5 filter=lfs diff=lfs merge=lfs -text
*.joblib filter=lfs diff=lfs merge=lfs -text
*.lfs.* filter=lfs diff=lfs merge=lfs -text
*.mlmodel filter=lfs diff=lfs merge=lfs -text
*.model filter=lfs diff=lfs merge=lfs -text
*.msgpack filter=lfs diff=lfs merge=lfs -text
*.npy filter=lfs diff=lfs merge=lfs -text
*.npz filter=lfs diff=lfs merge=lfs -text
*.onnx filter=lfs diff=lfs merge=lfs -text
*.ot filter=lfs diff=lfs merge=lfs -text
*.parquet filter=lfs diff=lfs merge=lfs -text
*.pb filter=lfs diff=lfs merge=lfs -text
*.pickle filter=lfs diff=lfs merge=lfs -text
*.pkl filter=lfs diff=lfs merge=lfs -text
*.pt filter=lfs diff=lfs merge=lfs -text
*.pth filter=lfs diff=lfs merge=lfs -text
*.rar filter=lfs diff=lfs merge=lfs -text
*.safetensors filter=lfs diff=lfs merge=lfs -text
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
*.tar.* filter=lfs diff=lfs merge=lfs -text
*.tar filter=lfs diff=lfs merge=lfs -text
*.tflite filter=lfs diff=lfs merge=lfs -text
*.tgz filter=lfs diff=lfs merge=lfs -text
*.wasm filter=lfs diff=lfs merge=lfs -text
*.xz 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

275
README.md Normal file
View File

@@ -0,0 +1,275 @@
---
library_name: transformers
license: apache-2.0
base_model: Qwen/Qwen3-4B-Instruct-2507
pipeline_tag: text-generation
tags:
- zip-rc
- adaptive-compute
- introspection
- reasoning
---
# Qwen3-4B-Instruct-2507-ZIP-RC
This model is a modified version of **[Qwen/Qwen3-4B-Instruct-2507](https://huggingface.co/Qwen/Qwen3-4B-Instruct-2507)**, trained to support **Zero-Overhead Introspection (ZIP-RC)** for adaptive test-time compute.
It was created as part of a **Paper Replication** experiment for:
**"Zero-Overhead Introspection for Adaptive Test-Time Compute"** (Manvi et al., 2025).
| **Links** | **Description** |
| :--- | :--- |
| [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/dataopsnick/paper-replication/blob/main/zip-rc_replication/dataopsnick_Qwen3_4B_Instruct_2507_zip_rc_QUICKSTART.ipynb) | **Quickstart Notebook:** Run adaptive inference immediately. |
| [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/dataopsnick/paper-replication/blob/main/zip-rc_replication/paper_replication_arxiv_org_abs_2512_01457.ipynb) | **Replication Notebook:** Full experiments and reproduction. |
## Model Description
This model retains the full reasoning capabilities of the base `Qwen3-4B-Instruct` model but features a **fine-tuned LM Head**. The head has been trained to repurpose unused logit space to predict a **joint distribution of Expected Reward (Correctness) and Remaining Generation Length** at every token step.
This allows the model to "introspect" during generation with **zero computational overhead**, enabling:
* **Adaptive Sampling:** Dynamically pruning low-quality trajectories.
* **Budget Management:** Balancing compute cost vs. accuracy.
* **Self-Correction:** Detecting when a reasoning path is failing before it finishes.
## Usage
### 1. Quick Start: Adaptive Inference
The easiest way to use the model is via the `ziprc` helper library, which handles the Meta-MDP logic (branching, pruning, and swapping).
```python
import torch
import sys
import os
from huggingface_hub import hf_hub_download
# 1. Download the helper script dynamically
script_path = hf_hub_download(repo_id="dataopsnick/Qwen3-4B-Instruct-2507-zip-rc", filename="ziprc.py")
sys.path.append(os.path.dirname(script_path))
# 2. Import the downloaded module
import ziprc
# 3. Run Inference
model = ziprc.ZIPRCModel(ziprc.ZIPRCConfig())
sampler = ziprc.ZIPRCSampler(model)
prompt = "Solve the following logic puzzle: Five adults check into a hotel with three dogs. How many shoes are they all wearing?"
trajectories = sampler.generate(prompt, initial_samples=2)
best = sampler.select_best_trajectory(trajectories)
print(f"Confidence: {best['final_score']:.2%}")
```
### 2. Advanced Usage: Streaming & Configuration
This example shows how to configure the pruning aggressiveness (`alpha`) and cost penalty (`beta`), and how to stream the result to see the introspection in action.
```python
import sys
import os
#import tqdm
from huggingface_hub import hf_hub_download
from tqdm import tqdm
# 1. Download the helper script dynamically from the repo
script_path = hf_hub_download(repo_id="dataopsnick/Qwen3-4B-Instruct-2507-zip-rc", filename="ziprc.py")
sys.path.append(os.path.dirname(script_path))
# 2. Import the module
from ziprc import ZIPRCModel, ZIPRCConfig, ZIPRCSampler
# 3. Configure and Load Model
# Note: The model weights are downloaded automatically here
cfg = ZIPRCConfig(
model_name="dataopsnick/Qwen3-4B-Instruct-2507-zip-rc",
alpha=0.1, # Threshold for pruning
beta=0.05, # Cost penalty
smoothing_window=3 # For stable predictions
)
model = ZIPRCModel(cfg)
sampler = ZIPRCSampler(model)
# 4. Generate with Introspection
prompt = "Solve the following logic puzzle: Five adults check into a hotel with three dogs. How many shoes are they all wearing?"
# generate_stream produces trajectories with introspection data
trajectories = sampler.generate_stream(prompt, initial_samples=2)
# Select the best answer based on the introspection score
best = sampler.select_best_trajectory(trajectories)
print(f"Confidence: {best['final_score']:.2%}")
print(f"Answer: {model.tokenizer.decode(best['ids'][0], skip_special_tokens=True)}")
```
### 3. Low-Level: Reading the Logits
You can manually decode the introspection signal (Reward and Cost) from the reserved tokens in the logits without using the sampler.
```python
import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer
model_id = "dataopsnick/Qwen3-4B-Instruct-2507-zip-rc"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="auto")
# Configuration used during training
reward_bins = 8
length_bins = 7
total_zip_tokens = 56
zip_start_offset = 56
# ZIP tokens are located at the very end of the vocabulary
zip_start_id = model.config.vocab_size - zip_start_offset
def get_introspection_probs(logits):
"""
Extracts the joint distribution P(Reward, Length) from the logits.
"""
# Slice the reserved ZIP logits
zip_logits = logits[:, zip_start_id : zip_start_id + total_zip_tokens]
# Softmax over the flat ZIP tokens to get valid probabilities
probs = F.softmax(zip_logits, dim=-1)
# Reshape to [Batch, Reward_Bins, Length_Bins]
return probs.view(-1, reward_bins, length_bins)
# Example Inference Step
prompt = "Solve the following logic puzzle: Five adults check into a hotel with three dogs. How many shoes are they all wearing?"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
with torch.no_grad():
outputs = model(inputs.input_ids)
next_token_logits = outputs.logits[:, -1, :]
# Get Introspection Signal (Zero Overhead)
joint_dist = get_introspection_probs(next_token_logits)
# 1. Marginalize over length to get P(Reward) distribution
p_reward = joint_dist.sum(dim=2) # Shape: [Batch, Reward_Bins]
# 2. Calculate Expected Reward (Confidence)
# The reward bins are linearly spaced [0, 1]. We use bin centers for the weighted sum.
# centers = 0.0625, 0.1875, ..., 0.9375
reward_grid = torch.linspace(0.0625, 0.9375, reward_bins).to(model.device)
# E[R] = sum(P(r) * r)
expected_reward = (p_reward * reward_grid).sum(dim=1).item()
print(f"Model Confidence: {expected_reward:.2%}")
```
### 4. OpenAI-Compatible Streaming (Async)
This method exposes the introspection data (`zip_rc` field) alongside standard text generation chunks, suitable for integration with frontends.
```python
import asyncio
import nest_asyncio
from ziprc import ZIPRCModel, ZIPRCConfig, ZIPRCSampler
# 1. Setup (Run once)
# This patch is required for running async loops in Colab/Jupyter
nest_asyncio.apply()
# Load Model
cfg = ZIPRCConfig(model_name="dataopsnick/Qwen3-4B-Instruct-2507-zip-rc")
model = ZIPRCModel(cfg)
sampler = ZIPRCSampler(model)
async def consume_inference_stream():
prompt = "Solve the following logic puzzle: Five adults check into a hotel with three dogs. How many shoes are they all wearing?"
print(f"User: {prompt}\n" + "-"*60)
print("Assistant (Streaming with Introspection):")
# 2. Get the OpenAI-compatible stream
# Returns an async generator yielding chunk objects
stream = sampler.openai(prompt, max_tokens=256)
final_clean_answer = ""
async for chunk in stream:
# --- Channel A: Standard Text (Compatible with standard UIs) ---
# Use .get() to handle the final chunk where delta is empty
# Use .get() to safely handle the final chunk where delta is empty
delta = chunk.choices[0].delta
content = delta.get("content", "")
if content:
print(content, end="", flush=True)
# --- Channel B: Zero-Overhead Introspection (The "Pareto" Gain) ---
# We access the side-channel data to see what the model is thinking
# without running separate reward model inference.
if hasattr(chunk, 'zip_rc'):
info = chunk.zip_rc
# If the model performs a meta-action (Branching/Pruning), log it
# Filter out 'finished' to avoid accessing missing utility/score fields
if info.action not in ['keep', 'finished']:
print(f"\n[⚙️ META-ACTION: {info.action} | Utility: {info.utility:.4f}] ", end="")
# Check for the Final Answer
if info.get('action') == 'finished' and 'final_text' in info:
final_clean_answer = info['final_text']
# Optional: Peek at the "Confidence" (Expected Correctness) in real-time
# if info.step % 10 == 0:
# print(f" (Conf: {info.lhs_score:.1%}) ", end="")
print("\n" + "-" * 40)
print("🏆 FINAL BEST ANSWER (Clean):")
print("-" * 40)
print(final_clean_answer)
# 3. Execution
loop = asyncio.get_event_loop()
loop.run_until_complete(consume_inference_stream())
```
### 5. Local Server Deployment
You can deploy an OpenAI-compatible API server that streams both text and introspection data.
```python
import sys
import os
import asyncio
import uvicorn
from huggingface_hub import hf_hub_download
# 1. Download server.py
script_path = hf_hub_download(repo_id="dataopsnick/Qwen3-4B-Instruct-2507-zip-rc", filename="server.py")
sys.path.append(os.path.dirname(script_path))
# 2. Import the app
# NOTE: This will load the model weights again if they aren't cached.
# If you are low on VRAM, restart your runtime before running this cell.
from server import app
# 3. Run the Server (Colab/Jupyter Compatible)
HOST = "0.0.0.0"
PORT = 8000
config = uvicorn.Config(app, host=HOST, port=PORT)
server = uvicorn.Server(config)
try:
# Check if we are in an existing loop (Colab)
loop = asyncio.get_running_loop()
print(f"🚀 Server running in background on http://{HOST}:{PORT}")
loop.create_task(server.serve())
except RuntimeError:
# Standard script execution
asyncio.run(server.serve())
```
## Citation
```bibtex
@article{manvi2025ziprc,
title={Zero-Overhead Introspection for Adaptive Test-Time Compute},
author={Manvi, Rohin and Hong, Joey and Seyde, Tim and Labonne, Maxime and Lechner, Mathias and Levine, Sergey},
journal={arXiv preprint arXiv:2512.01457},
year={2025}
}
```

28
added_tokens.json Normal file
View File

@@ -0,0 +1,28 @@
{
"</think>": 151668,
"</tool_call>": 151658,
"</tool_response>": 151666,
"<think>": 151667,
"<tool_call>": 151657,
"<tool_response>": 151665,
"<|box_end|>": 151649,
"<|box_start|>": 151648,
"<|endoftext|>": 151643,
"<|file_sep|>": 151664,
"<|fim_middle|>": 151660,
"<|fim_pad|>": 151662,
"<|fim_prefix|>": 151659,
"<|fim_suffix|>": 151661,
"<|im_end|>": 151645,
"<|im_start|>": 151644,
"<|image_pad|>": 151655,
"<|object_ref_end|>": 151647,
"<|object_ref_start|>": 151646,
"<|quad_end|>": 151651,
"<|quad_start|>": 151650,
"<|repo_name|>": 151663,
"<|video_pad|>": 151656,
"<|vision_end|>": 151653,
"<|vision_pad|>": 151654,
"<|vision_start|>": 151652
}

61
chat_template.jinja Normal file
View File

@@ -0,0 +1,61 @@
{%- if tools %}
{{- '<|im_start|>system\n' }}
{%- if messages[0].role == 'system' %}
{{- messages[0].content + '\n\n' }}
{%- endif %}
{{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
{%- for tool in tools %}
{{- "\n" }}
{{- tool | tojson }}
{%- endfor %}
{{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
{%- else %}
{%- if messages[0].role == 'system' %}
{{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
{%- endif %}
{%- endif %}
{%- for message in messages %}
{%- if message.content is string %}
{%- set content = message.content %}
{%- else %}
{%- set content = '' %}
{%- endif %}
{%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
{{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
{%- elif message.role == "assistant" %}
{{- '<|im_start|>' + message.role + '\n' + content }}
{%- if message.tool_calls %}
{%- for tool_call in message.tool_calls %}
{%- if (loop.first and content) or (not loop.first) %}
{{- '\n' }}
{%- endif %}
{%- if tool_call.function %}
{%- set tool_call = tool_call.function %}
{%- endif %}
{{- '<tool_call>\n{"name": "' }}
{{- tool_call.name }}
{{- '", "arguments": ' }}
{%- if tool_call.arguments is string %}
{{- tool_call.arguments }}
{%- else %}
{{- tool_call.arguments | tojson }}
{%- endif %}
{{- '}\n</tool_call>' }}
{%- endfor %}
{%- endif %}
{{- '<|im_end|>\n' }}
{%- elif message.role == "tool" %}
{%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
{{- '<|im_start|>user' }}
{%- endif %}
{{- '\n<tool_response>\n' }}
{{- content }}
{{- '\n</tool_response>' }}
{%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
{{- '<|im_end|>\n' }}
{%- endif %}
{%- endif %}
{%- endfor %}
{%- if add_generation_prompt %}
{{- '<|im_start|>assistant\n' }}
{%- endif %}

68
config.json Normal file
View File

@@ -0,0 +1,68 @@
{
"architectures": [
"Qwen3ForCausalLM"
],
"attention_bias": false,
"attention_dropout": 0.0,
"bos_token_id": 151643,
"dtype": "bfloat16",
"eos_token_id": 151645,
"head_dim": 128,
"hidden_act": "silu",
"hidden_size": 2560,
"initializer_range": 0.02,
"intermediate_size": 9728,
"layer_types": [
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention"
],
"max_position_embeddings": 262144,
"max_window_layers": 36,
"model_type": "qwen3",
"num_attention_heads": 32,
"num_hidden_layers": 36,
"num_key_value_heads": 8,
"rms_norm_eps": 1e-06,
"rope_scaling": null,
"rope_theta": 5000000,
"sliding_window": null,
"tie_word_embeddings": true,
"transformers_version": "4.57.3",
"use_cache": true,
"use_sliding_window": false,
"vocab_size": 151936
}

13
generation_config.json Normal file
View File

@@ -0,0 +1,13 @@
{
"bos_token_id": 151643,
"do_sample": true,
"eos_token_id": [
151645,
151643
],
"pad_token_id": 151643,
"temperature": 0.7,
"top_k": 20,
"top_p": 0.8,
"transformers_version": "4.57.3"
}

151388
merges.txt Normal file
View File

File diff suppressed because it is too large Load Diff

3
model-00001-of-00002.safetensors Normal file
View File

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

3
model-00002-of-00002.safetensors Normal file
View File

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

406
model.safetensors.index.json Normal file
View File

@@ -0,0 +1,406 @@
{
"metadata": {
"total_parameters": 4022468096,
"total_size": 8044936192
},
"weight_map": {
"model.embed_tokens.weight": "model-00001-of-00002.safetensors",
"model.layers.0.input_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.0.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.0.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.0.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.0.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.0.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.0.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.0.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.0.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.0.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.0.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.1.input_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.1.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.1.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.1.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.1.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.1.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.1.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.1.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.1.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.1.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.1.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.10.input_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.10.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.10.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.10.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.10.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.10.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.10.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.10.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.10.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.10.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.10.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.11.input_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.11.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.11.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.11.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.11.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.11.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.11.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.11.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.11.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.11.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.11.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.12.input_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.12.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.12.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.12.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.12.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.12.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.12.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.12.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.12.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.12.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.12.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.13.input_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.13.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.13.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.13.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.13.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.13.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.13.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.13.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.13.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.13.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.13.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.14.input_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.14.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.14.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.14.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.14.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.14.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.14.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.14.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.14.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.14.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.14.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.15.input_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.15.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.15.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.15.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.15.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.15.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.15.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.15.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.15.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.15.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.15.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.16.input_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.16.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.16.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.16.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.16.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.16.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.16.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.16.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.16.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.16.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.16.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.17.input_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.17.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.17.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.17.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.17.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.17.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.17.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.17.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.17.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.17.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.17.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.18.input_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.18.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.18.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.18.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.18.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.18.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.18.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.18.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.18.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.18.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.18.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.19.input_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.19.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.19.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.19.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.19.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.19.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.19.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.19.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.19.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.19.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.19.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.2.input_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.2.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.2.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.2.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.2.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.2.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.2.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.2.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.2.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.2.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.2.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.20.input_layernorm.weight": "model-00002-of-00002.safetensors",
"model.layers.20.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.20.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.20.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.20.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
"model.layers.20.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.20.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.20.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.20.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.20.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.20.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.21.input_layernorm.weight": "model-00002-of-00002.safetensors",
"model.layers.21.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.21.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.21.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.21.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
"model.layers.21.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
"model.layers.21.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.21.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.21.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
"model.layers.21.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.21.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.22.input_layernorm.weight": "model-00002-of-00002.safetensors",
"model.layers.22.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.22.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.22.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.22.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
"model.layers.22.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
"model.layers.22.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.22.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.22.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
"model.layers.22.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.22.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.23.input_layernorm.weight": "model-00002-of-00002.safetensors",
"model.layers.23.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.23.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.23.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.23.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
"model.layers.23.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
"model.layers.23.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.23.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.23.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
"model.layers.23.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.23.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.24.input_layernorm.weight": "model-00002-of-00002.safetensors",
"model.layers.24.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.24.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.24.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.24.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
"model.layers.24.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
"model.layers.24.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.24.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.24.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
"model.layers.24.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.24.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.25.input_layernorm.weight": "model-00002-of-00002.safetensors",
"model.layers.25.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.25.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.25.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.25.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
"model.layers.25.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
"model.layers.25.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.25.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.25.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
"model.layers.25.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.25.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.26.input_layernorm.weight": "model-00002-of-00002.safetensors",
"model.layers.26.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.26.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.26.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.26.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
"model.layers.26.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
"model.layers.26.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.26.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.26.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
"model.layers.26.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.26.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.27.input_layernorm.weight": "model-00002-of-00002.safetensors",
"model.layers.27.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.27.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.27.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.27.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
"model.layers.27.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
"model.layers.27.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.27.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.27.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
"model.layers.27.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.27.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.28.input_layernorm.weight": "model-00002-of-00002.safetensors",
"model.layers.28.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.28.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.28.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.28.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
"model.layers.28.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
"model.layers.28.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.28.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.28.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
"model.layers.28.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.28.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.29.input_layernorm.weight": "model-00002-of-00002.safetensors",
"model.layers.29.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.29.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.29.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.29.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
"model.layers.29.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
"model.layers.29.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.29.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.29.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
"model.layers.29.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.29.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.3.input_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.3.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.3.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.3.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.3.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.3.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.3.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.3.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.3.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.3.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.3.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.30.input_layernorm.weight": "model-00002-of-00002.safetensors",
"model.layers.30.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.30.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.30.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.30.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
"model.layers.30.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
"model.layers.30.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.30.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.30.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
"model.layers.30.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.30.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.31.input_layernorm.weight": "model-00002-of-00002.safetensors",
"model.layers.31.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.31.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.31.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.31.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
"model.layers.31.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
"model.layers.31.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.31.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.31.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
"model.layers.31.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.31.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.32.input_layernorm.weight": "model-00002-of-00002.safetensors",
"model.layers.32.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.32.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.32.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.32.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
"model.layers.32.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
"model.layers.32.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.32.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.32.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
"model.layers.32.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.32.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.33.input_layernorm.weight": "model-00002-of-00002.safetensors",
"model.layers.33.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.33.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.33.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.33.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
"model.layers.33.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
"model.layers.33.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.33.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.33.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
"model.layers.33.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.33.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.34.input_layernorm.weight": "model-00002-of-00002.safetensors",
"model.layers.34.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.34.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.34.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.34.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
"model.layers.34.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
"model.layers.34.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.34.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.34.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
"model.layers.34.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.34.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.35.input_layernorm.weight": "model-00002-of-00002.safetensors",
"model.layers.35.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.35.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.35.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.35.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
"model.layers.35.self_attn.k_norm.weight": "model-00002-of-00002.safetensors",
"model.layers.35.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.35.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.35.self_attn.q_norm.weight": "model-00002-of-00002.safetensors",
"model.layers.35.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.35.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
"model.layers.4.input_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.4.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.4.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.4.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.4.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.4.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.4.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.4.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.4.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.4.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.4.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.5.input_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.5.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.5.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.5.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.5.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.5.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.5.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.5.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.5.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.5.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.5.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.6.input_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.6.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.6.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.6.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.6.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.6.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.6.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.6.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.6.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.6.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.6.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.7.input_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.7.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.7.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.7.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.7.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.7.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.7.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.7.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.7.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.7.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.7.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.8.input_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.8.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.8.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.8.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.8.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.8.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.8.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.8.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.8.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.8.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.8.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.9.input_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.9.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.9.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.9.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.9.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
"model.layers.9.self_attn.k_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.9.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.9.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.9.self_attn.q_norm.weight": "model-00001-of-00002.safetensors",
"model.layers.9.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
"model.layers.9.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
"model.norm.weight": "model-00002-of-00002.safetensors"
}
}

58
server.py Normal file
View File

@@ -0,0 +1,58 @@
import uvicorn
import json
import asyncio
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
from ziprc import ZIPRCModel, ZIPRCConfig, ZIPRCSampler
# --- Configuration ---
HOST = "0.0.0.0"
PORT = 8000
MODEL_ID = "dataopsnick/Qwen3-4B-Instruct-2507-zip-rc"
# --- Load Model Once ---
print(f"Loading {MODEL_ID}...")
cfg = ZIPRCConfig(model_name=MODEL_ID)
model = ZIPRCModel(cfg)
sampler = ZIPRCSampler(model)
print("Model loaded. Starting server...")
app = FastAPI(title="ZIP-RC OpenAI Compatible API")
@app.post("/v1/chat/completions")
async def chat_completions(request: Request):
"""
Standard OpenAI Chat Completion endpoint.
Streams JSON chunks as Server-Sent Events (SSE).
"""
data = await request.json()
messages = data.get("messages", [])
max_tokens = data.get("max_tokens", 512)
# 1. Use the sampler's generator
stream = sampler.openai(messages, max_tokens=max_tokens)
# 2. Convert to SSE format
async def sse_generator():
async for chunk in stream:
# chunk is an OpenAIObject (dict-like)
payload = json.dumps(dict(chunk))
yield f"data: {payload}\n\n"
yield "data: [DONE]\n\n"
return StreamingResponse(sse_generator(), media_type="text/event-stream")
if __name__ == "__main__":
# Use direct Server instantiation to avoid nested-asyncio conflicts in Notebooks
config = uvicorn.Config(app, host=HOST, port=PORT)
server = uvicorn.Server(config)
try:
# Detect if we are already in an event loop (e.g. Colab/Jupyter)
loop = asyncio.get_running_loop()
print(f"Server started in background task on http://{HOST}:{PORT}")
loop.create_task(server.serve())
except RuntimeError:
# Standard script execution
print(f"Server starting on http://{HOST}:{PORT}")
asyncio.run(server.serve())

31
special_tokens_map.json Normal file
View File

@@ -0,0 +1,31 @@
{
"additional_special_tokens": [
"<|im_start|>",
"<|im_end|>",
"<|object_ref_start|>",
"<|object_ref_end|>",
"<|box_start|>",
"<|box_end|>",
"<|quad_start|>",
"<|quad_end|>",
"<|vision_start|>",
"<|vision_end|>",
"<|vision_pad|>",
"<|image_pad|>",
"<|video_pad|>"
],
"eos_token": {
"content": "<|im_end|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false
},
"pad_token": {
"content": "<|endoftext|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false
}
}

BIN
tokenizer.json (Stored with Git LFS) Normal file
View File

Binary file not shown.

239
tokenizer_config.json Normal file
View File

@@ -0,0 +1,239 @@
{
"add_bos_token": false,
"add_prefix_space": false,
"added_tokens_decoder": {
"151643": {
"content": "<|endoftext|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"151644": {
"content": "<|im_start|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"151645": {
"content": "<|im_end|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"151646": {
"content": "<|object_ref_start|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"151647": {
"content": "<|object_ref_end|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"151648": {
"content": "<|box_start|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"151649": {
"content": "<|box_end|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"151650": {
"content": "<|quad_start|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"151651": {
"content": "<|quad_end|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"151652": {
"content": "<|vision_start|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"151653": {
"content": "<|vision_end|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"151654": {
"content": "<|vision_pad|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"151655": {
"content": "<|image_pad|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"151656": {
"content": "<|video_pad|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"151657": {
"content": "<tool_call>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"151658": {
"content": "</tool_call>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"151659": {
"content": "<|fim_prefix|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"151660": {
"content": "<|fim_middle|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"151661": {
"content": "<|fim_suffix|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"151662": {
"content": "<|fim_pad|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"151663": {
"content": "<|repo_name|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"151664": {
"content": "<|file_sep|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"151665": {
"content": "<tool_response>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"151666": {
"content": "</tool_response>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"151667": {
"content": "<think>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"151668": {
"content": "</think>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
}
},
"additional_special_tokens": [
"<|im_start|>",
"<|im_end|>",
"<|object_ref_start|>",
"<|object_ref_end|>",
"<|box_start|>",
"<|box_end|>",
"<|quad_start|>",
"<|quad_end|>",
"<|vision_start|>",
"<|vision_end|>",
"<|vision_pad|>",
"<|image_pad|>",
"<|video_pad|>"
],
"bos_token": null,
"clean_up_tokenization_spaces": false,
"eos_token": "<|im_end|>",
"errors": "replace",
"extra_special_tokens": {},
"model_max_length": 1010000,
"pad_token": "<|endoftext|>",
"split_special_tokens": false,
"tokenizer_class": "Qwen2Tokenizer",
"unk_token": null
}

1
vocab.json Normal file
View File

File diff suppressed because one or more lines are too long

742
ziprc.py Normal file
View File

@@ -0,0 +1,742 @@
import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer
from dataclasses import dataclass
import copy
import asyncio
import uuid
import time
# Helper to allow dot-notation access (chunk.choices[0].delta.content)
class OpenAIObject(dict):
def __getattr__(self, name):
if name in self:
value = self[name]
if isinstance(value, dict):
return OpenAIObject(value)
if isinstance(value, list):
return [OpenAIObject(v) if isinstance(v, dict) else v for v in value]
return value
raise AttributeError(f"'OpenAIObject' object has no attribute '{name}'")
def __setattr__(self, name, value):
self[name] = value
@dataclass
class ZIPRCConfig:
model_name: str = "dataopsnick/Qwen3-4B-Instruct-2507-zip-rc"
reward_bins: int = 8
length_bins: int = 7
total_zip_tokens: int = 56
zip_start_offset: int = 56
alpha: float = 0.1
beta: float = 0.05
smoothing_window: int = 3
r_boundaries = torch.linspace(0, 1, 9)
l_boundaries = torch.tensor([0, 16, 32, 64, 128, 256, 512, 1024], dtype=torch.float32)
class ZIPRCMath:
@staticmethod
def get_bin_idx(val, boundaries):
for i in range(len(boundaries) - 1):
if boundaries[i] <= val < boundaries[i+1]:
return i
return len(boundaries) - 2
@staticmethod
def apply_horizon_capping(joint_probs, current_len, horizon, config):
"""Eq 25: Collapses mass where length > horizon into a failure state."""
B, R_bins, L_bins = joint_probs.shape
device = joint_probs.device
cutoff_l_idx = L_bins - 1
for i, bound in enumerate(config.l_boundaries):
if bound > horizon:
cutoff_l_idx = max(0, i - 1)
break
capped_probs = joint_probs.clone()
valid_mask = torch.zeros((L_bins), dtype=torch.bool, device=device)
valid_mask[:cutoff_l_idx+1] = True
kept_mass = capped_probs[:, :, valid_mask].sum(dim=(1, 2))
pruned_mass = 1.0 - kept_mass
capped_probs[:, :, ~valid_mask] = 0.0
capped_probs[:, 0, cutoff_l_idx] += pruned_mass
return capped_probs
@staticmethod
def get_marginals(joint_probs):
q_v = joint_probs.sum(dim=2)
q_l = joint_probs.sum(dim=1)
return q_v, q_l
@staticmethod
def compute_expected_max_value(marginals_list, values_per_bin):
if not marginals_list: return 0.0
stacked_marginals = torch.cat(marginals_list, dim=0)
cdfs = torch.cumsum(stacked_marginals, dim=1)
f_max = torch.prod(cdfs, dim=0)
f_max_shifted = torch.roll(f_max, 1)
f_max_shifted[0] = 0.0
p_max = f_max - f_max_shifted
expected_max = torch.sum(p_max * values_per_bin).item()
return expected_max
@staticmethod
def compute_sampling_utility(candidates, config):
"""Eq 19: Utility optimization for shared horizon."""
if not candidates: return -1e9
device = candidates[0]['joint_probs'].device
r_vals = (config.r_boundaries[:-1] + config.r_boundaries[1:]).to(device) / 2
l_vals = (config.l_boundaries[:-1] + config.l_boundaries[1:]).to(device) / 2
sum_est_total_len = 0.0
for cand in candidates:
qv, ql = ZIPRCMath.get_marginals(cand['joint_probs'].unsqueeze(0))
e_rem_len = torch.sum(ql * l_vals).item()
curr_prefix_len = cand['ids'].shape[1]
sum_est_total_len += (curr_prefix_len + e_rem_len)
b_bar = max(sum_est_total_len / len(candidates) if len(candidates) > 0 else 1.0, 1.0)
beta_tilde = config.beta / b_bar
best_util = -float('inf')
search_space = config.l_boundaries.tolist() + [2048]
for h in search_space:
h = int(h)
q_v_list, q_l_list = [], []
for cand in candidates:
capped_joint = ZIPRCMath.apply_horizon_capping(
cand['joint_probs'].unsqueeze(0), cand['current_len'], h, config
)
qv, ql = ZIPRCMath.get_marginals(capped_joint)
q_v_list.append(qv)
q_l_list.append(ql)
e_max_reward = ZIPRCMath.compute_expected_max_value(q_v_list, r_vals)
e_latency = ZIPRCMath.compute_expected_max_value(q_l_list, l_vals)
total_compute = sum(torch.sum(ql * l_vals).item() for ql in q_l_list)
cost = beta_tilde * (config.alpha * total_compute + (1 - config.alpha) * e_latency)
util = e_max_reward - cost
if util > best_util: best_util = util
return best_util
class PredictionBuffer:
def __init__(self, window_size):
self.window = window_size
self.history = []
def add(self, prob_tensor):
self.history.append(prob_tensor)
if len(self.history) > self.window: self.history.pop(0)
def get_smoothed(self):
stack = torch.stack(self.history)
return torch.mean(stack, dim=0)
class ZIPRCModel(torch.nn.Module):
def __init__(self, config: ZIPRCConfig):
super().__init__()
self.config = config
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.base_model = AutoModelForCausalLM.from_pretrained(
config.model_name, torch_dtype=torch.bfloat16, device_map=self.device
)
self.tokenizer = AutoTokenizer.from_pretrained(config.model_name)
self.zip_start_id = self.base_model.config.vocab_size - config.zip_start_offset
self.base_model.eval()
def get_joint_distribution(self, logits):
zip_logits = logits[:, self.zip_start_id : self.zip_start_id + self.config.total_zip_tokens]
probs = F.softmax(zip_logits, dim=-1)
return probs.view(-1, self.config.reward_bins, self.config.length_bins)
class ZIPRCSampler:
def __init__(self, model):
self.model = model
self.config = model.config
def select_best_trajectory(self, trajectories):
if not trajectories: return None
best_traj = None
best_score = -float('inf')
device = trajectories[0]['joint_probs'].device
r_vals = (self.config.r_boundaries[:-1] + self.config.r_boundaries[1:]).to(device) / 2
for traj in trajectories:
qv, _ = ZIPRCMath.get_marginals(traj['joint_probs'].unsqueeze(0))
score = torch.sum(qv * r_vals).item()
traj['final_score'] = score
if score > best_score:
best_score = score
best_traj = traj
return best_traj
async def openai(self, messages, max_tokens=512, initial_samples=2):
"""
Async generator that yields OpenAI-compatible chunks with added ZIP-RC introspection data.
"""
# 1. Handle Input (String or Messages List)
if isinstance(messages, str):
prompt = messages
else:
# Assumes generic chat template
prompt = self.model.tokenizer.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True
)
# 2. Setup Candidates
input_ids = self.model.tokenizer(prompt, return_tensors="pt").input_ids.to(self.model.device)
candidates = []
for i in range(initial_samples):
candidates.append({
'id': i, 'ids': input_ids.clone(), 'finished': False,
'buffer': PredictionBuffer(self.config.smoothing_window),
'joint_probs': None, 'current_len': 0
})
finished_trajectories = []
chat_id = f"chatcmpl-{uuid.uuid4()}"
created_ts = int(time.time())
model_name = self.config.model_name
# State for delta streaming
last_top1_id = -1
last_top1_len = input_ids.shape[1]
for step in range(max_tokens):
if not candidates: break
# --- [A] MODEL FORWARD (Async Wrapper) ---
active_ids = torch.cat([c['ids'] for c in candidates], dim=0)
# Simple wrapper to allow event loop to breathe
await asyncio.sleep(0)
with torch.no_grad():
outputs = self.model.base_model(active_ids)
next_token_logits = outputs.logits[:, -1, :]
raw_joint = self.model.get_joint_distribution(next_token_logits)
# --- Update Candidates ---
for i, c in enumerate(candidates):
c['buffer'].add(raw_joint[i])
c['joint_probs'] = c['buffer'].get_smoothed()
c['current_len'] = step
valid_logits = next_token_logits[i].clone()
valid_logits[self.model.zip_start_id : self.model.zip_start_id + self.config.total_zip_tokens] = -float('inf')
probs = F.softmax(valid_logits, dim=-1)
next_token = torch.multinomial(probs, 1).unsqueeze(0)
c['ids'] = torch.cat([c['ids'], next_token], dim=1)
if next_token.item() == self.model.tokenizer.eos_token_id:
c['finished'] = True
finished_trajectories.append(c)
candidates = [c for c in candidates if not c['finished']]
if not candidates: break
# --- [B] META-ACTIONS ---
cand_metrics = []
r_vals = (self.config.r_boundaries[:-1] + self.config.r_boundaries[1:]).to(self.model.device)/2
for i, c in enumerate(candidates):
qv, _ = ZIPRCMath.get_marginals(c['joint_probs'].unsqueeze(0))
e_r = torch.sum(qv * r_vals).item()
cand_metrics.append((i, e_r))
sorted_by_reward = sorted(cand_metrics, key=lambda x: x[1], reverse=True)
top_indices = [x[0] for x in sorted_by_reward]
possible_actions = [('keep', candidates)]
MAX_SAMPLES = 8
if len(candidates) < MAX_SAMPLES:
top_idx = top_indices[0]
new_set = copy.deepcopy(candidates)
clone = copy.deepcopy(new_set[top_idx])
clone['id'] = max([c['id'] for c in new_set], default=0) + 1
new_set.append(clone)
possible_actions.append(('branch_top1', new_set))
if len(candidates) >= 2 and len(candidates) + 1 < MAX_SAMPLES:
new_set_b2 = copy.deepcopy(candidates)
clone2 = copy.deepcopy(new_set_b2[top_indices[1]])
clone2['id'] = max([c['id'] for c in new_set_b2], default=0) + 1
new_set_b2.append(clone2)
possible_actions.append(('branch_top2', new_set_b2))
if len(candidates) > 1:
worst_idx = top_indices[-1]
new_set = [c for i, c in enumerate(candidates) if i != worst_idx]
possible_actions.append(('prune_bot1', new_set))
if len(candidates) > 1 and top_indices[0] != top_indices[-1]:
top_id = candidates[top_indices[0]]['id']
worst_idx = top_indices[-1]
new_set = copy.deepcopy(candidates)
new_set = [c for i, c in enumerate(new_set) if i != worst_idx]
source = next(c for c in new_set if c['id'] == top_id)
clone = copy.deepcopy(source)
clone['id'] = max([c['id'] for c in new_set], default=0) + 1
new_set.append(clone)
possible_actions.append(('swap', new_set))
best_action_name, best_util, best_next_candidates = 'keep', -float('inf'), candidates
for name, cand_set in possible_actions:
if not cand_set: continue
penalty = 0.0 if name == 'keep' else 0.01
util = ZIPRCMath.compute_sampling_utility(cand_set, self.config) - penalty
if util > best_util:
best_util, best_action_name, best_next_candidates = util, name, cand_set
candidates = best_next_candidates
# --- [C] PREPARE INTROSPECTION PAYLOAD ---
vis_metrics = []
for c in candidates:
qv, _ = ZIPRCMath.get_marginals(c['joint_probs'].unsqueeze(0))
e_r = torch.sum(qv * r_vals).item()
vis_metrics.append((c, e_r))
vis_sorted = sorted(vis_metrics, key=lambda x: x[1], reverse=True)
lhs_c = vis_sorted[0][0] if len(vis_sorted) > 0 else None
rhs_c = vis_sorted[1][0] if len(vis_sorted) > 1 else None
lhs_score = vis_sorted[0][1] if len(vis_sorted) > 0 else 0.0
rhs_score = vis_sorted[1][1] if len(vis_sorted) > 1 else 0.0
def get_text(c_obj):
if not c_obj: return ""
curr_ids = c_obj['ids'][0]
full_text = self.model.tokenizer.decode(curr_ids, skip_special_tokens=True)
return full_text
lhs_text = get_text(lhs_c)
rhs_text = get_text(rhs_c)
delta_content = ""
if lhs_c and lhs_c['id'] == last_top1_id:
new_len = len(lhs_text)
if new_len > last_top1_len:
delta_content = lhs_text[last_top1_len:]
last_top1_len = new_len
elif lhs_c:
last_top1_id = lhs_c['id']
last_top1_len = len(lhs_text)
delta_content = ""
chunk_dict = {
"id": chat_id,
"object": "chat.completion.chunk",
"created": created_ts,
"model": model_name,
"choices": [{"index": 0, "delta": {"content": delta_content}, "finish_reason": None}],
"zip_rc": {
"step": step,
"action": best_action_name,
"utility": best_util,
"lhs_text": lhs_text,
"rhs_text": rhs_text,
"lhs_score": lhs_score,
"rhs_score": rhs_score,
"lhs_id": lhs_c['id'] if lhs_c else -1,
"rhs_id": rhs_c['id'] if rhs_c else -1
}
}
yield OpenAIObject(chunk_dict)
# Calculate Final Best Answer (clean from swaps/backtracks)
# Include running candidates in case max_tokens was hit before EOS
all_trajs = finished_trajectories + candidates
best_traj = self.select_best_trajectory(all_trajs)
final_answer = ""
if best_traj:
# Decode only the generated response (exclude prompt)
prompt_len = input_ids.shape[1]
final_ids = best_traj['ids'][0][prompt_len:]
final_answer = self.model.tokenizer.decode(final_ids, skip_special_tokens=True)
yield OpenAIObject({
"id": chat_id,
"object": "chat.completion.chunk",
"created": created_ts,
"model": model_name,
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
"zip_rc": {
"action": "finished",
"final_text": final_answer
}
})
def generate_stream(self, prompt, max_new_tokens=512, initial_samples=2):
# Setup
input_ids = self.model.tokenizer.apply_chat_template(
[{"role": "user", "content": prompt}], tokenize=True, return_tensors="pt", add_generation_prompt=True
).to(self.model.device)
candidates = []
for i in range(initial_samples):
candidates.append({
'id': i, 'ids': input_ids.clone(), 'finished': False,
'buffer': PredictionBuffer(self.config.smoothing_window),
'joint_probs': None, 'current_len': 0
})
finished_trajectories = []
text_cache = {}
# UI State
dashboard_widget = None
last_cli_height = 0
# Check environment for widget vs CLI
try:
import ipywidgets as widgets
from IPython.display import display
ENV_MODE = 'notebook'
except ImportError:
ENV_MODE = 'cli'
if ENV_MODE == 'notebook':
import html
from tqdm.notebook import tqdm
dashboard_widget = widgets.HTML(value="Initialization...")
display(dashboard_widget)
pbar = tqdm(total=max_new_tokens, display=False)
else:
import sys, shutil, textwrap
from tqdm import tqdm
pbar = tqdm(total=max_new_tokens, dynamic_ncols=True, bar_format='{bar}| {n_fmt}/{total_fmt}')
try:
for step in range(max_new_tokens):
if not candidates: break
# --- [A] MODEL FORWARD ---
active_ids = torch.cat([c['ids'] for c in candidates], dim=0)
with torch.no_grad():
outputs = self.model.base_model(active_ids)
next_token_logits = outputs.logits[:, -1, :]
raw_joint = self.model.get_joint_distribution(next_token_logits)
for i, c in enumerate(candidates):
c['buffer'].add(raw_joint[i])
c['joint_probs'] = c['buffer'].get_smoothed()
c['current_len'] = step
valid_logits = next_token_logits[i].clone()
valid_logits[self.model.zip_start_id : self.model.zip_start_id + self.config.total_zip_tokens] = -float('inf')
probs = F.softmax(valid_logits, dim=-1)
next_token = torch.multinomial(probs, 1).unsqueeze(0)
c['ids'] = torch.cat([c['ids'], next_token], dim=1)
if next_token.item() == self.model.tokenizer.eos_token_id:
c['finished'] = True
finished_trajectories.append(c)
candidates = [c for c in candidates if not c['finished']]
if not candidates: break
# --- [B] META-ACTIONS (Restored Full Logic) ---
# 1. Metric Calculation & Sorting
cand_metrics = []
r_vals = (self.config.r_boundaries[:-1] + self.config.r_boundaries[1:]).to(self.model.device)/2
for i, c in enumerate(candidates):
qv, _ = ZIPRCMath.get_marginals(c['joint_probs'].unsqueeze(0))
e_r = torch.sum(qv * r_vals).item()
cand_metrics.append((i, e_r))
# Sort to identify Top-1, Top-2, and Worst
sorted_by_reward = sorted(cand_metrics, key=lambda x: x[1], reverse=True)
top_indices = [x[0] for x in sorted_by_reward]
# 2. Define Possible Actions
possible_actions = [('keep', candidates)]
MAX_SAMPLES = 8
# Action: Branch Top-1
if len(candidates) < MAX_SAMPLES:
top_idx = top_indices[0]
new_set = copy.deepcopy(candidates)
clone = copy.deepcopy(new_set[top_idx])
clone['id'] = max([c['id'] for c in new_set], default=0) + 1
new_set.append(clone)
possible_actions.append(('branch_top1', new_set))
# Action: Branch Top-2
if len(candidates) >= 2 and len(candidates) + 1 < MAX_SAMPLES:
new_set2 = copy.deepcopy(new_set) # Base off the set that already branched top1? No, independent action in original code.
# Original code treats them as distinct alternative meta-actions for the step.
# Re-building from clean candidates for branch_top2:
new_set_b2 = copy.deepcopy(candidates)
clone2 = copy.deepcopy(new_set_b2[top_indices[1]])
clone2['id'] = max([c['id'] for c in new_set_b2], default=0) + 1
new_set_b2.append(clone2)
possible_actions.append(('branch_top2', new_set_b2))
# Action: Prune Worst 1
if len(candidates) > 1:
worst_idx = top_indices[-1]
new_set = [c for i, c in enumerate(candidates) if i != worst_idx]
possible_actions.append(('prune_bot1', new_set))
# Action: Prune Worst 2
if len(candidates) > 2:
worst_indices = set(top_indices[-2:])
new_set = [c for i, c in enumerate(candidates) if i not in worst_indices]
possible_actions.append(('prune_bot2', new_set))
# Action: Swap (Prune Worst, Branch Best)
if len(candidates) > 1 and top_indices[0] != top_indices[-1]:
top_id = candidates[top_indices[0]]['id']
worst_idx = top_indices[-1]
new_set = copy.deepcopy(candidates)
new_set = [c for i, c in enumerate(new_set) if i != worst_idx]
source = next(c for c in new_set if c['id'] == top_id)
clone = copy.deepcopy(source)
clone['id'] = max([c['id'] for c in new_set], default=0) + 1
new_set.append(clone)
possible_actions.append(('swap', new_set))
# 3. Select Best Action via Utility
best_action_name, best_util, best_next_candidates = 'keep', -float('inf'), candidates
for name, cand_set in possible_actions:
if not cand_set: continue
penalty = 0.0 if name == 'keep' else 0.01
util = ZIPRCMath.compute_sampling_utility(cand_set, self.config) - penalty
if util > best_util:
best_util, best_action_name, best_next_candidates = util, name, cand_set
# Apply selection
candidates = best_next_candidates
# Re-evaluate metrics for visualization (indices might have shifted or sizes changed)
# We need to find the new Top-1 and Top-2 to display in the UI.
vis_metrics = []
for i, c in enumerate(candidates):
qv, _ = ZIPRCMath.get_marginals(c['joint_probs'].unsqueeze(0))
e_r = torch.sum(qv * r_vals).item()
vis_metrics.append((c, e_r))
vis_sorted = sorted(vis_metrics, key=lambda x: x[1], reverse=True)
# --- [C] ADVANCED DASHBOARD RENDERING ---
pbar.update(1)
# 1. Prepare Text (Always show current best 2)
lhs_c = vis_sorted[0][0] if len(vis_sorted) > 0 else None
rhs_c = vis_sorted[1][0] if len(vis_sorted) > 1 else None
def get_text(c_obj):
if not c_obj: return ""
curr_ids = c_obj['ids'][0]
if len(curr_ids) > text_cache.get(c_obj['id'], (None, 0))[1]:
full_text = self.model.tokenizer.decode(curr_ids, skip_special_tokens=True)
text_cache[c_obj['id']] = (full_text, len(curr_ids))
return full_text
return text_cache[c_obj['id']][0]
l_raw = get_text(lhs_c).replace('\n', ' ')
r_raw = get_text(rhs_c).replace('\n', ' ')
# Filter Action Display based on Prompt Requirements
# "only show branch_top1 and branch_top2 updates in the streaming view"
if best_action_name in ['branch_top1', 'branch_top2']:
display_action = best_action_name
else:
display_action = "" # Hide keep/prune/swap from the prominent label to reduce noise
if ENV_MODE == 'notebook':
l_esc = html.escape(l_raw[-2000:])
r_esc = html.escape(r_raw[-2000:])
lhs_head = f"LHS (Top-1) [ID:{lhs_c['id']}]"
rhs_head = f"RHS (Top-2) [ID:{rhs_c['id'] if rhs_c else '-'}]"
css = """
<style>
.zip-dash { font-family: 'Courier New', monospace; background: #1e1e1e; color: #d4d4d4; padding: 10px; border: 1px solid #333; }
.zip-header { border-bottom: 2px solid #444; padding-bottom: 5px; margin-bottom: 5px; font-weight: bold; }
.zip-meta { color: #569cd6; border-bottom: 1px dashed #444; margin-bottom: 10px; padding-bottom: 5px;}
.zip-grid { display: grid; grid-template-columns: 1fr 1px 1fr; gap: 10px; }
.zip-col { white-space: pre-wrap; word-wrap: break-word; overflow-y: hidden; max-height: 1000px; line-height: 1.3; }
.zip-sep { background: #444; height: 100%; }
</style>
"""
body = f"""
<div class="zip-dash">
<div class="zip-header">
<div style="display:flex; justify-content:space-between;">
<span>{lhs_head}</span>
<span>{rhs_head}</span>
</div>
</div>
<div class="zip-meta">
Step {step} | Action: {display_action} | Util: {best_util:.4f} | Candidates: {len(candidates)}
</div>
<div class="zip-grid">
<div class="zip-col">{l_esc}</div>
<div class="zip-sep"></div>
<div class="zip-col">{r_esc}</div>
</div>
</div>
"""
dashboard_widget.value = css + body
else:
# CLI Renderer
term_w = shutil.get_terminal_size((100, 24)).columns
col_w = (term_w - 7) // 2
line_sep = "=" * term_w
line_head = f"{f' LHS (Top-1) [ID:{lhs_c["id"]}]':<{col_w}} || {f' RHS (Top-2) [ID:{rhs_c["id"] if rhs_c else "-"}]':<{col_w}}"
line_meta = f" Step {step:<4} | Action: {display_action:<11} | Util: {best_util:.4f} | Pool: {len(candidates)}"
line_meta = f"{line_meta:<{col_w}} || "
tail_len = col_w * 12
l_lines = textwrap.wrap(l_raw[-tail_len:], width=col_w)
r_lines = textwrap.wrap(r_raw[-tail_len:], width=col_w)
max_h = max(len(l_lines), len(r_lines), 1)
l_lines += [" " * col_w] * (max_h - len(l_lines))
r_lines += [" " * col_w] * (max_h - len(r_lines))
buffer = []
buffer.append(line_sep)
buffer.append(line_head)
buffer.append(line_sep)
buffer.append(line_meta)
buffer.append(line_sep)
for i in range(max_h):
buffer.append(f"{l_lines[i]:<{col_w}} || {r_lines[i]:<{col_w}}")
buffer.append(line_sep)
final_str = "\n".join(buffer) + "\n"
if last_cli_height > 0:
sys.stdout.write(f"\033[{last_cli_height}A")
sys.stdout.write("\033[J")
sys.stdout.write(final_str)
sys.stdout.flush()
last_cli_height = len(buffer) + 1
finally:
pbar.close()
# Final Summary
best_traj = self.select_best_trajectory(finished_trajectories)
score = best_traj['final_score'] if best_traj else 0.0
full_text = self.model.tokenizer.decode(best_traj['ids'][0], skip_special_tokens=True)
answer = full_text.split("assistant")[-1].strip() if "assistant" in full_text else full_text
print("\n" + "=" * 100)
print(f"Best Trajectory (Top-1) Confidence: {score:.2%}")
print("=" * 100 + "\n")
print(answer)
print("\n" + "=" * 100)
return finished_trajectories
def generate(self, prompt, max_new_tokens=512, initial_samples=2):
input_ids = self.model.tokenizer.apply_chat_template(
[{"role": "user", "content": prompt}], tokenize=True, return_tensors="pt", add_generation_prompt=True
).to(self.model.device)
candidates = []
for i in range(initial_samples):
candidates.append({
'id': i, 'ids': input_ids.clone(), 'finished': False,
'buffer': PredictionBuffer(self.config.smoothing_window),
'joint_probs': None, 'current_len': 0
})
finished_trajectories = []
for step in range(max_new_tokens):
if not candidates: break
active_ids = torch.cat([c['ids'] for c in candidates], dim=0)
with torch.no_grad():
outputs = self.model.base_model(active_ids)
next_token_logits = outputs.logits[:, -1, :]
raw_joint = self.model.get_joint_distribution(next_token_logits)
for i, c in enumerate(candidates):
c['buffer'].add(raw_joint[i])
c['joint_probs'] = c['buffer'].get_smoothed()
c['current_len'] = step
valid_logits = next_token_logits[i].clone()
valid_logits[self.model.zip_start_id : self.model.zip_start_id + self.config.total_zip_tokens] = -float('inf')
probs = F.softmax(valid_logits, dim=-1)
next_token = torch.multinomial(probs, 1).unsqueeze(0)
c['ids'] = torch.cat([c['ids'], next_token], dim=1)
if next_token.item() == self.model.tokenizer.eos_token_id:
c['finished'] = True
finished_trajectories.append(c)
candidates = [c for c in candidates if not c['finished']]
if not candidates: break
# Meta-Actions (Branching/Pruning/Swapping)
cand_metrics = []
r_vals = (self.config.r_boundaries[:-1] + self.config.r_boundaries[1:]).to(self.model.device)/2
for i, c in enumerate(candidates):
qv, _ = ZIPRCMath.get_marginals(c['joint_probs'].unsqueeze(0))
e_r = torch.sum(qv * r_vals).item()
cand_metrics.append((i, e_r))
sorted_by_reward = sorted(cand_metrics, key=lambda x: x[1], reverse=True)
top_indices = [x[0] for x in sorted_by_reward]
possible_actions = [('keep', candidates)]
MAX_SAMPLES = 8
if len(candidates) < MAX_SAMPLES:
top_idx = top_indices[0]
new_set = copy.deepcopy(candidates)
clone = copy.deepcopy(new_set[top_idx])
clone['id'] = max([c['id'] for c in new_set], default=0) + 1
new_set.append(clone)
possible_actions.append(('branch_top1', new_set))
if len(candidates) >= 2 and len(candidates) + 1 < MAX_SAMPLES:
new_set2 = copy.deepcopy(new_set)
clone2 = copy.deepcopy(new_set2[top_indices[1]])
clone2['id'] = max([c['id'] for c in new_set2], default=0) + 1
new_set2.append(clone2)
possible_actions.append(('branch_top2', new_set2))
if len(candidates) > 1:
worst_idx = top_indices[-1]
new_set = [c for i, c in enumerate(candidates) if i != worst_idx]
possible_actions.append(('prune_bot1', new_set))
if len(candidates) > 2:
worst_indices = set(top_indices[-2:])
new_set = [c for i, c in enumerate(candidates) if i not in worst_indices]
possible_actions.append(('prune_bot2', new_set))
if len(candidates) > 1 and top_indices[0] != top_indices[-1]:
top_id = candidates[top_indices[0]]['id']
worst_idx = top_indices[-1]
new_set = copy.deepcopy(candidates)
new_set = [c for i, c in enumerate(new_set) if i != worst_idx]
source = next(c for c in new_set if c['id'] == top_id)
clone = copy.deepcopy(source)
clone['id'] = max([c['id'] for c in new_set], default=0) + 1
new_set.append(clone)
possible_actions.append(('swap', new_set))
best_action_name, best_util, best_next_candidates = 'keep', -float('inf'), candidates
for name, cand_set in possible_actions:
if not cand_set: continue
penalty = 0.0 if name == 'keep' else 0.01
util = ZIPRCMath.compute_sampling_utility(cand_set, self.config) - penalty
if util > best_util:
best_util, best_action_name, best_next_candidates = util, name, cand_set
if best_action_name != 'keep':
print(f"Step {step}: Meta-Action -> {best_action_name} (Util: {best_util:.4f}) | Pool: {len(best_next_candidates)}")
candidates = best_next_candidates
return finished_trajectories