Create qwen_simple_test_trainer.py

qwen_simple_test_trainer.py

@@ -0,0 +1,620 @@
+# =============================================================================
+# CELL 4: Qwen × Geometric Classifier Cross-Contrast Training
+# Requires: Cell 1 (constants), Cell 2 (model classes), Cell 3 (trained checkpoint)
+# Outputs: crosscontrast/ and qwen_embeddings/ on HF
+#
+# Features:
+#   - Loads geo classifier from Cell 3 checkpoint (no notebook scope dependency)
+#   - Uses model.forward()["features"] (no duplicated internals)
+#   - Dataset + Qwen embeddings cached to disk
+#   - CC model checkpointed with resume
+# =============================================================================
+
+import math, time, json, os
+from pathlib import Path
+from tqdm.auto import tqdm
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+HF_REPO = "AbstractPhil/grid-geometric-classifier-proto"
+CKPT_DIR = Path("./checkpoints")
+CC_CKPT_DIR = Path("./cc_checkpoints")
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+torch.backends.cudnn.benchmark = True
+
+use_amp = device.type == "cuda"
+amp_dtype = torch.bfloat16 if (device.type == "cuda" and
+            torch.cuda.is_bf16_supported()) else torch.float16
+
+
+# =============================================================================
+# Shape Descriptions
+# =============================================================================
+
+SHAPE_DESCRIPTIONS = {
+    "point": "A zero-dimensional geometric primitive occupying a single discrete location in three-dimensional space with no extent along any axis.",
+    "line_x": "A one-dimensional line segment extending along the horizontal x-axis, connecting two endpoints with uniform spacing between occupied voxels.",
+    "line_y": "A one-dimensional line segment extending along the vertical y-axis, a straight structure rising upward through the grid.",
+    "line_z": "A one-dimensional line segment extending along the depth z-axis, projecting straight backward through the voxel grid.",
+    "line_diag": "A one-dimensional diagonal line segment cutting across multiple axes simultaneously, connecting opposite corners of the grid.",
+    "cross": "Two perpendicular line segments intersecting at their midpoints forming a plus-shaped cross pattern in a single plane.",
+    "l_shape": "Two connected line segments meeting at a right angle to form an L-shaped corner, like two edges of a rectangle.",
+    "collinear": "Three or more points arranged along a single straight line with equal spacing, demonstrating perfect linear alignment.",
+    "triangle_xy": "A flat triangular outline formed by three connected edges lying in the horizontal xy-plane, the simplest polygon.",
+    "triangle_xz": "A flat triangular outline formed by three connected edges lying in the vertical xz-plane, a triangle standing upright.",
+    "triangle_3d": "A triangular outline with vertices at different heights, forming a non-planar triangle tilted in three-dimensional space.",
+    "square_xy": "A square outline formed by four equal edges in the xy-plane, a regular quadrilateral with right angles at each corner.",
+    "square_xz": "A square outline formed by four equal edges in the xz-plane, a square standing vertically like a window frame.",
+    "rectangle": "A rectangular outline with two pairs of parallel edges of different lengths, wider than it is tall.",
+    "coplanar": "A set of points all lying in the same plane but not forming a regular polygon, a scattered planar arrangement.",
+    "plane": "A solid flat surface filling an entire plane with occupied voxels, a two-dimensional sheet extending across the grid.",
+    "tetrahedron": "A three-dimensional simplex with four triangular faces meeting at four vertices and six edges, the simplest polyhedron.",
+    "pyramid": "A solid with a square base and four triangular faces converging to a single apex point above the base center.",
+    "pentachoron": "A four-dimensional simplex projected into three dimensions, consisting of five tetrahedral cells sharing faces.",
+    "cube": "A regular hexahedron with six identical square faces, twelve edges, and eight vertices forming a perfect box shape.",
+    "cuboid": "A rectangular box with six rectangular faces, similar to a cube but with at least one pair of edges longer than the others.",
+    "triangular_prism": "A solid with two parallel triangular faces connected by three rectangular faces, like a tent or Toblerone shape.",
+    "octahedron": "A regular polyhedron with eight equilateral triangular faces, twelve edges, and six vertices, like two pyramids base-to-base.",
+    "arc": "A curved one-dimensional segment forming part of a circle, a smooth bend connecting two endpoints along a circular path.",
+    "helix": "A three-dimensional spiral curve that winds around a central axis while advancing along it, like a corkscrew or spring.",
+    "circle": "A closed curved outline where every point is equidistant from the center, forming a perfect round ring in a plane.",
+    "ellipse": "A closed curved outline forming an elongated circle, an oval shape with two focal points and varying curvature.",
+    "disc": "A solid filled circular region, a flat round plate occupying all voxels within a circular boundary in a plane.",
+    "sphere": "A perfectly round three-dimensional solid where every surface point is equidistant from the center, fully filled inside.",
+    "hemisphere": "Half of a sphere cut along a great circle, a dome shape with a flat base and a convex curved upper surface.",
+    "cylinder": "A solid with two parallel circular faces connected by a curved rectangular surface, like a can or pillar.",
+    "cone": "A solid tapering smoothly from a circular base to a single apex point, with a curved surface of decreasing radius.",
+    "capsule": "A cylinder capped with hemispheres at both ends, a smooth elongated pill shape with no sharp edges.",
+    "torus": "A donut-shaped solid formed by revolving a circle around an external axis, with a hole through the center.",
+    "shell": "A hollow spherical surface with no interior fill, an empty ball where only the outer boundary layer is occupied.",
+    "tube": "A hollow cylindrical surface with no interior fill, an empty pipe where only the curved wall is occupied.",
+    "bowl": "A concave open surface curving inward like a dish, the bottom half of a hollow sphere with the opening facing up.",
+    "saddle": "A hyperbolic surface that curves upward along one axis and downward along the perpendicular axis, like a horse saddle.",
+}
+assert set(SHAPE_DESCRIPTIONS.keys()) == set(CLASS_NAMES), "Description/class mismatch!"
+
+
+# =============================================================================
+# Qwen Embedding Extractor
+# =============================================================================
+
+class QwenEmbeddingExtractor:
+    MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"
+    HIDDEN_DIM = 1536
+
+    def __init__(self, device="cuda"):
+        self.device = device
+        self.model = None
+        self.tokenizer = None
+
+    def load_model(self):
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+        print(f"Loading {self.MODEL_NAME}...")
+        self.tokenizer = AutoTokenizer.from_pretrained(self.MODEL_NAME, trust_remote_code=True)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            self.MODEL_NAME, dtype=torch.float16,
+            device_map=self.device, trust_remote_code=True)
+        self.model.eval()
+        print(f"Qwen loaded: {self.HIDDEN_DIM}-dim hidden states")
+
+    def _build_encode_prompt(self, description):
+        messages = [
+            {"role": "system", "content": "You are a geometric shape analyst."},
+            {"role": "user", "content": f"Analyze this shape: {description}"},
+        ]
+        return self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+
+    @torch.no_grad()
+    def extract_embedding(self, text):
+        prompt = self._build_encode_prompt(text)
+        inputs = self.tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True).to(self.device)
+        outputs = self.model(**inputs, output_hidden_states=True)
+        hidden = outputs.hidden_states[-1]
+        return hidden.mean(dim=1).squeeze(0).float()
+
+    def cache_all_embeddings(self, class_names):
+        print(f"Extracting embeddings for {len(class_names)} classes...")
+        embeddings = {}
+        for name in class_names:
+            embeddings[name] = self.extract_embedding(SHAPE_DESCRIPTIONS[name])
+        emb_tensor = torch.stack([embeddings[n] for n in class_names])
+        normed = F.normalize(emb_tensor, dim=-1)
+        sim = normed @ normed.T
+        mean_sim = (sim.sum() - sim.trace()) / (len(class_names) * (len(class_names) - 1))
+        print(f"Cached: {emb_tensor.shape} | mean cross-class sim: {mean_sim:.4f}")
+        return emb_tensor
+
+    def unload(self):
+        del self.model, self.tokenizer
+        self.model = self.tokenizer = None
+        torch.cuda.empty_cache()
+        print("Qwen unloaded")
+
+
+# =============================================================================
+# Projection Heads + Cross-Contrast Model
+# =============================================================================
+
+class TextProjection(nn.Module):
+    def __init__(self, text_dim=1536, latent_dim=256):
+        super().__init__()
+        self.proj = nn.Sequential(
+            nn.Linear(text_dim, latent_dim * 2), nn.GELU(),
+            nn.Linear(latent_dim * 2, latent_dim), nn.GELU(),
+            nn.Linear(latent_dim, latent_dim))
+        self.norm = nn.LayerNorm(latent_dim)
+    def forward(self, x): return self.norm(self.proj(x))
+
+class VoxelProjection(nn.Module):
+    def __init__(self, voxel_dim=645, latent_dim=256):
+        super().__init__()
+        self.proj = nn.Sequential(
+            nn.Linear(voxel_dim, latent_dim * 2), nn.GELU(),
+            nn.Linear(latent_dim * 2, latent_dim), nn.GELU(),
+            nn.Linear(latent_dim, latent_dim))
+        self.norm = nn.LayerNorm(latent_dim)
+    def forward(self, x): return self.norm(self.proj(x))
+
+
+class CrossContrastModel(nn.Module):
+    def __init__(self, text_dim=1536, voxel_dim=645, latent_dim=256,
+                 n_classes=38, temperature=0.07):
+        super().__init__()
+        self.text_proj = TextProjection(text_dim, latent_dim)
+        self.voxel_proj = VoxelProjection(voxel_dim, latent_dim)
+        self.log_temperature = nn.Parameter(torch.tensor(math.log(1.0 / temperature)))
+
+    @property
+    def temperature(self):
+        return torch.exp(-self.log_temperature)
+
+    def forward(self, voxel_features, class_labels, text_embeddings_table):
+        text_emb = text_embeddings_table[class_labels]
+        z_text = F.normalize(self.text_proj(text_emb), dim=-1)
+        z_voxel = F.normalize(self.voxel_proj(voxel_features), dim=-1)
+
+        temp = self.temperature
+        logits_v2t = z_voxel @ z_text.T / temp
+        logits_t2v = z_text @ z_voxel.T / temp
+
+        labels_matrix = (class_labels.unsqueeze(0) == class_labels.unsqueeze(1)).float()
+        labels_matrix = labels_matrix / labels_matrix.sum(dim=1, keepdim=True).clamp(min=1)
+
+        loss_v2t = (-labels_matrix * F.log_softmax(logits_v2t, dim=1)).sum(dim=1).mean()
+        loss_t2v = (-labels_matrix * F.log_softmax(logits_t2v, dim=1)).sum(dim=1).mean()
+        loss = (loss_v2t + loss_t2v) / 2.0
+
+        with torch.no_grad():
+            v2t_preds = logits_v2t.argmax(dim=1)
+            pred_classes = class_labels[v2t_preds]
+            acc = (pred_classes == class_labels).float().mean()
+            pos_sim = (z_voxel * z_text).sum(dim=-1).mean()
+            neg_mask = ~(class_labels.unsqueeze(0) == class_labels.unsqueeze(1))
+            neg_sim = (z_voxel @ z_text.T)[neg_mask].mean() if neg_mask.any() else torch.tensor(0.0)
+
+        return loss, {"acc": acc.item(), "pos_sim": pos_sim.item(),
+                      "neg_sim": neg_sim.item(), "temperature": temp.item()}
+
+
+# =============================================================================
+# Helpers
+# =============================================================================
+
+def get_hf_token():
+    try:
+        from google.colab import userdata
+        return userdata.get('HF_TOKEN')
+    except Exception:
+        return os.environ.get('HF_TOKEN')
+
+
+def load_geo_model(ckpt_dir=CKPT_DIR):
+    """Load trained GeometricShapeClassifier from Cell 3 checkpoint."""
+    latest = ckpt_dir / "latest.pt"
+    if not latest.exists():
+        raise FileNotFoundError(
+            f"No checkpoint at {latest}. Run Cell 3 first to train the classifier.")
+    print(f"Loading geo classifier from {latest}...")
+    ckpt = torch.load(latest, weights_only=False, map_location=device)
+    geo = GeometricShapeClassifier().to(device)
+    geo.load_state_dict(ckpt["model_state_dict"])
+    geo.eval()
+    for p in geo.parameters():
+        p.requires_grad = False
+    print(f"Loaded: epoch {ckpt['epoch']}, val_acc={ckpt['best_val_acc']:.4f}, "
+          f"{sum(p.numel() for p in geo.parameters()):,} params (frozen)")
+    return geo
+
+
+@torch.no_grad()
+def extract_voxel_features(geo_model, grid):
+    """Extract pre-classifier features using model.forward()['features']."""
+    with torch.amp.autocast('cuda', enabled=use_amp, dtype=amp_dtype):
+        out = geo_model(grid)
+    return out["features"].float()
+
+
+def save_cc_checkpoint(cc_model, cc_opt, cc_sched, epoch, best_acc, ckpt_dir=CC_CKPT_DIR):
+    ckpt_dir.mkdir(parents=True, exist_ok=True)
+    ckpt = {
+        "epoch": epoch,
+        "best_cc_acc": best_acc,
+        "cc_model_state_dict": cc_model.state_dict(),
+        "cc_optimizer_state_dict": cc_opt.state_dict(),
+        "cc_scheduler_state_dict": cc_sched.state_dict(),
+    }
+    torch.save(ckpt, ckpt_dir / "latest.pt")
+
+
+def load_cc_checkpoint(cc_model, cc_opt, cc_sched, ckpt_dir=CC_CKPT_DIR):
+    latest = ckpt_dir / "latest.pt"
+    if not latest.exists():
+        return 0, 0.0
+    print(f"Resuming CC from {latest}...")
+    ckpt = torch.load(latest, weights_only=False)
+    cc_model.load_state_dict(ckpt["cc_model_state_dict"])
+    cc_opt.load_state_dict(ckpt["cc_optimizer_state_dict"])
+    cc_sched.load_state_dict(ckpt["cc_scheduler_state_dict"])
+    start = ckpt["epoch"] + 1
+    best = ckpt["best_cc_acc"]
+    print(f"Resumed: epoch {start}, best_cc_acc={best:.4f}")
+    return start, best
+
+
+def upload_cc_to_hf(cc_model, best_acc, epoch, token, reason="periodic",
+                    text_dim_=None, voxel_dim_=None, latent_dim_=None):
+    """Upload crosscontrast weights + config to HF. Called mid-training and at end."""
+    if not token:
+        return
+    try:
+        from huggingface_hub import HfApi, create_repo
+        from safetensors.torch import save_file as st_save
+
+        staging = Path("./hf_staging/crosscontrast")
+        staging.mkdir(parents=True, exist_ok=True)
+
+        st_save(cc_model.text_proj.state_dict(), str(staging / "text_proj.safetensors"))
+        st_save(cc_model.voxel_proj.state_dict(), str(staging / "voxel_proj.safetensors"))
+        st_save({"log_temperature": cc_model.log_temperature.data.unsqueeze(0)},
+                str(staging / "temperature.safetensors"))
+
+        # Write config so uploads are self-contained
+        if text_dim_ and voxel_dim_ and latent_dim_:
+            cfg = {
+                "model_type": "CrossContrastModel",
+                "text_dim": text_dim_, "voxel_dim": voxel_dim_,
+                "latent_dim": latent_dim_,
+                "best_val_accuracy": best_acc,
+                "epoch": epoch + 1, "upload_reason": reason,
+                "temperature": cc_model.temperature.item(),
+            }
+            with open(staging / "config.json", "w") as f:
+                json.dump(cfg, f, indent=2)
+
+        api = HfApi(token=token)
+        create_repo(HF_REPO, token=token, exist_ok=True)
+        api.upload_folder(
+            folder_path=str(staging), repo_id=HF_REPO,
+            path_in_repo="crosscontrast", token=token,
+            commit_message=f"crosscontrast ep{epoch+1} | acc={best_acc:.4f} | {reason}")
+        tqdm.write(f"  ✓ HF upload ({reason}): ep{epoch+1} acc={best_acc:.4f}")
+    except Exception as e:
+        tqdm.write(f"  ⚠ HF upload failed: {e}")
+
+
+# =============================================================================
+# PHASE 1: Qwen Embeddings (cached)
+# =============================================================================
+
+print("=" * 70)
+print("Phase 1: Qwen Embeddings")
+print("=" * 70)
+
+cache_path = Path("qwen_geo_cache.pt")
+if cache_path.exists():
+    _cache = torch.load(cache_path, map_location="cpu", weights_only=True)
+    text_embeddings = _cache["embeddings"]
+    print(f"Loaded cached: {text_embeddings.shape}")
+else:
+    extractor = QwenEmbeddingExtractor(device=str(device))
+    extractor.load_model()
+    text_embeddings = extractor.cache_all_embeddings(CLASS_NAMES)
+    torch.save({"embeddings": text_embeddings.cpu(), "class_names": CLASS_NAMES}, cache_path)
+    extractor.unload()
+
+text_embeddings = text_embeddings.to(device)
+text_dim = text_embeddings.shape[1]
+
+# Upload qwen_embeddings/ to HF
+hf_token = get_hf_token()
+qwen_staging = Path("./hf_staging/qwen_embeddings")
+qwen_staging.mkdir(parents=True, exist_ok=True)
+
+qwen_config = {
+    "model_name": QwenEmbeddingExtractor.MODEL_NAME,
+    "hidden_dim": QwenEmbeddingExtractor.HIDDEN_DIM,
+    "extraction_method": "mean_pool_last_layer",
+    "prompt_style": "2shot_geometric",
+    "num_classes": NUM_CLASSES,
+    "class_names": CLASS_NAMES,
+    "embedding_shape": list(text_embeddings.shape),
+}
+with open(qwen_staging / "config.json", "w") as f:
+    json.dump(qwen_config, f, indent=2)
+with open(qwen_staging / "descriptions.json", "w") as f:
+    json.dump(SHAPE_DESCRIPTIONS, f, indent=2)
+
+try:
+    from safetensors.torch import save_file as st_save
+    st_save({"embeddings": text_embeddings.cpu()}, str(qwen_staging / "embeddings.safetensors"))
+except ImportError:
+    torch.save({"embeddings": text_embeddings.cpu()}, qwen_staging / "embeddings.pt")
+
+if hf_token:
+    try:
+        from huggingface_hub import HfApi, create_repo
+        api = HfApi(token=hf_token)
+        create_repo(HF_REPO, token=hf_token, exist_ok=True)
+        api.upload_folder(
+            folder_path=str(qwen_staging), repo_id=HF_REPO,
+            path_in_repo="qwen_embeddings", token=hf_token,
+            commit_message=f"qwen_embeddings | {QwenEmbeddingExtractor.MODEL_NAME} | {NUM_CLASSES} classes")
+        print(f"Uploaded: https://huggingface.co/{HF_REPO}/tree/main/qwen_embeddings")
+    except Exception as e:
+        print(f"Upload failed: {e}")
+else:
+    print("No HF_TOKEN — saved locally")
+
+
+# =============================================================================
+# PHASE 2: Load Geo Model + Dataset
+# =============================================================================
+
+print("\n" + "=" * 70)
+print("Phase 2: Geo Model + Voxel Data")
+print("=" * 70)
+
+geo_model = load_geo_model()
+
+CC_SAMPLES = 500000
+CC_BATCH = 4096
+CC_EPOCHS = 40
+CC_LR = 2e-3
+CC_LATENT = 256
+
+# Reuse Cell 3's cached dataset
+DATASET_PATH = Path("./cached_dataset.pt")
+if not DATASET_PATH.exists():
+    raise FileNotFoundError("No cached_dataset.pt — run Cell 3 first.")
+
+print(f"Loading dataset from {DATASET_PATH}...")
+_cached = torch.load(DATASET_PATH, weights_only=True)
+cc_train_ds = ShapeDataset.__new__(ShapeDataset)
+cc_val_ds = ShapeDataset.__new__(ShapeDataset)
+for k in ["grids", "labels", "dim_conf", "peak_dim", "volume", "cm_det", "is_curved", "curvature"]:
+    setattr(cc_train_ds, k, _cached["train"][k])
+    setattr(cc_val_ds, k, _cached["val"][k])
+print(f"Loaded {len(cc_train_ds)} train + {len(cc_val_ds)} val (from Cell 3 cache)")
+
+cc_train_loader = torch.utils.data.DataLoader(
+    cc_train_ds, batch_size=CC_BATCH, shuffle=True,
+    num_workers=4, pin_memory=True, persistent_workers=True)
+cc_val_loader = torch.utils.data.DataLoader(
+    cc_val_ds, batch_size=CC_BATCH, shuffle=False,
+    num_workers=4, pin_memory=True, persistent_workers=True)
+
+# Probe voxel feature dim
+with torch.no_grad():
+    _dummy = torch.zeros(1, GS, GS, GS, device=device)
+    voxel_dim = extract_voxel_features(geo_model, _dummy).shape[1]
+print(f"Voxel dim: {voxel_dim} | Text dim: {text_dim}")
+
+
+# =============================================================================
+# PHASE 3: Cross-Contrast Training
+# =============================================================================
+
+print("\n" + "=" * 70)
+print("Phase 3: Cross-Contrast Training")
+print("=" * 70)
+
+cc_model = CrossContrastModel(
+    text_dim=text_dim, voxel_dim=voxel_dim,
+    latent_dim=CC_LATENT, n_classes=NUM_CLASSES
+).to(device)
+cc_params = sum(p.numel() for p in cc_model.parameters())
+print(f"CrossContrast: {cc_params:,} params | latent={CC_LATENT}")
+
+cc_opt = torch.optim.AdamW(cc_model.parameters(), lr=CC_LR, weight_decay=1e-4)
+_warmup = 3
+def _cc_lr(ep):
+    if ep < _warmup: return (ep + 1) / _warmup
+    return 0.5 * (1 + math.cos(math.pi * (ep - _warmup) / (CC_EPOCHS - _warmup)))
+cc_sched = torch.optim.lr_scheduler.LambdaLR(cc_opt, _cc_lr)
+
+# Resume CC training if checkpoint exists
+cc_start, best_cc_acc = load_cc_checkpoint(cc_model, cc_opt, cc_sched)
+
+t_start = time.time()
+epoch_bar = tqdm(range(cc_start, CC_EPOCHS), desc="Training", unit="ep")
+
+for epoch in epoch_bar:
+    t0 = time.time()
+    cc_model.train()
+    tot_loss, tot_acc, nb = 0, 0, 0
+
+    batch_bar = tqdm(cc_train_loader, desc=f"Ep {epoch+1}/{CC_EPOCHS} train",
+                     leave=False, unit="batch")
+    for grid, label, *_ in batch_bar:
+        grid = grid.to(device, non_blocking=True)
+        label = label.to(device, non_blocking=True)
+        vf = extract_voxel_features(geo_model, grid)
+
+        cc_opt.zero_grad(set_to_none=True)
+        with torch.amp.autocast('cuda', enabled=use_amp, dtype=amp_dtype):
+            loss, metrics = cc_model(vf, label, text_embeddings)
+
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(cc_model.parameters(), 1.0)
+        cc_opt.step()
+        tot_loss += loss.item(); tot_acc += metrics["acc"]; nb += 1
+        batch_bar.set_postfix(loss=f"{loss.item():.4f}", acc=f"{metrics['acc']:.3f}")
+
+    cc_sched.step()
+
+    cc_model.eval()
+    vl, va, vps, vns, vnb = 0, 0, 0, 0, 0
+    with torch.no_grad(), torch.amp.autocast('cuda', enabled=use_amp, dtype=amp_dtype):
+        for grid, label, *_ in tqdm(cc_val_loader, desc=f"Ep {epoch+1}/{CC_EPOCHS} val",
+                                     leave=False, unit="batch"):
+            grid = grid.to(device, non_blocking=True)
+            label = label.to(device, non_blocking=True)
+            vf = extract_voxel_features(geo_model, grid)
+            loss, m = cc_model(vf, label, text_embeddings)
+            vl += loss.item(); va += m["acc"]; vps += m["pos_sim"]; vns += m["neg_sim"]; vnb += 1
+
+    tl = tot_loss/max(nb,1); ta = tot_acc/max(nb,1)
+    vl_ = vl/max(vnb,1); va_ = va/max(vnb,1)
+    ps = vps/max(vnb,1); ns = vns/max(vnb,1)
+    temp = cc_model.temperature.item()
+    dt = time.time() - t0
+    mk = " *" if va_ > best_cc_acc else ""
+    if va_ > best_cc_acc: best_cc_acc = va_
+
+    save_cc_checkpoint(cc_model, cc_opt, cc_sched, epoch, best_cc_acc)
+
+    # Upload to HF on new best or every 10 epochs
+    is_new_best = mk == " *"
+    periodic = (epoch + 1) % 10 == 0
+    if is_new_best:
+        upload_cc_to_hf(cc_model, best_cc_acc, epoch, hf_token, reason="new_best",
+                        text_dim_=text_dim, voxel_dim_=voxel_dim, latent_dim_=CC_LATENT)
+    elif periodic:
+        upload_cc_to_hf(cc_model, best_cc_acc, epoch, hf_token, reason="periodic",
+                        text_dim_=text_dim, voxel_dim_=voxel_dim, latent_dim_=CC_LATENT)
+
+    epoch_bar.set_postfix(loss=f"{vl_:.4f}", acc=f"{va_:.3f}", best=f"{best_cc_acc:.3f}",
+                          tau=f"{temp:.4f}")
+
+    if (epoch+1) % 5 == 0 or epoch == cc_start or mk:
+        tqdm.write(f"Ep {epoch+1:3d}/{CC_EPOCHS} [{dt:.1f}s] | "
+                   f"loss {tl:.4f}/{vl_:.4f} | acc {ta:.3f}/{va_:.3f} | "
+                   f"pos {ps:.3f} neg {ns:.3f} | τ {temp:.4f}{mk}")
+
+    if epoch == cc_start and device.type == "cuda":
+        tqdm.write(f"VRAM peak: {torch.cuda.max_memory_allocated()/1e9:.2f}GB")
+
+tt = time.time() - t_start
+print(f"\nDone in {tt:.0f}s ({tt/60:.1f}min) | Best acc: {best_cc_acc:.4f}")
+
+
+# =============================================================================
+# Per-Class Alignment Analysis
+# =============================================================================
+
+print("\n" + "=" * 70)
+print("Per-Class Alignment")
+print("=" * 70)
+
+cc_model.eval()
+cls_vox = {n: [] for n in CLASS_NAMES}
+with torch.no_grad():
+    text_proj_all = F.normalize(cc_model.text_proj(text_embeddings), dim=-1)
+
+with torch.no_grad(), torch.amp.autocast('cuda', enabled=use_amp, dtype=amp_dtype):
+    for grid, label, *_ in cc_val_loader:
+        grid = grid.to(device, non_blocking=True)
+        label = label.to(device, non_blocking=True)
+        vf = extract_voxel_features(geo_model, grid)
+        zv = F.normalize(cc_model.voxel_proj(vf), dim=-1)
+        for k in range(len(label)):
+            cls_vox[CLASS_NAMES[label[k].item()]].append(zv[k].cpu())
+
+print(f"\n{'Class':22s} | {'Align':>6s} | {'N':>5s} | {'Nearest':22s} | {'OK':>3s}")
+print("-" * 70)
+correct = 0; total_c = 0
+for name in CLASS_NAMES:
+    if not cls_vox[name]: continue
+    mv = F.normalize(torch.stack(cls_vox[name]).mean(dim=0), dim=-1)
+    own = text_proj_all[CLASS_NAMES.index(name)].cpu()
+    align = (mv * own).sum().item()
+    sims = (mv.unsqueeze(0) @ text_proj_all.cpu().T).squeeze(0)
+    ni = sims.argmax().item(); nn_ = CLASS_NAMES[ni]
+    ok = "Y" if nn_ == name else f"X->{nn_}"
+    if nn_ == name: correct += 1
+    total_c += 1
+    print(f"  {name:20s} | {align:.4f} | {len(cls_vox[name]):5d} | {nn_:22s} | {ok}")
+print(f"\nNearest-text accuracy: {correct}/{total_c} = {correct/max(total_c,1):.1%}")
+
+print("\nTop 10 Text-Space Confusions:")
+with torch.no_grad():
+    sim = (text_proj_all @ text_proj_all.T).cpu().numpy()
+import numpy as np
+confusions = []
+for i in range(len(CLASS_NAMES)):
+    for j in range(i+1, len(CLASS_NAMES)):
+        confusions.append((CLASS_NAMES[i], CLASS_NAMES[j], sim[i,j]))
+confusions.sort(key=lambda x: x[2], reverse=True)
+for a, b, s in confusions[:10]:
+    print(f"  {a:20s} <-> {b:20s} | {s:.4f}")
+
+
+# =============================================================================
+# Upload crosscontrast/ to HuggingFace (final — with full configs)
+# =============================================================================
+
+print("\n" + "=" * 70)
+print("Saving crosscontrast/ to HuggingFace")
+print("=" * 70)
+
+cc_staging = Path("./hf_staging/crosscontrast")
+cc_staging.mkdir(parents=True, exist_ok=True)
+
+# Write detailed architecture config (overwrites minimal mid-training config)
+cc_arch = {
+    "model_type": "CrossContrastModel",
+    "text_dim": text_dim,
+    "voxel_dim": voxel_dim,
+    "latent_dim": CC_LATENT,
+    "n_classes": NUM_CLASSES,
+    "text_proj_layers": [text_dim, CC_LATENT * 2, CC_LATENT, CC_LATENT],
+    "voxel_proj_layers": [voxel_dim, CC_LATENT * 2, CC_LATENT, CC_LATENT],
+    "activation": "GELU",
+    "normalization": "LayerNorm",
+    "total_params": cc_params,
+    "class_names": CLASS_NAMES,
+    "text_encoder": QwenEmbeddingExtractor.MODEL_NAME,
+    "voxel_encoder": "GeometricShapeClassifier_v8",
+}
+with open(cc_staging / "config.json", "w") as f:
+    json.dump(cc_arch, f, indent=2)
+
+cc_train_cfg = {
+    "n_samples": CC_SAMPLES, "epochs": CC_EPOCHS, "batch_size": CC_BATCH,
+    "lr": CC_LR, "weight_decay": 1e-4, "optimizer": "AdamW",
+    "scheduler": "cosine_with_warmup", "warmup_epochs": _warmup,
+    "loss": "symmetric_InfoNCE", "initial_temperature": 0.07,
+    "final_temperature": cc_model.temperature.item(),
+    "amp_dtype": str(amp_dtype),
+    "best_val_accuracy": best_cc_acc,
+    "nearest_text_accuracy": correct / max(total_c, 1),
+    "total_training_time_seconds": tt,
+}
+with open(cc_staging / "training_config.json", "w") as f:
+    json.dump(cc_train_cfg, f, indent=2)
+
+# Final upload with full configs + weights
+upload_cc_to_hf(cc_model, best_cc_acc, CC_EPOCHS - 1, hf_token, reason="final",
+                text_dim_=text_dim, voxel_dim_=voxel_dim, latent_dim_=CC_LATENT)
+if not hf_token:
+    print("No HF_TOKEN — saved locally at ./hf_staging/crosscontrast/")
+
+print(f"\nAll three subdirectories staged:")
+print(f"  geometric_classifier/ — from Cell 3")
+print(f"  qwen_embeddings/      — {text_embeddings.shape}")
+print(f"  crosscontrast/        — latent={CC_LATENT}, acc={best_cc_acc:.4f}")
+print(f"  Repo: https://huggingface.co/{HF_REPO}")