liqiang
b119135836
Build latest book artifacts / build (push) Canceled after 0s
dependency resolution / resolve (3.11) (push) Canceled after 0s
dependency resolution / resolve (3.13) (push) Canceled after 0s
deploy-pages / build (push) Canceled after 0s
deploy-pages / deploy (push) Canceled after 0s
i18n consistency check / check (push) Canceled after 0s
provider adoption tests / test (chapter2/context-compression) (push) Canceled after 0s
provider adoption tests / test (chapter2/prompt-injection) (push) Canceled after 0s
provider adoption tests / test (chapter2/system-hint) (push) Canceled after 0s
provider adoption tests / test (chapter3/log-sanitization) (push) Canceled after 0s
web-search-agent tests / test (push) Canceled after 0s
web-search-agent tests / agentbook (push) Canceled after 0s
753 lines
32 KiB
Python
753 lines
32 KiB
Python
#!/usr/bin/env python3
|
|
"""Run the manuscript's Experiment 8-8 with canonical, leakage-free evidence.
|
|
|
|
The campaign deliberately uses a small student so that *real parameter training*
|
|
can run on Apple Silicon. It still preserves the experiment's causal contrast:
|
|
|
|
* teacher: long task prompt + a thinking model (Moonshot ``kimi-k3``);
|
|
* student: raw user text only + a non-thinking 135M model;
|
|
* data: disjoint public train/test splits with independent gold labels;
|
|
* evidence: provider response IDs/usage, adapter weights, hashes, held-out
|
|
quality, actual wall-clock latency, tokenizer counts and provider cost.
|
|
|
|
No API key or credential value is ever written to disk.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import os
|
|
import random
|
|
import re
|
|
import statistics
|
|
import sys
|
|
import time
|
|
from dataclasses import dataclass
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from typing import Any, Iterable, Optional
|
|
|
|
|
|
EXPERIMENT_ID = "8-8"
|
|
DATASET_ID = "papluca/language-identification"
|
|
DATASET_REVISION = "aa56583bf2bc52b0565770607d6fc3faebecf9e2"
|
|
STUDENT_MODEL = "HuggingFaceTB/SmolLM2-135M-Instruct"
|
|
STUDENT_REVISION = "12fd25f77366fa6b3b4b768ec3050bf629380bac"
|
|
TEACHER_MODEL = "kimi-k3"
|
|
TEACHER_BASE_URL = "https://api.moonshot.cn/v1"
|
|
TEACHER_KEY_ENV = "MOONSHOT_API_KEY"
|
|
|
|
# Dated native Kimi K3 pricing copied from the verified Chapter 6 campaign
|
|
# configuration. We preserve both the native rate and the dated FX rate.
|
|
PRICING = {
|
|
"input_per_million": 20.0,
|
|
"cached_input_per_million": 2.0,
|
|
"output_per_million": 100.0,
|
|
"currency": "CNY",
|
|
"source_url": "https://platform.kimi.com/docs/pricing/chat-k3.md",
|
|
"as_of": "2026-07-29",
|
|
"usd_per_currency_unit": 0.1477922077922078,
|
|
"fx_source_url": "https://www.ecb.europa.eu/stats/eurofxref/eurofxref-daily.xml",
|
|
"fx_as_of": "2026-07-29",
|
|
}
|
|
|
|
TARGET_LABELS = {"ar", "de", "el", "en", "es", "fr", "hi", "ru", "tr", "ur", "vi", "zh"}
|
|
VALID_LABELS = TARGET_LABELS | {"ot"}
|
|
|
|
LANGUAGE_CLASSIFICATION_PROMPT = """You are a precise language classifier.
|
|
|
|
Goal: Classify the language of the provided text into exactly one label:
|
|
ar (Arabic), de (German), el (Greek), en (English), es (Spanish), fr (French),
|
|
hi (Hindi), ru (Russian), tr (Turkish), ur (Urdu), vi (Vietnamese),
|
|
zh (Chinese), or ot (all other languages / unknown).
|
|
|
|
Rules:
|
|
1. Ignore URLs, email addresses, numbers, emoji, punctuation, and code syntax.
|
|
2. Use native script first. Distinguish Arabic from Urdu by Urdu-specific
|
|
letters and words. Treat both Simplified and Traditional Chinese as zh.
|
|
3. For Latin scripts use vocabulary, function words, and diacritics together;
|
|
do not decide from one proper noun.
|
|
4. Deliberate mixed-language text or a language outside the listed set is ot.
|
|
5. Very short text is ot unless the language signal is unambiguous.
|
|
6. Think through ambiguity before answering, but keep that reasoning private.
|
|
|
|
Text to classify:
|
|
{text}
|
|
|
|
The visible response must be exactly one line: Final Answer: xx
|
|
"""
|
|
|
|
|
|
def utc_now() -> str:
|
|
return datetime.now(timezone.utc).isoformat()
|
|
|
|
|
|
def sha256_file(path: Path) -> str:
|
|
h = hashlib.sha256()
|
|
with path.open("rb") as f:
|
|
for chunk in iter(lambda: f.read(1024 * 1024), b""):
|
|
h.update(chunk)
|
|
return h.hexdigest()
|
|
|
|
|
|
def canonical_hash(value: Any) -> str:
|
|
payload = json.dumps(value, sort_keys=True, ensure_ascii=False, separators=(",", ":"))
|
|
return hashlib.sha256(payload.encode("utf-8")).hexdigest()
|
|
|
|
|
|
def write_json(path: Path, value: Any) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_text(json.dumps(value, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
|
|
|
|
|
|
def write_jsonl(path: Path, rows: Iterable[dict[str, Any]]) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
with path.open("w", encoding="utf-8") as f:
|
|
for row in rows:
|
|
f.write(json.dumps(row, ensure_ascii=False) + "\n")
|
|
|
|
|
|
def read_jsonl(path: Path) -> list[dict[str, Any]]:
|
|
if not path.exists():
|
|
return []
|
|
with path.open(encoding="utf-8") as f:
|
|
return [json.loads(line) for line in f if line.strip()]
|
|
|
|
|
|
def map_gold(source_label: str) -> str:
|
|
return source_label if source_label in TARGET_LABELS else "ot"
|
|
|
|
|
|
def _stable_sample(rows: list[dict[str, str]], n: int, seed: int, split: str) -> list[dict[str, str]]:
|
|
grouped: dict[str, list[dict[str, str]]] = {}
|
|
for row in rows:
|
|
grouped.setdefault(row["labels"], []).append(row)
|
|
selected: list[dict[str, str]] = []
|
|
for source_label in sorted(grouped):
|
|
candidates = sorted(
|
|
grouped[source_label],
|
|
key=lambda r: hashlib.sha256(f"{seed}:{split}:{source_label}:{r['text']}".encode()).hexdigest(),
|
|
)
|
|
for rank, row in enumerate(candidates[:n]):
|
|
selected.append(
|
|
{
|
|
"id": f"{split}-{source_label}-{rank:03d}",
|
|
"split": split,
|
|
"source_label": source_label,
|
|
"gold_label": map_gold(source_label),
|
|
"text": row["text"],
|
|
}
|
|
)
|
|
return sorted(selected, key=lambda r: r["id"])
|
|
|
|
|
|
def prepare_benchmark(run_dir: Path, train_per_language: int, test_per_language: int, seed: int) -> dict[str, Any]:
|
|
from datasets import load_dataset
|
|
|
|
train_ds = load_dataset(DATASET_ID, revision=DATASET_REVISION, split="train")
|
|
test_ds = load_dataset(DATASET_ID, revision=DATASET_REVISION, split="test")
|
|
train = _stable_sample(list(train_ds), train_per_language, seed, "train")
|
|
test = _stable_sample(list(test_ds), test_per_language, seed, "test")
|
|
train_texts = {r["text"] for r in train}
|
|
test_texts = {r["text"] for r in test}
|
|
overlap = train_texts & test_texts
|
|
if overlap:
|
|
raise RuntimeError(f"train/test leakage detected: {len(overlap)} texts")
|
|
write_jsonl(run_dir / "benchmark_train_gold.jsonl", train)
|
|
write_jsonl(run_dir / "benchmark_test_gold.jsonl", test)
|
|
provenance = {
|
|
"dataset": DATASET_ID,
|
|
"revision": DATASET_REVISION,
|
|
"train_source_split": "train",
|
|
"test_source_split": "test",
|
|
"train_fingerprint": train_ds._fingerprint,
|
|
"test_fingerprint": test_ds._fingerprint,
|
|
"seed": seed,
|
|
"train_per_source_language": train_per_language,
|
|
"test_per_source_language": test_per_language,
|
|
"train_rows": len(train),
|
|
"test_rows": len(test),
|
|
"train_unique_texts": len(train_texts),
|
|
"test_unique_texts": len(test_texts),
|
|
"exact_text_overlap": len(overlap),
|
|
"train_rows_sha256": canonical_hash(train),
|
|
"test_rows_sha256": canonical_hash(test),
|
|
}
|
|
write_json(run_dir / "dataset_provenance.json", provenance)
|
|
return provenance
|
|
|
|
|
|
def parse_label(text: Optional[str]) -> Optional[str]:
|
|
if not text:
|
|
return None
|
|
matches = re.findall(r"Final Answer:\s*([a-z]{2})", text, re.IGNORECASE)
|
|
if matches and matches[-1].lower() in VALID_LABELS:
|
|
return matches[-1].lower()
|
|
stripped = text.strip().lower()
|
|
return stripped if stripped in VALID_LABELS else None
|
|
|
|
|
|
def usage_cost(usage: dict[str, Any]) -> dict[str, float]:
|
|
prompt = int(usage.get("prompt_tokens") or 0)
|
|
completion = int(usage.get("completion_tokens") or 0)
|
|
details = usage.get("prompt_tokens_details") or {}
|
|
cached = int(details.get("cached_tokens") or usage.get("cached_tokens") or 0)
|
|
uncached = max(0, prompt - cached)
|
|
cny = (
|
|
uncached * PRICING["input_per_million"]
|
|
+ cached * PRICING["cached_input_per_million"]
|
|
+ completion * PRICING["output_per_million"]
|
|
) / 1_000_000
|
|
return {"cny": cny, "usd": cny * PRICING["usd_per_currency_unit"]}
|
|
|
|
|
|
async def collect_teacher(run_dir: Path, concurrency: int, max_retries: int) -> dict[str, Any]:
|
|
from openai import AsyncOpenAI
|
|
|
|
key = os.getenv(TEACHER_KEY_ENV)
|
|
if not key:
|
|
raise RuntimeError(f"{TEACHER_KEY_ENV} is not configured")
|
|
train = read_jsonl(run_dir / "benchmark_train_gold.jsonl")
|
|
test = read_jsonl(run_dir / "benchmark_test_gold.jsonl")
|
|
rows = train + test
|
|
raw_path = run_dir / "teacher_receipts.jsonl"
|
|
existing = {r["id"]: r for r in read_jsonl(raw_path)}
|
|
client = AsyncOpenAI(api_key=key, base_url=TEACHER_BASE_URL, timeout=90)
|
|
sem = asyncio.Semaphore(concurrency)
|
|
lock = asyncio.Lock()
|
|
|
|
async def one(row: dict[str, Any]) -> dict[str, Any]:
|
|
async with sem:
|
|
error: Optional[str] = None
|
|
for attempt in range(max_retries + 1):
|
|
started = utc_now()
|
|
t0 = time.perf_counter()
|
|
try:
|
|
response = await client.chat.completions.create(
|
|
model=TEACHER_MODEL,
|
|
messages=[
|
|
{
|
|
"role": "system",
|
|
"content": "Reason privately. The visible response must be exactly: Final Answer: xx",
|
|
},
|
|
{"role": "user", "content": LANGUAGE_CLASSIFICATION_PROMPT.format(text=row["text"])},
|
|
],
|
|
max_tokens=256,
|
|
extra_body={"reasoning_effort": "low"},
|
|
)
|
|
latency = time.perf_counter() - t0
|
|
message = response.choices[0].message
|
|
content = message.content or ""
|
|
reasoning = getattr(message, "reasoning_content", None) or getattr(message, "reasoning", None)
|
|
usage = response.usage.model_dump() if response.usage else {}
|
|
record = {
|
|
**row,
|
|
"provider": "moonshot",
|
|
"base_url": TEACHER_BASE_URL,
|
|
"model_requested": TEACHER_MODEL,
|
|
"response_id": response.id,
|
|
"response_model": response.model,
|
|
"response_created": response.created,
|
|
"request_started_at": started,
|
|
"latency_seconds": latency,
|
|
"attempt": attempt,
|
|
"content": content,
|
|
"reasoning_content": reasoning,
|
|
"prediction": parse_label(content),
|
|
"usage": usage,
|
|
"calculated_cost": usage_cost(usage),
|
|
"request": {
|
|
"max_tokens": 256,
|
|
"reasoning_effort": "low",
|
|
"prompt_sha256": hashlib.sha256(
|
|
LANGUAGE_CLASSIFICATION_PROMPT.format(text=row["text"]).encode("utf-8")
|
|
).hexdigest(),
|
|
},
|
|
"error": None,
|
|
}
|
|
async with lock:
|
|
with raw_path.open("a", encoding="utf-8") as f:
|
|
f.write(json.dumps(record, ensure_ascii=False) + "\n")
|
|
return record
|
|
except Exception as exc: # noqa: BLE001
|
|
error = f"{type(exc).__name__}: {exc}"
|
|
record = {**row, "provider": "moonshot", "model_requested": TEACHER_MODEL, "error": error}
|
|
async with lock:
|
|
with raw_path.open("a", encoding="utf-8") as f:
|
|
f.write(json.dumps(record, ensure_ascii=False) + "\n")
|
|
return record
|
|
|
|
pending = [row for row in rows if row["id"] not in existing or existing[row["id"]].get("error")]
|
|
if pending:
|
|
results = await asyncio.gather(*(one(row) for row in pending))
|
|
existing.update({r["id"]: r for r in results})
|
|
ordered = [existing[row["id"]] for row in rows]
|
|
write_jsonl(raw_path, ordered)
|
|
|
|
accepted_train = [
|
|
{
|
|
"id": r["id"],
|
|
"messages": [
|
|
{"role": "user", "content": r["text"]},
|
|
{"role": "assistant", "content": r["prediction"]},
|
|
],
|
|
"teacher_response_id": r.get("response_id"),
|
|
"gold_label": r["gold_label"],
|
|
}
|
|
for r in ordered
|
|
if r["split"] == "train" and not r.get("error") and r.get("prediction") == r["gold_label"]
|
|
]
|
|
write_jsonl(run_dir / "student_train.jsonl", accepted_train)
|
|
|
|
def split_summary(split: str) -> dict[str, Any]:
|
|
rs = [r for r in ordered if r["split"] == split]
|
|
valid = [r for r in rs if not r.get("error") and r.get("prediction")]
|
|
costs = [r.get("calculated_cost") or {} for r in valid]
|
|
return {
|
|
"rows": len(rs),
|
|
"valid_receipts": len(valid),
|
|
"unique_response_ids": len({r.get("response_id") for r in valid}),
|
|
"correct": sum(r.get("prediction") == r["gold_label"] for r in valid),
|
|
"gold_accuracy": sum(r.get("prediction") == r["gold_label"] for r in valid) / len(rs),
|
|
"prompt_tokens": sum((r.get("usage") or {}).get("prompt_tokens", 0) for r in valid),
|
|
"completion_tokens": sum((r.get("usage") or {}).get("completion_tokens", 0) for r in valid),
|
|
"latency_seconds_total": sum(r.get("latency_seconds", 0) for r in valid),
|
|
"latency_seconds_mean": statistics.mean(r.get("latency_seconds", 0) for r in valid) if valid else None,
|
|
"provider_cost_cny": sum(c.get("cny", 0) for c in costs),
|
|
"provider_cost_usd": sum(c.get("usd", 0) for c in costs),
|
|
}
|
|
|
|
summary = {
|
|
"teacher": {"provider": "moonshot", "model": TEACHER_MODEL, "pricing": PRICING},
|
|
"train": split_summary("train"),
|
|
"test": split_summary("test"),
|
|
"accepted_student_train_rows": len(accepted_train),
|
|
}
|
|
write_json(run_dir / "teacher_summary.json", summary)
|
|
return summary
|
|
|
|
|
|
@dataclass
|
|
class EncodedRow:
|
|
input_ids: list[int]
|
|
labels: list[int]
|
|
|
|
|
|
def _chat_template_ids(encoded: Any) -> list[int]:
|
|
"""Return one token-id list across Transformers 4.x and 5.x APIs."""
|
|
if isinstance(encoded, dict) or hasattr(encoded, "keys"):
|
|
encoded = encoded["input_ids"]
|
|
if hasattr(encoded, "tolist"):
|
|
encoded = encoded.tolist()
|
|
if encoded and isinstance(encoded[0], list):
|
|
if len(encoded) != 1:
|
|
raise ValueError("expected one chat-template sequence")
|
|
encoded = encoded[0]
|
|
if not isinstance(encoded, list) or not all(isinstance(token, int) for token in encoded):
|
|
raise TypeError("chat template did not return a one-dimensional integer token sequence")
|
|
return encoded
|
|
|
|
|
|
def _local_device(torch: Any) -> Any:
|
|
if torch.cuda.is_available():
|
|
return torch.device("cuda")
|
|
if torch.backends.mps.is_available():
|
|
return torch.device("mps")
|
|
return torch.device("cpu")
|
|
|
|
|
|
def _encode_sft(tokenizer: Any, messages: list[dict[str, str]], max_length: int) -> EncodedRow:
|
|
prompt = messages[:-1]
|
|
prompt_ids = _chat_template_ids(
|
|
tokenizer.apply_chat_template(prompt, tokenize=True, add_generation_prompt=True)
|
|
)
|
|
full_ids = _chat_template_ids(
|
|
tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=False)
|
|
)
|
|
full_ids = full_ids[:max_length]
|
|
labels = [-100] * min(len(prompt_ids), len(full_ids)) + full_ids[min(len(prompt_ids), len(full_ids)) :]
|
|
return EncodedRow(input_ids=full_ids, labels=labels)
|
|
|
|
|
|
def train_student(run_dir: Path, epochs: int, batch_size: int, learning_rate: float, seed: int) -> dict[str, Any]:
|
|
import torch
|
|
from peft import LoraConfig, get_peft_model
|
|
from torch.utils.data import DataLoader
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
|
|
random.seed(seed)
|
|
torch.manual_seed(seed)
|
|
device = _local_device(torch)
|
|
tokenizer = AutoTokenizer.from_pretrained(STUDENT_MODEL, revision=STUDENT_REVISION)
|
|
if tokenizer.pad_token_id is None:
|
|
tokenizer.pad_token = tokenizer.eos_token
|
|
rows = read_jsonl(run_dir / "student_train.jsonl")
|
|
if not rows:
|
|
# A teacher campaign may be interrupted after individual receipts have
|
|
# been durably appended but before its final derived files are written.
|
|
# Rebuild the training view from successful, gold-verified real
|
|
# receipts so the documented phase-resume boundary is actually usable.
|
|
receipts = read_jsonl(run_dir / "teacher_receipts.jsonl")
|
|
rows = [
|
|
{
|
|
"id": receipt["id"],
|
|
"messages": [
|
|
{"role": "user", "content": receipt["text"]},
|
|
{"role": "assistant", "content": receipt["prediction"]},
|
|
],
|
|
"teacher_response_id": receipt.get("response_id"),
|
|
"gold_label": receipt["gold_label"],
|
|
}
|
|
for receipt in receipts
|
|
if receipt.get("split") == "train"
|
|
and not receipt.get("error")
|
|
and receipt.get("response_id")
|
|
and receipt.get("prediction") == receipt.get("gold_label")
|
|
]
|
|
if rows:
|
|
write_jsonl(run_dir / "student_train.jsonl", rows)
|
|
if not rows:
|
|
raise RuntimeError("no accepted teacher rows available for training")
|
|
encoded = [_encode_sft(tokenizer, r["messages"], 192) for r in rows]
|
|
|
|
def collate(items: list[EncodedRow]) -> dict[str, torch.Tensor]:
|
|
length = max(len(x.input_ids) for x in items)
|
|
input_ids, labels, masks = [], [], []
|
|
for item in items:
|
|
pad = length - len(item.input_ids)
|
|
input_ids.append(item.input_ids + [tokenizer.pad_token_id] * pad)
|
|
labels.append(item.labels + [-100] * pad)
|
|
masks.append([1] * len(item.input_ids) + [0] * pad)
|
|
return {
|
|
"input_ids": torch.tensor(input_ids, dtype=torch.long),
|
|
"labels": torch.tensor(labels, dtype=torch.long),
|
|
"attention_mask": torch.tensor(masks, dtype=torch.long),
|
|
}
|
|
|
|
generator = torch.Generator().manual_seed(seed)
|
|
loader = DataLoader(encoded, batch_size=batch_size, shuffle=True, collate_fn=collate, generator=generator)
|
|
base = AutoModelForCausalLM.from_pretrained(
|
|
STUDENT_MODEL,
|
|
revision=STUDENT_REVISION,
|
|
dtype=torch.bfloat16 if device.type == "cuda" else torch.float32,
|
|
)
|
|
config = LoraConfig(
|
|
r=16,
|
|
lora_alpha=32,
|
|
lora_dropout=0.0,
|
|
target_modules=["q_proj", "v_proj"],
|
|
bias="none",
|
|
task_type="CAUSAL_LM",
|
|
)
|
|
model = get_peft_model(base, config).to(device)
|
|
model.train()
|
|
params = [p for p in model.parameters() if p.requires_grad]
|
|
optimizer = torch.optim.AdamW(params, lr=learning_rate)
|
|
losses: list[float] = []
|
|
started = utc_now()
|
|
t0 = time.perf_counter()
|
|
for epoch in range(epochs):
|
|
for batch in loader:
|
|
optimizer.zero_grad(set_to_none=True)
|
|
batch = {k: v.to(device) for k, v in batch.items()}
|
|
output = model(**batch)
|
|
output.loss.backward()
|
|
torch.nn.utils.clip_grad_norm_(params, 1.0)
|
|
optimizer.step()
|
|
losses.append(float(output.loss.detach().cpu()))
|
|
if device.type == "mps":
|
|
torch.mps.synchronize()
|
|
runtime = time.perf_counter() - t0
|
|
adapter_dir = run_dir / "student_adapter"
|
|
adapter_dir.mkdir(parents=True, exist_ok=True)
|
|
model.save_pretrained(adapter_dir, safe_serialization=True)
|
|
tokenizer.save_pretrained(adapter_dir)
|
|
receipt = {
|
|
"started_at": started,
|
|
"finished_at": utc_now(),
|
|
"runtime_seconds": runtime,
|
|
"device": str(device),
|
|
"torch_version": torch.__version__,
|
|
"base_model": STUDENT_MODEL,
|
|
"base_model_revision": STUDENT_REVISION,
|
|
"training_rows": len(rows),
|
|
"epochs": epochs,
|
|
"batch_size": batch_size,
|
|
"optimizer": "AdamW",
|
|
"learning_rate": learning_rate,
|
|
"lora": {"r": 16, "alpha": 32, "targets": ["q_proj", "v_proj"]},
|
|
"trainable_parameters": sum(p.numel() for p in params),
|
|
"total_parameters_with_adapter": sum(p.numel() for p in model.parameters()),
|
|
"optimizer_steps": len(losses),
|
|
"first_loss": losses[0],
|
|
"final_loss": losses[-1],
|
|
"mean_loss": statistics.mean(losses),
|
|
"losses": losses,
|
|
}
|
|
write_json(run_dir / "training_receipt.json", receipt)
|
|
del model, base
|
|
if device.type == "mps":
|
|
torch.mps.empty_cache()
|
|
elif device.type == "cuda":
|
|
torch.cuda.empty_cache()
|
|
return receipt
|
|
|
|
|
|
def _sync(device: Any) -> None:
|
|
import torch
|
|
|
|
if device.type == "mps":
|
|
torch.mps.synchronize()
|
|
elif device.type == "cuda":
|
|
torch.cuda.synchronize(device)
|
|
|
|
|
|
def summarize_retained_teacher_receipts(run_dir: Path) -> dict[str, Any]:
|
|
"""Summarize durable teacher receipts, including interrupted campaigns."""
|
|
expected = read_jsonl(run_dir / "benchmark_train_gold.jsonl") + read_jsonl(
|
|
run_dir / "benchmark_test_gold.jsonl"
|
|
)
|
|
expected_by_split = {
|
|
split: [row for row in expected if row["split"] == split]
|
|
for split in ("train", "test")
|
|
}
|
|
receipts = read_jsonl(run_dir / "teacher_receipts.jsonl")
|
|
|
|
def split_summary(split: str) -> dict[str, Any]:
|
|
expected_rows = expected_by_split[split]
|
|
rs = [r for r in receipts if r.get("split") == split]
|
|
valid = [r for r in rs if not r.get("error") and r.get("prediction")]
|
|
costs = [r.get("calculated_cost") or {} for r in valid]
|
|
correct = sum(r.get("prediction") == r.get("gold_label") for r in valid)
|
|
return {
|
|
"rows": len(expected_rows),
|
|
"retained_receipts": len(rs),
|
|
"coverage": len(valid) / len(expected_rows) if expected_rows else 0.0,
|
|
"valid_receipts": len(valid),
|
|
"unique_response_ids": len({r.get("response_id") for r in valid}),
|
|
"correct": correct,
|
|
# Missing calls do not silently disappear from the campaign score.
|
|
"gold_accuracy": correct / len(expected_rows) if expected_rows else 0.0,
|
|
"observed_gold_accuracy": correct / len(valid) if valid else 0.0,
|
|
"prompt_tokens": sum((r.get("usage") or {}).get("prompt_tokens", 0) for r in valid),
|
|
"completion_tokens": sum((r.get("usage") or {}).get("completion_tokens", 0) for r in valid),
|
|
"latency_seconds_total": sum(r.get("latency_seconds", 0) for r in valid),
|
|
"latency_seconds_mean": statistics.mean(r.get("latency_seconds", 0) for r in valid) if valid else None,
|
|
"provider_cost_cny": sum(c.get("cny", 0) for c in costs),
|
|
"provider_cost_usd": sum(c.get("usd", 0) for c in costs),
|
|
}
|
|
|
|
summary = {
|
|
"teacher": {"provider": "moonshot", "model": TEACHER_MODEL, "pricing": PRICING},
|
|
"train": split_summary("train"),
|
|
"test": split_summary("test"),
|
|
"campaign_complete": len(receipts) == len(expected)
|
|
and all(not r.get("error") and r.get("response_id") for r in receipts),
|
|
}
|
|
write_json(run_dir / "teacher_summary.json", summary)
|
|
return summary
|
|
|
|
|
|
def evaluate_local_arm(run_dir: Path, arm: str) -> dict[str, Any]:
|
|
import torch
|
|
from peft import PeftModel
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
|
|
device = _local_device(torch)
|
|
tokenizer = AutoTokenizer.from_pretrained(STUDENT_MODEL, revision=STUDENT_REVISION)
|
|
if tokenizer.pad_token_id is None:
|
|
tokenizer.pad_token = tokenizer.eos_token
|
|
base = AutoModelForCausalLM.from_pretrained(
|
|
STUDENT_MODEL,
|
|
revision=STUDENT_REVISION,
|
|
dtype=torch.bfloat16 if device.type == "cuda" else torch.float32,
|
|
).to(device)
|
|
model = base if arm == "baseline" else PeftModel.from_pretrained(base, run_dir / "student_adapter").to(device)
|
|
model.eval()
|
|
rows = read_jsonl(run_dir / "benchmark_test_gold.jsonl")
|
|
|
|
# Warmup is measured separately and excluded from reported case latency.
|
|
warm = tokenizer.apply_chat_template(
|
|
[{"role": "user", "content": "Hello world."}], tokenize=True, add_generation_prompt=True, return_tensors="pt"
|
|
)
|
|
if isinstance(warm, dict) or hasattr(warm, "keys"):
|
|
warm = warm["input_ids"]
|
|
warm = warm.to(device)
|
|
with torch.no_grad():
|
|
_ = model.generate(warm, max_new_tokens=4, do_sample=False, pad_token_id=tokenizer.eos_token_id)
|
|
_sync(device)
|
|
|
|
results: list[dict[str, Any]] = []
|
|
for row in rows:
|
|
ids = tokenizer.apply_chat_template(
|
|
[{"role": "user", "content": row["text"]}],
|
|
tokenize=True,
|
|
add_generation_prompt=True,
|
|
return_tensors="pt",
|
|
)
|
|
if isinstance(ids, dict) or hasattr(ids, "keys"):
|
|
ids = ids["input_ids"]
|
|
ids = ids.to(device)
|
|
_sync(device)
|
|
t0 = time.perf_counter()
|
|
with torch.no_grad():
|
|
out = model.generate(
|
|
ids,
|
|
max_new_tokens=8,
|
|
do_sample=False,
|
|
pad_token_id=tokenizer.eos_token_id,
|
|
)
|
|
_sync(device)
|
|
latency = time.perf_counter() - t0
|
|
generated = out[0, ids.shape[1] :]
|
|
text = tokenizer.decode(generated, skip_special_tokens=True).strip()
|
|
pred = parse_label(text) or (text.lower() if text.lower() in VALID_LABELS else None)
|
|
results.append(
|
|
{
|
|
"id": row["id"],
|
|
"gold_label": row["gold_label"],
|
|
"source_label": row["source_label"],
|
|
"prediction": pred,
|
|
"response": text,
|
|
"correct": pred == row["gold_label"],
|
|
"input_tokens": int(ids.shape[1]),
|
|
"output_tokens": int(generated.shape[0]),
|
|
"latency_seconds": latency,
|
|
}
|
|
)
|
|
output_path = run_dir / f"student_{arm}_heldout.jsonl"
|
|
write_jsonl(output_path, results)
|
|
summary = {
|
|
"arm": arm,
|
|
"rows": len(results),
|
|
"correct": sum(r["correct"] for r in results),
|
|
"accuracy": sum(r["correct"] for r in results) / len(results),
|
|
"parse_rate": sum(r["prediction"] is not None for r in results) / len(results),
|
|
"input_tokens": sum(r["input_tokens"] for r in results),
|
|
"output_tokens": sum(r["output_tokens"] for r in results),
|
|
"latency_seconds_total": sum(r["latency_seconds"] for r in results),
|
|
"latency_seconds_mean": statistics.mean(r["latency_seconds"] for r in results),
|
|
"latency_seconds_median": statistics.median(r["latency_seconds"] for r in results),
|
|
"provider_charge_usd": 0.0,
|
|
"provider_charge_note": "Local inference made no provider API calls; electricity/hardware amortization is excluded.",
|
|
}
|
|
write_json(run_dir / f"student_{arm}_summary.json", summary)
|
|
del model, base
|
|
if device.type == "mps":
|
|
torch.mps.empty_cache()
|
|
elif device.type == "cuda":
|
|
torch.cuda.empty_cache()
|
|
return summary
|
|
|
|
|
|
def package_evidence(run_dir: Path, command: str) -> dict[str, Any]:
|
|
teacher = summarize_retained_teacher_receipts(run_dir)
|
|
baseline = json.loads((run_dir / "student_baseline_summary.json").read_text(encoding="utf-8"))
|
|
trained = json.loads((run_dir / "student_trained_summary.json").read_text(encoding="utf-8"))
|
|
train_rows = read_jsonl(run_dir / "benchmark_train_gold.jsonl")
|
|
test_rows = read_jsonl(run_dir / "benchmark_test_gold.jsonl")
|
|
train_ids = {r["text"] for r in train_rows}
|
|
test_ids = {r["text"] for r in test_rows}
|
|
teacher_test = [r for r in read_jsonl(run_dir / "teacher_receipts.jsonl") if r["split"] == "test"]
|
|
trained_rows = {r["id"]: r for r in read_jsonl(run_dir / "student_trained_heldout.jsonl")}
|
|
agreement = sum(
|
|
trained_rows[r["id"]]["prediction"] == r.get("prediction")
|
|
for r in teacher_test
|
|
if r["id"] in trained_rows and r.get("prediction")
|
|
) / len(teacher_test)
|
|
comparison = {
|
|
"heldout_rows": len(test_rows),
|
|
"train_test_exact_text_overlap": len(train_ids & test_ids),
|
|
"teacher_gold_accuracy": teacher["test"]["gold_accuracy"],
|
|
"baseline_gold_accuracy": baseline["accuracy"],
|
|
"trained_gold_accuracy": trained["accuracy"],
|
|
"trained_absolute_uplift": trained["accuracy"] - baseline["accuracy"],
|
|
"trained_teacher_agreement": agreement,
|
|
"teacher_mean_latency_seconds": teacher["test"]["latency_seconds_mean"],
|
|
"trained_mean_latency_seconds": trained["latency_seconds_mean"],
|
|
"latency_speedup": teacher["test"]["latency_seconds_mean"] / trained["latency_seconds_mean"],
|
|
"teacher_input_tokens": teacher["test"]["prompt_tokens"],
|
|
"trained_input_tokens": trained["input_tokens"],
|
|
"input_token_reduction": 1 - trained["input_tokens"] / teacher["test"]["prompt_tokens"],
|
|
"teacher_provider_cost_usd": teacher["test"]["provider_cost_usd"],
|
|
"trained_provider_charge_usd": trained["provider_charge_usd"],
|
|
"cost_scope_note": trained["provider_charge_note"],
|
|
}
|
|
write_json(run_dir / "comparison.json", comparison)
|
|
required_files = [
|
|
p
|
|
for p in run_dir.rglob("*")
|
|
if p.is_file() and p.name != "manifest.json"
|
|
]
|
|
artifact_hashes = {
|
|
str(p.relative_to(run_dir)): {"sha256": sha256_file(p), "bytes": p.stat().st_size}
|
|
for p in sorted(required_files)
|
|
}
|
|
receipts = read_jsonl(run_dir / "teacher_receipts.jsonl")
|
|
gates = {
|
|
"disjoint_heldout": len(train_ids & test_ids) == 0,
|
|
"real_teacher_receipts": len(receipts) == len(train_rows) + len(test_rows)
|
|
and len({r.get("response_id") for r in receipts}) == len(receipts),
|
|
"real_parameter_training": (run_dir / "student_adapter" / "adapter_model.safetensors").exists(),
|
|
"before_after_quality": trained["accuracy"] > baseline["accuracy"],
|
|
"quality_near_teacher": trained["accuracy"] >= teacher["test"]["gold_accuracy"] - 0.15,
|
|
"measured_latency": teacher["test"]["latency_seconds_mean"] > 0 and trained["latency_seconds_mean"] > 0,
|
|
"measured_tokens": teacher["test"]["prompt_tokens"] > 0 and trained["input_tokens"] > 0,
|
|
"dollar_accounting": teacher["test"]["provider_cost_usd"] > 0,
|
|
}
|
|
manifest = {
|
|
"schema_version": 1,
|
|
"experiment_id": EXPERIMENT_ID,
|
|
"status": "complete" if all(gates.values()) else "incomplete",
|
|
"created_at": utc_now(),
|
|
"command": command,
|
|
"host": {"platform": sys.platform, "python": sys.version, "credential_env": TEACHER_KEY_ENV},
|
|
"student": {"model": STUDENT_MODEL, "revision": STUDENT_REVISION},
|
|
"teacher": {"model": TEACHER_MODEL, "base_url": TEACHER_BASE_URL, "pricing": PRICING},
|
|
"gates": gates,
|
|
"comparison": comparison,
|
|
"artifacts": artifact_hashes,
|
|
"credential_values_retained": False,
|
|
}
|
|
write_json(run_dir / "manifest.json", manifest)
|
|
return manifest
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Run canonical Experiment 8-8")
|
|
parser.add_argument("--run-dir", default="./validation/exp8-8-real")
|
|
parser.add_argument("--phase", choices=["all", "prepare", "teacher", "train", "evaluate", "package"], default="all")
|
|
parser.add_argument("--train-per-language", type=int, default=8)
|
|
parser.add_argument("--test-per-language", type=int, default=4)
|
|
parser.add_argument("--seed", type=int, default=78)
|
|
parser.add_argument("--concurrency", type=int, default=12)
|
|
parser.add_argument("--max-retries", type=int, default=2)
|
|
parser.add_argument("--epochs", type=int, default=5)
|
|
parser.add_argument("--batch-size", type=int, default=4)
|
|
parser.add_argument("--learning-rate", type=float, default=8e-4)
|
|
args = parser.parse_args()
|
|
run_dir = Path(args.run_dir).resolve()
|
|
run_dir.mkdir(parents=True, exist_ok=True)
|
|
phases = {args.phase} if args.phase != "all" else {"prepare", "teacher", "train", "evaluate", "package"}
|
|
if "prepare" in phases:
|
|
prepare_benchmark(run_dir, args.train_per_language, args.test_per_language, args.seed)
|
|
if "teacher" in phases:
|
|
asyncio.run(collect_teacher(run_dir, args.concurrency, args.max_retries))
|
|
if "train" in phases:
|
|
train_student(run_dir, args.epochs, args.batch_size, args.learning_rate, args.seed)
|
|
if "evaluate" in phases:
|
|
evaluate_local_arm(run_dir, "baseline")
|
|
evaluate_local_arm(run_dir, "trained")
|
|
if "package" in phases:
|
|
command = " ".join([sys.executable, *sys.argv])
|
|
manifest = package_evidence(run_dir, command)
|
|
print(json.dumps({"run_dir": str(run_dir), "status": manifest["status"], "comparison": manifest["comparison"]}, indent=2))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|