初始化项目，由ModelHub XC社区提供模型
Model: dataopsnick/Qwen3-4B-Instruct-2507-zip-rc Source: Original Platform
This commit is contained in:
ModelHub XC
36
.gitattributes
vendored
Normal file
36
.gitattributes
vendored
Normal 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
275
README.md
Normal 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** |
|
||||
| :--- | :--- |
|
||||
| [](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. |
|
||||
| [](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
28
added_tokens.json
Normal 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
61
chat_template.jinja
Normal 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
68
config.json
Normal 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
13
generation_config.json
Normal 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
151388
merges.txt
Normal file
File diff suppressed because it is too large
Load Diff
3
model-00001-of-00002.safetensors
Normal file
3
model-00001-of-00002.safetensors
Normal 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
3
model-00002-of-00002.safetensors
Normal 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
406
model.safetensors.index.json
Normal 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
58
server.py
Normal 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
31
special_tokens_map.json
Normal 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
BIN
tokenizer.json
(Stored with Git LFS)
Normal file
Binary file not shown.
239
tokenizer_config.json
Normal file
239
tokenizer_config.json
Normal 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
1
vocab.json
Normal file
File diff suppressed because one or more lines are too long
742
ziprc.py
Normal file
742
ziprc.py
Normal 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
|
||||
Reference in New Issue
Block a user