Create classifier_trainer.py

classifier_trainer.py

@@ -0,0 +1,476 @@
+# =============================================================================
+# CELL 3: Train Geometric Classifier + Upload to HuggingFace
+# Requires: Cell 1 (generator/constants), Cell 2 (model classes)
+# Outputs: `model` in notebook scope + geometric_classifier/ on HF
+#
+# Features:
+#   - Dataset cached to disk (skip regeneration on resume)
+#   - Checkpoint saved every epoch (model, optimizer, scheduler, epoch, best_acc)
+#   - Auto-resume from latest checkpoint
+# =============================================================================
+
+import json, time, os, shutil
+from pathlib import Path
+
+HF_REPO = "AbstractPhil/grid-geometric-classifier-proto"
+CKPT_DIR = Path("./checkpoints")
+DATASET_PATH = Path("./cached_dataset.pt")
+
+# --- Loss Functions ---
+
+def _safe_bce(inp, tgt):
+    """BCE that forces fp32 and clamps to prevent log(0) from BF16 sigmoid saturation."""
+    with torch.amp.autocast('cuda', enabled=False):
+        return F.binary_cross_entropy(
+            inp.float().clamp(1e-7, 1 - 1e-7),
+            tgt.float())
+
+def capacity_fill_loss(fr, dt): return _safe_bce(fr, dt)
+def overflow_reg(on, dt):
+    """Vectorized overflow penalty — no Python loops, no .item() calls."""
+    pk = dt.sum(dim=-1).long().clamp(min=0)  # (B,) peak dim index
+    n_caps = on.shape[1]
+    arange = torch.arange(n_caps, device=on.device).unsqueeze(0)  # (1, n_caps)
+    mask = (arange >= pk.unsqueeze(1)).float()  # (B, n_caps)
+    return (on * mask).sum() / (on.shape[0] + 1e-8)
+def cap_diversity(c): return -c.var()
+def peak_loss(l, t): return F.cross_entropy(l, t)
+def cm_loss(p, t): return F.mse_loss(p, torch.sign(t))
+def curved_bce(p, t): return _safe_bce(p.squeeze(-1), t)
+def ctype_loss(l, t): return F.cross_entropy(l, t)
+
+# --- Dataset Cache ---
+
+def get_or_generate_dataset(n_samples, seed, path=DATASET_PATH):
+    """Load cached dataset from disk, or generate + cache it."""
+    if path.exists():
+        print(f"Loading cached dataset from {path}...")
+        t0 = time.time()
+        cached = torch.load(path, weights_only=True)
+        if cached["n_samples"] == n_samples and cached["seed"] == seed:
+            train_ds = ShapeDataset.__new__(ShapeDataset)
+            val_ds = ShapeDataset.__new__(ShapeDataset)
+            for k in ["grids", "labels", "dim_conf", "peak_dim", "volume", "cm_det", "is_curved", "curvature"]:
+                setattr(train_ds, k, cached["train"][k])
+                setattr(val_ds, k, cached["val"][k])
+            dt = time.time() - t0
+            print(f"Loaded {len(train_ds)} train + {len(val_ds)} val in {dt:.1f}s (cached)")
+            return train_ds, val_ds
+        else:
+            print(f"Cache mismatch (n={cached['n_samples']}, seed={cached['seed']}) — regenerating")
+
+    all_samples = generate_parallel(n_samples, seed=seed, n_workers=8)
+    n_train = int(len(all_samples) * 0.8)
+    train_ds = ShapeDataset(all_samples[:n_train])
+    val_ds = ShapeDataset(all_samples[n_train:])
+
+    print(f"Caching dataset to {path}...")
+    cache_data = {
+        "n_samples": n_samples, "seed": seed,
+        "train": {k: getattr(train_ds, k) for k in ["grids", "labels", "dim_conf", "peak_dim", "volume", "cm_det", "is_curved", "curvature"]},
+        "val": {k: getattr(val_ds, k) for k in ["grids", "labels", "dim_conf", "peak_dim", "volume", "cm_det", "is_curved", "curvature"]},
+    }
+    torch.save(cache_data, path)
+    size_mb = path.stat().st_size / 1e6
+    print(f"Cached: {size_mb:.0f}MB")
+    return train_ds, val_ds
+
+# --- Checkpoint helpers ---
+
+def save_checkpoint(model, optimizer, scheduler, epoch, best_val_acc, ckpt_dir=CKPT_DIR):
+    ckpt_dir.mkdir(parents=True, exist_ok=True)
+    raw = model._orig_mod if hasattr(model, '_orig_mod') else model
+    ckpt = {
+        "epoch": epoch,
+        "best_val_acc": best_val_acc,
+        "model_state_dict": raw.state_dict(),
+        "optimizer_state_dict": optimizer.state_dict(),
+        "scheduler_state_dict": scheduler.state_dict(),
+    }
+    path = ckpt_dir / f"epoch_{epoch:03d}.pt"
+    torch.save(ckpt, path)
+    latest = ckpt_dir / "latest.pt"
+    torch.save(ckpt, latest)
+    return path
+
+
+def load_checkpoint(model, optimizer, scheduler, ckpt_dir=CKPT_DIR):
+    latest = ckpt_dir / "latest.pt"
+    if not latest.exists():
+        return 0, 0.0
+    print(f"Resuming from {latest}...")
+    ckpt = torch.load(latest, weights_only=False)
+    raw = model._orig_mod if hasattr(model, '_orig_mod') else model
+    raw.load_state_dict(ckpt["model_state_dict"])
+    optimizer.load_state_dict(ckpt["optimizer_state_dict"])
+    scheduler.load_state_dict(ckpt["scheduler_state_dict"])
+    start_epoch = ckpt["epoch"] + 1
+    best_val_acc = ckpt["best_val_acc"]
+    print(f"Resumed: epoch {start_epoch}, best_val_acc={best_val_acc:.4f}")
+    return start_epoch, best_val_acc
+
+
+# --- Training ---
+
+def train(n_samples=500000, epochs=80, batch_size=4096, lr=3e-3, seed=42):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Device: {device}")
+
+    if device.type == "cuda":
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.allow_tf32 = True
+        if hasattr(torch.backends.cuda.matmul, 'fp32_precision'):
+            torch.backends.cuda.matmul.fp32_precision = 'tf32'
+        if hasattr(torch.backends.cudnn, 'conv') and hasattr(torch.backends.cudnn.conv, 'fp32_precision'):
+            torch.backends.cudnn.conv.fp32_precision = 'tf32'
+        torch.backends.cudnn.benchmark = True
+        props = torch.cuda.get_device_properties(0)
+        print(f"GPU: {props.name} | {props.total_memory / 1e9:.1f}GB | SM {props.major}.{props.minor}")
+        print(f"TF32: enabled | cuDNN benchmark: enabled | batch: {batch_size}")
+
+    train_ds, val_ds = get_or_generate_dataset(n_samples, seed)
+    print(f"Train: {len(train_ds)} | Val: {len(val_ds)} | {NUM_CLASSES} classes | pre-tensored")
+
+    train_loader = torch.utils.data.DataLoader(
+        train_ds, batch_size=batch_size, shuffle=True,
+        num_workers=4, pin_memory=True, persistent_workers=True)
+    val_loader = torch.utils.data.DataLoader(
+        val_ds, batch_size=batch_size, shuffle=False,
+        num_workers=4, pin_memory=True, persistent_workers=True)
+
+    model = GeometricShapeClassifier().to(device)
+    n_params = sum(p.numel() for p in model.parameters())
+    print(f"Model: {n_params:,} parameters")
+    if device.type == "cuda":
+        print(f"VRAM after model load: {torch.cuda.memory_allocated()/1e9:.2f}GB / "
+              f"{torch.cuda.get_device_properties(0).total_memory/1e9:.1f}GB")
+
+    use_amp = device.type == "cuda"
+    amp_dtype = torch.bfloat16 if (device.type == "cuda" and
+                torch.cuda.is_bf16_supported()) else torch.float16
+    use_scaler = use_amp and amp_dtype == torch.float16
+    scaler = torch.amp.GradScaler('cuda', enabled=use_scaler)
+
+    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)
+    warmup_epochs = 5
+    def lr_lambda(epoch):
+        if epoch < warmup_epochs:
+            return (epoch + 1) / warmup_epochs
+        return 0.5 * (1 + math.cos(math.pi * (epoch - warmup_epochs) / (epochs - warmup_epochs)))
+    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+
+    # Resume from checkpoint (loads into model BEFORE compile)
+    start_epoch, best_val_acc = load_checkpoint(model, optimizer, scheduler)
+
+    # Compile AFTER loading checkpoint weights
+    if device.type == "cuda" and hasattr(torch, 'compile'):
+        try:
+            model = torch.compile(model, mode="default")
+            print("torch.compile: enabled (default mode)")
+        except Exception as e:
+            print(f"torch.compile: skipped ({e})")
+
+    print(f"AMP: {'bf16' if amp_dtype == torch.bfloat16 else 'fp16'}" +
+          (f" (scaler: {'on' if use_scaler else 'off'})" if use_amp else " disabled"))
+
+    w = {"cls": 1.0, "fill": 0.3, "peak": 0.3, "ovf": 0.05,
+         "div": 0.02, "vol": 0.1, "cm": 0.1, "curved": 0.2, "ctype": 0.2,
+         "arb_cls": 0.8, "arb_traj": 0.2, "arb_conf": 0.1, "flow": 0.5}
+
+    epoch_start = time.time()
+
+    for epoch in range(start_epoch, epochs):
+        t0 = time.time()
+        model.train()
+        correct, total = 0, 0
+        correct_init, correct_ref = 0, 0
+
+        for batch_idx, (grid, label, dc, pd, vol, cm, ic, ct) in enumerate(train_loader):
+            grid, label = grid.to(device, non_blocking=True), label.to(device, non_blocking=True)
+            dc, pd = dc.to(device, non_blocking=True), pd.to(device, non_blocking=True)
+            vol, cm = vol.to(device, non_blocking=True), cm.to(device, non_blocking=True)
+            ic, ct = ic.to(device, non_blocking=True), ct.to(device, non_blocking=True)
+
+            grid = deform_grid(grid, p_dropout=0.05, p_add=0.05, p_shift=0.08)
+            optimizer.zero_grad(set_to_none=True)
+
+            try:
+                with torch.amp.autocast('cuda', enabled=use_amp, dtype=amp_dtype):
+                    out = model(grid, labels=label)
+
+                    loss_first = (w["cls"] * F.cross_entropy(out["initial_logits"], label) +
+                                  w["fill"] * capacity_fill_loss(out["fill_ratios"], dc) +
+                                  w["peak"] * peak_loss(out["peak_logits"], pd) +
+                                  w["ovf"] * overflow_reg(out["overflows"], dc) +
+                                  w["div"] * cap_diversity(out["capacities"]) +
+                                  w["vol"] * F.mse_loss(out["volume_pred"], torch.log1p(vol)) +
+                                  w["cm"] * cm_loss(out["cm_pred"], cm) +
+                                  w["curved"] * curved_bce(out["is_curved_pred"], ic) +
+                                  w["ctype"] * ctype_loss(out["curv_type_logits"], ct))
+
+                    loss_arb = w["arb_cls"] * F.cross_entropy(out["refined_logits"], label)
+                    traj_loss = 0
+                    for step_i, step_logits in enumerate(out["trajectory_logits"]):
+                        step_weight = (step_i + 1) / len(out["trajectory_logits"])
+                        traj_loss += step_weight * F.cross_entropy(step_logits, label)
+                    traj_loss /= len(out["trajectory_logits"])
+                    loss_arb += w["arb_traj"] * traj_loss
+                    loss_arb += w["flow"] * out["flow_loss"]
+
+                    with torch.no_grad():
+                        is_correct = (out["refined_logits"].argmax(1) == label).float()
+                    loss_arb += w["arb_conf"] * _safe_bce(
+                        out["refined_confidence"].squeeze(-1), is_correct)
+
+                    with torch.no_grad():
+                        init_correct = (out["initial_logits"].argmax(1) == label).float()
+                        ref_correct = (out["refined_logits"].argmax(1) == label).float()
+                        blend_target = torch.where(init_correct >= ref_correct,
+                            torch.ones_like(init_correct) * 0.8,
+                            torch.ones_like(init_correct) * 0.2)
+                    loss_arb += 0.1 * _safe_bce(out["blend_weight"], blend_target)
+
+                    loss_blend = w["cls"] * F.cross_entropy(out["class_logits"], label)
+                    loss = loss_first + loss_arb + loss_blend
+
+                # NaN guard: skip batch if loss is non-finite
+                if not torch.isfinite(loss).item():
+                    optimizer.zero_grad(set_to_none=True)
+                    total += grid.size(0)
+                    continue
+                scaler.scale(loss).backward()
+                scaler.unscale_(optimizer)
+                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+                scaler.step(optimizer)
+                scaler.update()
+
+            except RuntimeError as e:
+                if "CUDA" in str(e) or "device-side" in str(e):
+                    print(f"\n!!! CUDA error at epoch {epoch+1}, batch {batch_idx} !!!")
+                    print(f"    Error: {e}")
+                    print(f"    label range: [{label.min().item()}, {label.max().item()}]")
+                    print(f"    pd range: [{pd.min().item()}, {pd.max().item()}]")
+                    print(f"    ct range: [{ct.min().item()}, {ct.max().item()}]")
+                    print(f"    Checkpoint saved at epoch {epoch-1}")
+                    print(f"    To diagnose: add os.environ['CUDA_LAUNCH_BLOCKING']='1' before training")
+                raise
+
+            correct += (out["class_logits"].argmax(1) == label).sum().item()
+            correct_init += (out["initial_logits"].argmax(1) == label).sum().item()
+            correct_ref += (out["refined_logits"].argmax(1) == label).sum().item()
+            total += grid.size(0)
+
+        scheduler.step()
+        train_acc = correct / total
+
+        if epoch == start_epoch and device.type == "cuda":
+            peak = torch.cuda.max_memory_allocated() / 1e9
+            print(f"VRAM peak: {peak:.2f}GB | throughput: {total/(time.time()-t0):.0f} samples/s")
+
+        model.eval()
+        vc, vt, vcc, vct = 0, 0, 0, 0
+        vc_init, vc_ref = 0, 0
+        val_fills, val_alts, val_confs, val_blends = [], [], [], []
+
+        with torch.no_grad(), torch.amp.autocast('cuda', enabled=use_amp, dtype=amp_dtype):
+            for grid, label, dc, pd, vol, cm, ic, ct in val_loader:
+                grid, label = grid.to(device, non_blocking=True), label.to(device, non_blocking=True)
+                ic = ic.to(device, non_blocking=True)
+                out = model(grid)
+                vc += (out["class_logits"].argmax(1) == label).sum().item()
+                vc_init += (out["initial_logits"].argmax(1) == label).sum().item()
+                vc_ref += (out["refined_logits"].argmax(1) == label).sum().item()
+                vt += grid.size(0)
+                vcc += ((out["is_curved_pred"].squeeze(-1) > 0.5).float() == ic).sum().item()
+                vct += grid.size(0)
+                val_fills.append(out["fill_ratios"].cpu())
+                val_alts.append(out["alternation"].cpu())
+                val_confs.append(out["confidence"].cpu())
+                val_blends.append(out["blend_weight"].cpu())
+
+        val_acc = vc / vt; val_init = vc_init / vt; val_ref = vc_ref / vt
+        curved_acc = vcc / vct
+        mf = torch.cat(val_fills).mean(dim=0)
+        mc = torch.cat(val_confs).mean().item()
+        mb = torch.cat(val_blends).mean().item()
+        marker = " *" if val_acc > best_val_acc else ""
+        if val_acc > best_val_acc: best_val_acc = val_acc
+
+        raw_model = model._orig_mod if hasattr(model, '_orig_mod') else model
+        with torch.no_grad():
+            caps = [F.softplus(getattr(raw_model, f"dim{d}")._raw_capacity).item() for d in range(4)]
+
+        dt = time.time() - t0
+
+        # Save checkpoint every epoch
+        save_checkpoint(model, optimizer, scheduler, epoch, best_val_acc)
+
+        if (epoch + 1) % 10 == 0 or epoch == start_epoch or marker:
+            if (epoch + 1) % 10 == 0 or epoch == start_epoch:
+                print(f"Epoch {epoch+1:3d}/{epochs} [{dt:.1f}s {total/dt:.0f} s/s] | "
+                      f"blend {val_acc:.3f} init {val_init:.3f} arb {val_ref:.3f} | "
+                      f"conf {mc:.3f} blend_w {mb:.2f} | curved {curved_acc:.3f} | "
+                      f"fill [{mf[0]:.2f} {mf[1]:.2f} {mf[2]:.2f} {mf[3]:.2f}] | "
+                      f"cap [{caps[0]:.2f} {caps[1]:.2f} {caps[2]:.2f} {caps[3]:.2f}]{marker}")
+            elif marker:
+                print(f"Epoch {epoch+1:3d}/{epochs} [{dt:.1f}s] | "
+                      f"blend {val_acc:.3f} init {val_init:.3f} arb {val_ref:.3f} | conf {mc:.3f}{marker}")
+
+    total_time = time.time() - epoch_start
+    print(f"\nTraining complete in {total_time:.0f}s ({total_time/60:.1f}min)")
+    print(f"Best val accuracy: {best_val_acc:.4f}")
+    raw_model = model._orig_mod if hasattr(model, '_orig_mod') else model
+
+    # --- Per-class breakdown ---
+    cc_b, cc_i, cc_r = {n: 0 for n in CLASS_NAMES}, {n: 0 for n in CLASS_NAMES}, {n: 0 for n in CLASS_NAMES}
+    ct_c = {n: 0 for n in CLASS_NAMES}
+    cf = {n: [] for n in CLASS_NAMES}
+    cconf = {n: [] for n in CLASS_NAMES}
+    cblend = {n: [] for n in CLASS_NAMES}
+
+    with torch.no_grad(), torch.amp.autocast('cuda', enabled=use_amp, dtype=amp_dtype):
+        for grid, label, *_ in val_loader:
+            grid, label = grid.to(device, non_blocking=True), label.to(device, non_blocking=True)
+            out = model(grid)
+            pb = out["class_logits"].argmax(1)
+            pi = out["initial_logits"].argmax(1)
+            pr = out["refined_logits"].argmax(1)
+            for k in range(len(label)):
+                name = CLASS_NAMES[label[k].item()]
+                cc_b[name] += (pb[k] == label[k]).item()
+                cc_i[name] += (pi[k] == label[k]).item()
+                cc_r[name] += (pr[k] == label[k]).item()
+                ct_c[name] += 1
+                cf[name].append(out["fill_ratios"][k].cpu().numpy())
+                cconf[name].append(out["confidence"][k].item())
+                cblend[name].append(out["blend_weight"][k].item())
+
+    print(f"\n{'Class':22s} | {'Blend':>5s} {'Init':>5s} {'Arb':>5s} | "
+          f"{'Conf':>5s} {'Bld':>4s} | {'Corr':>4s}/{'Tot':>4s} | "
+          f"{'Fill Ratios':22s} | {'Type':8s} Curvature")
+    print("-" * 110)
+    for name in CLASS_NAMES:
+        if ct_c[name] == 0: continue
+        ab = cc_b[name]/ct_c[name]; ai = cc_i[name]/ct_c[name]; ar = cc_r[name]/ct_c[name]
+        mfv = np.mean(cf[name], axis=0)
+        mconf = np.mean(cconf[name]); mblend = np.mean(cblend[name])
+        info = SHAPE_CATALOG[name]
+        arb_flag = f" +{ar-ai:+.3f}" if ar-ai > 0.01 else ""
+        print(f"  {name:20s} | {ab:.3f} {ai:.3f} {ar:.3f} | "
+              f"{mconf:.3f} {mblend:.2f} | {cc_b[name]:4d}/{ct_c[name]:4d} | "
+              f"[{mfv[0]:.2f} {mfv[1]:.2f} {mfv[2]:.2f} {mfv[3]:.2f}] | "
+              f"{'CURVED' if info['curved'] else 'rigid':8s} {info['curvature']}{arb_flag}")
+
+    print(f"\n--- Arbiter Impact Summary ---")
+    imps = [(n, cc_i[n]/ct_c[n], cc_r[n]/ct_c[n], cc_b[n]/ct_c[n], cc_r[n]/ct_c[n]-cc_i[n]/ct_c[n])
+            for n in CLASS_NAMES if ct_c[n] > 0]
+    imps.sort(key=lambda x: x[4], reverse=True)
+    print(f"  {'Class':20s} | {'Init':>5s} {'Arb':>5s} {'Blend':>5s} | {'Δ':>6s}")
+    for name, ai, ar, ab, delta in imps[:10]:
+        print(f"  {name:20s} | {ai:.3f} {ar:.3f} {ab:.3f} | {delta:+.3f}")
+
+    # =========================================================================
+    # Upload geometric_classifier/ to HuggingFace
+    # =========================================================================
+    print("\n" + "=" * 70)
+    print("Saving geometric_classifier/ to HuggingFace")
+    print("=" * 70)
+
+    raw_model = model._orig_mod if hasattr(model, '_orig_mod') else model
+    staging = Path("./hf_staging/geometric_classifier")
+    staging.mkdir(parents=True, exist_ok=True)
+
+    arch_config = {
+        "model_type": "GeometricShapeClassifier",
+        "version": "v8",
+        "grid_size": GS,
+        "num_classes": NUM_CLASSES,
+        "class_names": CLASS_NAMES,
+        "curvature_types": CURVATURE_TYPES,
+        "embed_dim": 128,
+        "n_tracers": 5,
+        "capacity_dims": [64, 64, 64, 64],
+        "curvature_embed_dim": 128,
+        "arbiter_latent_dim": 128,
+        "arbiter_flow_steps": 4,
+        "total_params": sum(p.numel() for p in raw_model.parameters()),
+        "shape_catalog": {k: v for k, v in SHAPE_CATALOG.items()},
+    }
+    with open(staging / "config.json", "w") as f:
+        json.dump(arch_config, f, indent=2)
+
+    train_cfg = {
+        "n_samples": n_samples, "epochs": epochs, "batch_size": batch_size,
+        "lr": lr, "seed": seed, "optimizer": "AdamW", "weight_decay": 1e-4,
+        "scheduler": "cosine_with_warmup", "warmup_epochs": warmup_epochs,
+        "amp_dtype": str(amp_dtype), "loss_weights": w,
+        "best_val_accuracy": best_val_acc, "learned_capacities": caps,
+        "total_training_time_seconds": total_time,
+    }
+    with open(staging / "training_config.json", "w") as f:
+        json.dump(train_cfg, f, indent=2)
+
+    try:
+        from safetensors.torch import save_file as st_save
+        st_save(raw_model.state_dict(), str(staging / "model.safetensors"))
+        print(f"  Saved: model.safetensors")
+    except ImportError:
+        torch.save(raw_model.state_dict(), staging / "model.pt")
+        print(f"  Saved: model.pt (install safetensors for .safetensors)")
+
+    try:
+        from huggingface_hub import HfApi, create_repo
+        token = None
+        try:
+            from google.colab import userdata
+            token = userdata.get('HF_TOKEN')
+        except Exception:
+            token = os.environ.get('HF_TOKEN')
+
+        if token:
+            api = HfApi(token=token)
+            create_repo(HF_REPO, token=token, exist_ok=True)
+            readme = Path("./hf_staging/README.md")
+            readme.write_text(f"""---
+license: mit
+tags:
+  - geometric-deep-learning
+  - voxel-classifier
+  - cross-contrast
+  - pentachoron
+---
+
+# Grid Geometric Classifier Proto
+
+Geometric primitive classifier using 5x5x5 binary voxel grids with capacity cascade,
+curvature analysis, differentiation gates, and rectified flow arbiter.
+
+## Structure
+
+```
+geometric_classifier/   # Voxel classifier (v8, ~1.85M params)
+crosscontrast/           # Text-Voxel alignment heads
+qwen_embeddings/         # Cached Qwen 2.5-1.5B embeddings
+```
+
+## 38 Shape Classes
+
+Rigid 0D-3D: point, lines, triangles, quads, tetrahedra, cubes, prisms, octahedra, pentachoron
+Curved 1D-3D: arcs, helices, circles, ellipses, discs, spheres, hemispheres, cylinders, cones, capsules, tori, shells, tubes, bowls, saddles
+""")
+            api.upload_file(path_or_fileobj=str(readme), path_in_repo="README.md",
+                           repo_id=HF_REPO, token=token, commit_message="README")
+            api.upload_folder(
+                folder_path=str(staging), repo_id=HF_REPO,
+                path_in_repo="geometric_classifier", token=token,
+                commit_message=f"geometric_classifier v8 | acc={best_val_acc:.4f} | {sum(p.numel() for p in raw_model.parameters()):,} params")
+            print(f"Uploaded: https://huggingface.co/{HF_REPO}/tree/main/geometric_classifier")
+        else:
+            print("No HF_TOKEN — saved locally at ./hf_staging/geometric_classifier/")
+    except Exception as e:
+        print(f"HF upload failed: {e}\n  Weights at ./hf_staging/geometric_classifier/")
+
+    return model
+
+
+model = train()