File size: 7,901 Bytes
53e24ca | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 | """Supervised head training on frozen features; no RL or calibration claims."""
import argparse
import hashlib
import json
import random
from pathlib import Path
import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as optim
from gemma_rlcd import DecisionEngine, Independent, Noul, State
from gemma_rlcd.calibration import evaluate
from gemma_rlcd.core import TokenScores, parse_question, softmax
from gemma_rlcd.decision_head import grouped_cross_entropy
from gemma_rlcd.head_backend import DecisionHeadBackend
def read_rows(path):
rows = [json.loads(line) for line in path.read_text().splitlines() if line.strip()]
if not rows:
raise ValueError(f"Empty dataset: {path}")
for row in rows:
if set(row) != {"id", "state", "questions", "targets"}:
raise ValueError("Each JSONL record requires id, state, questions, and targets")
if set(row["questions"]) != set(row["targets"]):
raise ValueError("Every question requires exactly one target")
for kind in ("images", "audio", "videos"):
if kind in row["state"]:
row["state"][kind] = [
str((path.parent / value).resolve()) for value in row["state"][kind]
]
return rows
def check_splits(splits):
seen_ids, seen_states = set(), set()
for name, rows in splits.items():
for row in rows:
state_hash = hashlib.sha256(
json.dumps(row["state"], sort_keys=True).encode()
).hexdigest()
if row["id"] in seen_ids or state_hash in seen_states:
raise ValueError(f"Duplicate source id or exact state in split {name}")
seen_ids.add(row["id"])
seen_states.add(state_hash)
def compile_row(backend, row):
state = State(**row["state"])
questions = {key: parse_question(value) for key, value in row["questions"].items()}
requests = []
class Capture:
symbols = backend.symbols
def score_batch(self, state, batch):
requests.extend(batch)
return [TokenScores((0.0,) * len(request.symbols), None, 0) for request in batch]
DecisionEngine(Capture()).system_one(state, questions)
target_names = []
for key, question in questions.items():
target = row["targets"][key]
if isinstance(question, Independent):
if not isinstance(target, dict) or set(target) != set(question.criteria):
raise ValueError("Independent labels require one boolean target per label")
for label in question.criteria:
if type(target[label]) is not bool:
raise ValueError("Independent targets must be booleans")
target_names.append("yes" if target[label] else "no")
elif isinstance(question, Noul) and type(target) is bool:
target_names.append("true" if target else "false")
else:
target_names.append(str(target))
targets = tuple(
[name for name, _ in request.criteria].index(target)
for request, target in zip(requests, target_names, strict=True)
)
return state, requests, targets
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--model", required=True)
parser.add_argument("--train", required=True, type=Path)
parser.add_argument("--validation", required=True, type=Path)
parser.add_argument("--test", required=True, type=Path)
parser.add_argument("--output", required=True, type=Path)
parser.add_argument("--steps", type=int, default=600)
parser.add_argument("--learning-rate", type=float, default=3e-4)
parser.add_argument("--seed", type=int, default=11)
parser.add_argument("--state-layers", type=int, default=35)
parser.add_argument("--dtype", choices=["float16", "float32", "bfloat16"], default="float32")
parser.add_argument("--reference", action="store_true")
args = parser.parse_args()
if args.steps < 1 or args.learning_rate <= 0:
parser.error("steps and learning rate must be positive")
if args.output.exists() and any(args.output.iterdir()):
parser.error("output directory must be empty to preserve prior experiments")
splits = {name: read_rows(getattr(args, name)) for name in ("train", "validation", "test")}
check_splits(splits)
mx.random.seed(args.seed)
rng = random.Random(args.seed)
backend = DecisionHeadBackend(
args.model, state_layers=args.state_layers, compute_dtype=args.dtype
)
encoded = {}
compiled = {}
for name, rows in splits.items():
compiled[name] = [compile_row(backend, row) for row in rows]
encoded[name] = [
(backend.encode(state, requests), targets)
for state, requests, targets in compiled[name]
]
print(json.dumps({"encoded_split": name, "source_items": len(rows)}), flush=True)
def measure(split):
rows, targets = [], []
backend.head.eval()
for inputs, expected in encoded[split]:
logits = backend.forward(inputs).tolist()
for start, end in zip(inputs.offsets[:-1], inputs.offsets[1:], strict=True):
rows.append(softmax(logits[start:end]))
targets.extend(expected)
return evaluate(rows, targets)
initial_validation = measure("validation")
optimizer = optim.AdamW(learning_rate=args.learning_rate)
def loss(head, inputs, targets):
logits = head(*inputs.arrays(), len(inputs.offsets) - 1)
return grouped_cross_entropy(logits, inputs.offsets, targets)
value_and_grad = nn.value_and_grad(backend.head, loss)
trace = []
order = list(range(len(encoded["train"])))
for step in range(args.steps):
if step % len(order) == 0:
rng.shuffle(order)
inputs, targets = encoded["train"][order[step % len(order)]]
backend.head.train()
value, gradients = value_and_grad(backend.head, inputs, targets)
optimizer.update(backend.head, gradients)
mx.eval(backend.head.parameters(), optimizer.state, value)
if not bool(mx.isfinite(value).item()):
raise ValueError(f"Non-finite training loss at update {step + 1}")
if (step + 1) % 100 == 0 or step + 1 == args.steps:
record = {"update": step + 1, "loss": float(value.item())}
trace.append(record)
print(json.dumps(record), flush=True)
training = {
"method": "supervised_categorical_nll_frozen_backbone",
"updates": args.steps,
"learning_rate": args.learning_rate,
"seed": args.seed,
"source_counts": {name: len(rows) for name, rows in splits.items()},
"dataset_sha256": {
name: hashlib.sha256(getattr(args, name).read_bytes()).hexdigest() for name in splits
},
}
report = {
"status": "pilot_evidence_only_not_generalization_or_calibration_validation",
"training": training,
"initial_validation": initial_validation,
"train": measure("train"),
"validation": measure("validation"),
"test": measure("test"),
"trace": trace,
}
if args.reference:
from gemma_rlcd.cached_backend import CachedMLXBackend
reference = CachedMLXBackend(args.model)
rows, targets = [], []
for state, requests, expected in compiled["test"]:
rows.extend(softmax(score.logits) for score in reference.score_batch(state, requests))
targets.extend(expected)
report["frozen_decoder_test"] = evaluate(rows, targets)
backend.head.eval()
backend.save(args.output, training)
(args.output / "training-report.json").write_text(json.dumps(report, indent=2) + "\n")
print(json.dumps({key: value for key, value in report.items() if key != "trace"}), flush=True)
if __name__ == "__main__":
main()
|