初始化项目，由ModelHub XC社区提供模型
Model: dataopsnick/Qwen3-4B-Instruct-2507-zip-rc Source: Original Platform
This commit is contained in:
ModelHub XC
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