import collections
import json
import os
import random
import re
from pathlib import Path
from huggingface_hub import hf_hub_download, list_repo_files
import pyarrow.parquet as pq
os.makedirs("/content/calib_parts", exist_ok=True)
if not Path("/content/calib_parts/ukisai/qwen38-agent-sft-combined-100pct.parquet").exists():
hf_hub_download("ukisai/Qwen3.8-27B-multi-turn-agent-sft", "qwen38-agent-sft-combined-100pct.parquet", repo_type="dataset", local_dir="/content/calib_parts/ukisai")
if not Path("/content/calib_parts/whittle/data/train-00000-of-00051.parquet").exists():
hf_hub_download("logic65/whittle-distill-qwen3.8-27b", "data/train-00000-of-00051.parquet", repo_type="dataset", local_dir="/content/calib_parts/whittle")
if not Path("/content/calib_parts/chitako/calibration_ja.txt").exists():
hf_hub_download("ChiTako/japanese-imatrix-calibration", "calibration_ja.txt", repo_type="dataset", local_dir="/content/calib_parts/chitako")
if not Path("/content/calib_parts/hotchpotch/data/test-00000-of-00001.parquet").exists():
hf_hub_download("hotchpotch/japanese-qa-reasoning-100k", "data/test-00000-of-00001.parquet", repo_type="dataset", local_dir="/content/calib_parts/hotchpotch")
rng = random.Random(20261014)
# 1. Load Ukisai and deduplicate by trial_name (keep final episode), then group-split by base task ID
tbl_uki = pq.read_table("/content/calib_parts/ukisai/qwen38-agent-sft-combined-100pct.parquet")
n_uki = tbl_uki.num_rows
tasks_col = tbl_uki["task"].to_pylist()
trials_col = tbl_uki["trial_name"].to_pylist()
episodes_col = tbl_uki["episode"].to_pylist()
def ep_num(ep_str):
m = re.search(r"(\d+)$", str(ep_str))
return int(m.group(1)) if m else 0
def get_base_task(t_str):
s = str(t_str)
s = re.sub(r"(__.*|_run\d+)$", "", s)
return s
# Keep only the max episode index per trial_name
best_row_for_trial = {}
for idx in range(n_uki):
tr = trials_col[idx]
en = ep_num(episodes_col[idx])
if tr not in best_row_for_trial or en > best_row_for_trial[tr][0]:
best_row_for_trial[tr] = (en, idx)
# Group trials by base_task ( keep only 1 trial per base_task so zero base_task duplicates! )
task_to_row = {}
for tr, (_en, idx) in best_row_for_trial.items():
bt = get_base_task(tasks_col[idx])
if bt not in task_to_row:
task_to_row[bt] = idx
nl2bash_tasks = sorted([bt for bt in task_to_row if bt.startswith("nl2bash")])
other_tasks = sorted([bt for bt in task_to_row if not bt.startswith("nl2bash")])
rng.shuffle(nl2bash_tasks)
rng.shuffle(other_tasks)
# Disjoint split by base_task:
calib_uki_tasks = nl2bash_tasks[:240] + other_tasks[:60]
eval_uki_tasks = nl2bash_tasks[240:275] + other_tasks[60:70]
assert set(calib_uki_tasks).isdisjoint(set(eval_uki_tasks))
rng.shuffle(calib_uki_tasks)
rng.shuffle(eval_uki_tasks)
def fmt_uki_row(row_idx, keep_preamble=False):
conv = tbl_uki["conversations"][row_idx].as_py()
parts = []
for m_i, msg in enumerate(conv):
role = msg.get("role", "user")
content = str(msg.get("content", ""))
if m_i == 0 and not keep_preamble and "Task Description:" in content:
content = "Task Description:" + content.split("Task Description:", 1)[1]
parts.append(f"<|im_start|>{role}\n{content}<|im_end|>")
full = "\n".join(parts)
# Keep tail if very long so multi-turn interactions and final turns are preserved
if len(full) > 5000:
full = full[:800] + "\n...\n" + full[-4200:]
return full
uki_calib_docs = [fmt_uki_row(task_to_row[bt], keep_preamble=(i < 2)) for i, bt in enumerate(calib_uki_tasks)]
uki_eval_docs = [fmt_uki_row(task_to_row[bt], keep_preamble=False) for bt in eval_uki_tasks]
# 2. Terminal-Bench 2.1: use xhigh/traces/<task>/prompt.txt, split strictly by <task>
files_tb = list_repo_files("Lottolabs/terminal-bench-2.1-qwen3.8-27b-traces", repo_type="dataset")
def get_tb_task_name(f):
return re.sub(r"-[0-9a-f]{8,}$", "", f.split("/")[2])
tb_by_task = {}
for f in sorted(files_tb):
if f.startswith("xhigh/traces/") and f.endswith("/prompt.txt"):
tname = get_tb_task_name(f)
if tname not in tb_by_task:
tb_by_task[tname] = f
xhigh_prompts = list(tb_by_task.values())
rng.shuffle(xhigh_prompts)
tb_calib_files = xhigh_prompts[:24]
tb_eval_files = xhigh_prompts[24:32]
assert set(map(get_tb_task_name, tb_calib_files)).isdisjoint(set(map(get_tb_task_name, tb_eval_files)))
def clean_tb_prompt(txt: str) -> str:
if "\n\nExecution contract:" in txt:
txt = txt.split("\n\nExecution contract:", 1)[0].strip()
return txt[-4500:] if len(txt) > 4500 else txt
tb_calib_docs = []
for f in tb_calib_files:
p = hf_hub_download("Lottolabs/terminal-bench-2.1-qwen3.8-27b-traces", f, repo_type="dataset", local_dir="/content/calib_parts/tb")
txt = Path(p).read_text(encoding="utf-8", errors="replace")
# Keep tb_calib_docs identical to Phase 1a so calibration_mix_v2.txt is 100% reproducible
tb_calib_docs.append(txt[-4500:] if len(txt) > 4500 else txt)
tb_eval_docs = []
for f in tb_eval_files:
p = hf_hub_download("Lottolabs/terminal-bench-2.1-qwen3.8-27b-traces", f, repo_type="dataset", local_dir="/content/calib_parts/tb")
txt = Path(p).read_text(encoding="utf-8", errors="replace")
tb_eval_docs.append(clean_tb_prompt(txt))
# 3. Whittle: use prompt + response directly (already templated!), stratify by kind, truncate prompt from left
tbl_whi = pq.read_table("/content/calib_parts/whittle/data/train-00000-of-00051.parquet")
whi_by_kind = collections.defaultdict(list)
for idx in range(tbl_whi.num_rows):
k = str(tbl_whi["kind"][idx].as_py())
whi_by_kind[k].append(idx)
for k in whi_by_kind:
rng.shuffle(whi_by_kind[k])
def fmt_whi_row(idx):
p = str(tbl_whi["prompt"][idx].as_py() or "")
r = str(tbl_whi["response"][idx].as_py() or "")
if not r.strip():
reas = str(tbl_whi["reasoning"][idx].as_py() or "")
ans = str(tbl_whi["answer"][idx].as_py() or "")
r = f"{reas}\n</think>\n{ans}<|im_end|>"
if len(p) > 4200:
# Keep system header + tail of prompt containing task & latest tool output
p = p[:500] + "\n...\n" + p[-3700:]
if len(r) > 3800:
r = r[:1800] + "\n...\n" + r[-2000:]
return p + r
# Stratified selection: all agentx (69), agent1 (39), json (47), longctx (64), chat (30), plus 24 math/gsm, 40 code
whi_calib_idxs = []
whi_eval_idxs = []
quotas = [("agentx", 55, 12), ("agent1", 31, 8), ("json", 38, 9), ("longctx", 52, 12), ("chat", 24, 6), ("code", 35, 8), ("math", 15, 4), ("gsm", 10, 2)]
for k, n_cal, n_ev in quotas:
pool = whi_by_kind.get(k, [])
whi_calib_idxs.extend(pool[:n_cal])
whi_eval_idxs.extend(pool[n_cal : n_cal + n_ev])
assert set(whi_calib_idxs).isdisjoint(set(whi_eval_idxs))
rng.shuffle(whi_calib_idxs)
rng.shuffle(whi_eval_idxs)
whi_calib_docs = [fmt_whi_row(i) for i in whi_calib_idxs]
whi_eval_docs = [fmt_whi_row(i) for i in whi_eval_idxs]
# 4. Japanese:
# - Calibration JA = deduplicated ChiTako + hotchpotch[:200] (by distinct url/id)
# - Held-out Eval JA = hotchpotch[200:260] (strictly url/id disjoint, and 0 shingle overlap with ChiTako!)
ja_raw = Path("/content/calib_parts/chitako/calibration_ja.txt").read_text(encoding="utf-8", errors="replace")
raw_paras = [b.strip() for b in ja_raw.split("\n\n") if len(b.strip()) > 120]
def get_shingles(s, k=20):
s_norm = re.sub(r"\s+", " ", s)
if len(s_norm) < k:
return {s_norm}
return {s_norm[i : i + k] for i in range(0, len(s_norm) - k + 1, 5)}
dedup_chitako = []
seen_shingles = set()
for p in raw_paras:
sh = get_shingles(p)
if len(sh & seen_shingles) / max(1, len(sh)) > 0.30:
continue
seen_shingles |= sh
dedup_chitako.append(p[:2000])
rng.shuffle(dedup_chitako)
print(f"ChiTako raw paras: {len(raw_paras)} -> deduplicated: {len(dedup_chitako)}")
tbl_hp = pq.read_table("/content/calib_parts/hotchpotch/data/test-00000-of-00001.parquet")
url_to_hp = {}
for idx in range(tbl_hp.num_rows):
u = str(tbl_hp["url"][idx].as_py() or tbl_hp["id"][idx].as_py() or idx)
if u not in url_to_hp:
url_to_hp[u] = idx
hp_urls = sorted(url_to_hp.keys())
rng.shuffle(hp_urls)
hp_calib_docs = []
hp_eval_docs = []
for u in hp_urls:
idx = url_to_hp[u]
q = str(tbl_hp["question"][idx].as_py() or "")
r = str(tbl_hp["reasoning"][idx].as_py() or "")
a = str(tbl_hp["answer"][idx].as_py() or "")
ctx = str(tbl_hp["text"][idx].as_py() or "")[:1200]
doc = f"<|im_start|>user\n以下の文章を参考に質問に答えてください。\n{ctx}\n\n質問: {q}<|im_end|>\n<|im_start|>assistant\n<think>\n{r[:1600]}\n</think>\n{a}<|im_end|>"
sh = get_shingles(doc)
if len(sh & seen_shingles) / max(1, len(sh)) > 0.25:
continue
seen_shingles |= sh
if len(hp_calib_docs) < 180:
hp_calib_docs.append(doc)
elif len(hp_eval_docs) < 45:
hp_eval_docs.append(doc)
# Group ChiTako paragraphs into ~1,600-char (~1,000-token) JA blocks
chitako_grouped = []
cur_g = []
cur_len = 0
for p in dedup_chitako:
cur_g.append(p)
cur_len += len(p)
if cur_len >= 1500:
chitako_grouped.append("\n\n".join(cur_g))
cur_g = []
cur_len = 0
if cur_g:
chitako_grouped.append("\n\n".join(cur_g))
# Interleave round-robin at document granularity across all 5 sources, token-budgeted per round:
# In each round i:
# - 1 Ukisai doc (~3,500 chars EN/bash ≈ 950 tokens)
# - every 4th round: 1 Terminal-Bench xhigh prompt (~4,000 chars ≈ 1,100 tokens)
# - 1 Whittle doc (~3,600 chars EN/code/json/think ≈ 1,000 tokens)
# - 1 JA doc alternating ChiTako (~1,600 chars JA ≈ 1,000 tokens) and Hotchpotch (~1,600 chars JA ≈ 1,000 tokens)
ja_interleaved_docs = []
for i in range(max(len(chitako_grouped), len(hp_calib_docs))):
if i < len(chitako_grouped): ja_interleaved_docs.append(chitako_grouped[i])
if i < len(hp_calib_docs): ja_interleaved_docs.append(hp_calib_docs[i])
n_rounds = max(len(uki_calib_docs), len(whi_calib_docs), len(ja_interleaved_docs))
rounds_text = []
for r in range(n_rounds):
if r < len(uki_calib_docs):
rounds_text.append(uki_calib_docs[r])
if r % 4 == 0 and (r // 4) < len(tb_calib_docs):
rounds_text.append(tb_calib_docs[r // 4])
if r < len(whi_calib_docs):
rounds_text.append(whi_calib_docs[r])
if r < len(ja_interleaved_docs):
rounds_text.append(ja_interleaved_docs[r])
calib_v2 = "\n\n".join(rounds_text)
Path("/content/calibration_mix_v2.txt").write_text(calib_v2, encoding="utf-8")
eval_ops_v2 = "\n\n".join(uki_eval_docs[:25] + tb_eval_docs)[:80_000]
eval_chat_v2 = "\n\n".join(whi_eval_docs[:30])[:80_000]
eval_ja_v2 = "\n\n".join(hp_eval_docs[:35])[:80_000]
Path("/content/eval_ops_tools.txt").write_text(eval_ops_v2, encoding="utf-8")
Path("/content/eval_chat_json.txt").write_text(eval_chat_v2, encoding="utf-8")
Path("/content/eval_ja.txt").write_text(eval_ja_v2, encoding="utf-8")
import hashlib
print(
f"Per-source calib chars: ukisai={sum(len(x) for x in uki_calib_docs):,}, "
f"tb={sum(len(x) for x in tb_calib_docs):,}, whittle={sum(len(x) for x in whi_calib_docs):,}, "
f"chitako={sum(len(x) for x in chitako_grouped):,}, hotchpotch={sum(len(x) for x in hp_calib_docs):,}"
)
print(
f"Per-source eval_ops chars before 80K cap: ukisai_25={sum(len(x) for x in uki_eval_docs[:25]):,}, "
f"tb_8={sum(len(x) for x in tb_eval_docs):,}"
)
for p in ["/content/calibration_mix_v2.txt", "/content/eval_ops_tools.txt", "/content/eval_chat_json.txt", "/content/eval_ja.txt"]:
raw_b = Path(p).read_bytes()
h256 = hashlib.sha256(raw_b).hexdigest()
print(f"Wrote {p:35s}: {len(raw_b) / 1024:8.1f} KB | sha256={h256}")