azure home scripts: data gen, training, misc

close_the_loop.py

@@ -0,0 +1,703 @@
+#!/usr/bin/env python3
+"""
+=============================================================================
+CLOSING THE AI-BRAIN LOOP
+Using TRIBE v2 to identify architectural gaps between AI models and the brain
+=============================================================================
+
+Methodology:
+  Phase 1 — Load TRIBE v2, run inference, capture per-layer AI features
+  Phase 2 — Layer-wise encoding analysis: which AI layers predict which brain regions
+  Phase 3 — Modality ablation: which encoder drives which brain area
+  Phase 4 — RSA: representational similarity between AI layers and brain ROIs
+  Phase 5 — Divergence mapping: where the brain does something AI can't capture
+  Phase 6 — Architectural implications: what's missing in current AI
+
+Output: /home/azureuser/loop_results/
+"""
+
+import os
+import sys
+import logging
+import warnings
+import time
+
+import numpy as np
+import torch
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+from matplotlib.gridspec import GridSpec
+import pandas as pd
+from pathlib import Path
+from scipy import stats
+from scipy.spatial.distance import pdist, squareform
+from einops import rearrange
+
+warnings.filterwarnings("ignore")
+logging.basicConfig(level=logging.INFO, format="[%(asctime)s] %(message)s", datefmt="%H:%M:%S")
+log = logging.getLogger(__name__)
+
+OUT = Path("/home/azureuser/loop_results")
+OUT.mkdir(exist_ok=True)
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# PHASE 1: Load model, run inference, capture intermediate representations
+# ═══════════════════════════════════════════════════════════════════════════════
+
+log.info("PHASE 1: Loading TRIBE v2 and running inference")
+
+from tribev2 import TribeModel
+
+CACHE = "/home/azureuser/cache"
+VIDEO = "/home/azureuser/test_stimulus.mp4"
+
+model = TribeModel.from_pretrained("facebook/tribev2", cache_folder=CACHE)
+fmri_model = model._model
+device = fmri_model.device
+log.info(f"Model on {device}, n_outputs={fmri_model.n_outputs}")
+
+# Print model structure
+log.info("Model structure:")
+log.info(f"  Projectors: {list(fmri_model.projectors.keys())}")
+log.info(f"  Hidden dim: {fmri_model.config.hidden}")
+log.info(f"  Layer aggregation: {fmri_model.config.layer_aggregation}")
+log.info(f"  Extractor aggregation: {fmri_model.config.extractor_aggregation}")
+if hasattr(fmri_model, 'encoder'):
+    log.info(f"  Encoder: {type(fmri_model.encoder).__name__}")
+if hasattr(fmri_model, 'low_rank_head'):
+    log.info(f"  Low-rank head: {fmri_model.low_rank_head}")
+
+# Build events
+log.info(f"Processing video: {VIDEO}")
+events = model.get_events_dataframe(video_path=VIDEO)
+log.info(f"Events: {len(events)} rows, types: {events.type.unique().tolist()}")
+
+# Get data loader
+loader = model.data.get_loaders(events=events, split_to_build="all")["all"]
+
+# ── Collect per-layer features and brain predictions ──
+# We intercept batch.data to get the raw per-layer features (B, L, D, T)
+# and run the model to get brain predictions
+
+all_features = {}  # modality -> list of (L, D, T) arrays
+all_projected = {}  # modality -> list of (T, H) arrays
+all_post_encoder = []  # list of (T, H) arrays
+all_brain_preds = []  # list of (T, V) arrays — V = n_vertices
+
+# Register hooks on projectors to capture projected features
+proj_captures = {}
+proj_hooks = []
+
+def make_proj_hook(name):
+    def hook(mod, inp, out):
+        proj_captures[name] = out.detach().cpu().numpy()
+    return hook
+
+for mod_name, proj in fmri_model.projectors.items():
+    proj_hooks.append(proj.register_forward_hook(make_proj_hook(mod_name)))
+
+# Hook encoder output
+encoder_capture = [None]
+def enc_hook(mod, inp, out):
+    encoder_capture[0] = out.detach().cpu().numpy()
+if hasattr(fmri_model, 'encoder') and fmri_model.encoder is not None:
+    proj_hooks.append(fmri_model.encoder.register_forward_hook(enc_hook))
+
+log.info("Running inference with hooks...")
+t0 = time.time()
+
+with torch.inference_mode():
+    for batch_idx, batch in enumerate(loader):
+        batch = batch.to(device)
+
+        # Capture raw per-layer features before the model touches them
+        for mod_name in fmri_model.projectors.keys():
+            if mod_name in batch.data:
+                feat = batch.data[mod_name].detach().cpu().numpy()  # (B, L, D, T) or (B, D, T)
+                if feat.ndim == 3:
+                    feat = feat[:, np.newaxis, :, :]  # ensure 4D
+                if mod_name not in all_features:
+                    all_features[mod_name] = []
+                all_features[mod_name].append(feat)
+
+        # Forward pass (triggers hooks)
+        y_pred = fmri_model(batch).detach().cpu().numpy()  # (B, V, T')
+        y_pred = rearrange(y_pred, 'b v t -> (b t) v')
+        all_brain_preds.append(y_pred)
+
+        # Save projected features
+        for mod_name in proj_captures:
+            if mod_name not in all_projected:
+                all_projected[mod_name] = []
+            all_projected[mod_name].append(proj_captures[mod_name])
+        proj_captures.clear()
+
+        # Save encoder output
+        if encoder_capture[0] is not None:
+            all_post_encoder.append(encoder_capture[0])
+            encoder_capture[0] = None
+
+# Clean hooks
+for h in proj_hooks:
+    h.remove()
+
+elapsed = time.time() - t0
+log.info(f"Inference done in {elapsed:.1f}s")
+
+# Concatenate results
+brain_preds = np.concatenate(all_brain_preds, axis=0)  # (T_total, V)
+log.info(f"Brain predictions: {brain_preds.shape}")
+
+for mod in all_features:
+    all_features[mod] = np.concatenate(all_features[mod], axis=0)  # (B_total, L, D, T)
+    log.info(f"Raw features [{mod}]: {all_features[mod].shape}")
+
+for mod in all_projected:
+    all_projected[mod] = np.concatenate(all_projected[mod], axis=0)  # (B_total, T, H)
+    log.info(f"Projected features [{mod}]: {all_projected[mod].shape}")
+
+if all_post_encoder:
+    post_encoder = np.concatenate(all_post_encoder, axis=0)  # (B_total, T, H)
+    log.info(f"Post-encoder features: {post_encoder.shape}")
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# PHASE 2: Brain parcellation — map vertices to functional regions
+# ═══════════════════════════════════════════════════════════════════════════════
+
+log.info("PHASE 2: Loading brain parcellation (Destrieux atlas, fsaverage5)")
+
+from nilearn import datasets
+from nilearn.surface import load_surf_data
+
+fsaverage5 = datasets.fetch_surf_fsaverage("fsaverage5")
+
+# Load Destrieux parcellation labels
+labels_lh = load_surf_data(fsaverage5["annot_left_destrieux"])
+labels_rh = load_surf_data(fsaverage5["annot_right_destrieux"])
+
+N_VERT = 10242  # vertices per hemisphere in fsaverage5
+all_labels = np.concatenate([labels_lh, labels_rh])  # (20484,)
+
+# Destrieux atlas label names
+destrieux = datasets.fetch_atlas_destrieux_2009()
+label_names_raw = destrieux["labels"]
+label_names = {}
+for i, name in enumerate(label_names_raw):
+    if isinstance(name, bytes):
+        name = name.decode("utf-8")
+    label_names[i] = name
+
+# Build region info
+regions = {}
+for lid in np.unique(all_labels):
+    mask = all_labels == lid
+    n = mask.sum()
+    if n < 5:
+        continue
+    name = label_names.get(int(lid), f"region_{lid}")
+    if name == "Unknown" or name == "Medial_wall":
+        continue
+    regions[int(lid)] = {"name": name, "mask": mask, "n_vertices": int(n)}
+
+log.info(f"Found {len(regions)} usable brain regions")
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# PHASE 3: Layer-wise encoding analysis
+# ═══════════════════════════════════════════════════════════════════════════════
+
+log.info("PHASE 3: Layer-wise encoding analysis")
+
+# For each modality and each cached layer, compute how well that layer's
+# features correlate with predicted brain activity in each region.
+#
+# Method: For each region, average brain predictions across vertices to get
+# a time series. For each AI layer, average features across the feature dim
+# to get a time series. Compute Pearson correlation.
+
+# First, build time-aligned representations
+# brain_preds is (T_total, V) — each row is one TR's brain prediction
+T_total = brain_preds.shape[0]
+V = brain_preds.shape[1]
+
+# Per-region mean brain activity over time
+region_timeseries = {}
+for lid, rinfo in regions.items():
+    region_timeseries[lid] = brain_preds[:, rinfo["mask"]].mean(axis=1)  # (T_total,)
+
+# Per-modality, per-layer feature time series
+# all_features[mod] is (B, L, D, T_batch) — we need to flatten B and T_batch
+layer_timeseries = {}
+for mod, feats in all_features.items():
+    B, L, D, T_batch = feats.shape
+    # Reshape to (B*T_batch, L, D)
+    feats_flat = rearrange(feats, 'b l d t -> (b t) l d')
+    # Trim to match brain_preds length
+    min_t = min(feats_flat.shape[0], T_total)
+    feats_flat = feats_flat[:min_t]
+
+    layer_timeseries[mod] = {}
+    for l in range(L):
+        # Mean across feature dim to get single time series per layer
+        layer_timeseries[mod][l] = feats_flat[:, l, :].mean(axis=1)  # (T,)
+
+log.info("Computing layer-brain correlations...")
+
+# Correlation matrix: (modality, layer) x (region)
+layer_brain_corr = {}
+for mod in layer_timeseries:
+    L = len(layer_timeseries[mod])
+    for l in range(L):
+        key = f"{mod}_L{l}"
+        layer_brain_corr[key] = {}
+        layer_ts = layer_timeseries[mod][l]
+        min_t = min(len(layer_ts), T_total)
+        for lid, rinfo in regions.items():
+            brain_ts = region_timeseries[lid][:min_t]
+            lt = layer_ts[:min_t]
+            if np.std(lt) < 1e-10 or np.std(brain_ts) < 1e-10:
+                r = 0.0
+            else:
+                r, _ = stats.pearsonr(lt, brain_ts)
+            layer_brain_corr[key][lid] = r
+
+# Build matrix for visualization
+all_layer_keys = sorted(layer_brain_corr.keys())
+all_region_ids = sorted(regions.keys())
+region_names_list = [regions[lid]["name"] for lid in all_region_ids]
+
+corr_matrix = np.zeros((len(all_layer_keys), len(all_region_ids)))
+for i, lk in enumerate(all_layer_keys):
+    for j, lid in enumerate(all_region_ids):
+        corr_matrix[i, j] = layer_brain_corr[lk].get(lid, 0)
+
+log.info(f"Correlation matrix: {corr_matrix.shape} (AI layers x brain regions)")
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# PHASE 4: Modality ablation — which encoder drives which brain area
+# ═══════════════════════════════════════════════════════════════════════════════
+
+log.info("PHASE 4: Modality ablation analysis")
+
+modalities = list(fmri_model.projectors.keys())
+log.info(f"Ablating modalities: {modalities}")
+
+# Re-run inference with each modality zeroed out
+ablation_preds = {"full": brain_preds}
+
+with torch.inference_mode():
+    for mod_to_ablate in modalities:
+        log.info(f"  Ablating: {mod_to_ablate}")
+        preds_list = []
+        loader = model.data.get_loaders(events=events, split_to_build="all")["all"]
+        for batch in loader:
+            batch = batch.to(device)
+            # Zero out target modality
+            if mod_to_ablate in batch.data:
+                original = batch.data[mod_to_ablate].clone()
+                batch.data[mod_to_ablate] = torch.zeros_like(original)
+                y = fmri_model(batch).detach().cpu().numpy()
+                y = rearrange(y, 'b v t -> (b t) v')
+                preds_list.append(y)
+                batch.data[mod_to_ablate] = original
+            else:
+                y = fmri_model(batch).detach().cpu().numpy()
+                y = rearrange(y, 'b v t -> (b t) v')
+                preds_list.append(y)
+        ablation_preds[mod_to_ablate] = np.concatenate(preds_list, axis=0)
+        log.info(f"    shape: {ablation_preds[mod_to_ablate].shape}")
+
+# Compute per-region modality importance
+# Importance = mean absolute change when modality is removed
+region_mod_importance = {}
+for lid, rinfo in regions.items():
+    mask = rinfo["mask"]
+    full = ablation_preds["full"][:, mask]
+    imp = {}
+    for mod in modalities:
+        ablated = ablation_preds[mod][:, mask]
+        # Use both MSE change and correlation change
+        delta_mse = np.mean((full - ablated) ** 2)
+        imp[mod] = float(delta_mse)
+    total = sum(imp.values()) + 1e-12
+    imp_norm = {k: v / total for k, v in imp.items()}
+    region_mod_importance[lid] = imp_norm
+
+log.info("Modality ablation done")
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# PHASE 5: RSA — Representational Similarity Analysis
+# ═══════════════════════════════════════════════════════════════════════════════
+
+log.info("PHASE 5: Representational Similarity Analysis")
+
+# Split predictions into temporal segments (2-second windows)
+# and compute RDMs for AI features and brain regions
+
+SEGMENT_SIZE = 2  # TRs per segment (adjustable)
+n_segments = T_total // SEGMENT_SIZE
+
+# Build segment-level representations
+def build_segments(timeseries, n_segments, seg_size):
+    """Average timeseries within segments."""
+    segs = []
+    for i in range(n_segments):
+        start = i * seg_size
+        end = start + seg_size
+        if end <= len(timeseries):
+            segs.append(timeseries[start:end].mean(axis=0) if timeseries.ndim > 1 else timeseries[start:end].mean())
+    return np.array(segs)
+
+# Brain RDMs per region
+log.info("Building brain RDMs per region...")
+brain_rdms = {}
+for lid, rinfo in regions.items():
+    region_data = brain_preds[:, rinfo["mask"]]  # (T, n_verts)
+    seg_data = []
+    for i in range(n_segments):
+        s, e = i * SEGMENT_SIZE, (i + 1) * SEGMENT_SIZE
+        if e <= region_data.shape[0]:
+            seg_data.append(region_data[s:e].mean(axis=0))
+    if len(seg_data) < 3:
+        continue
+    seg_data = np.array(seg_data)  # (n_segments, n_verts)
+    if seg_data.std() < 1e-10:
+        brain_rdms[lid] = np.zeros((len(seg_data), len(seg_data)))
+    else:
+        brain_rdms[lid] = squareform(pdist(seg_data, metric="correlation"))
+
+# AI feature RDMs per modality per layer
+log.info("Building AI feature RDMs per layer...")
+ai_rdms = {}
+for mod, feats in all_features.items():
+    B, L, D, T_batch = feats.shape
+    feats_flat = rearrange(feats, 'b l d t -> (b t) l d')
+    min_t = min(feats_flat.shape[0], T_total)
+    feats_flat = feats_flat[:min_t]
+
+    for l in range(L):
+        layer_data = feats_flat[:, l, :]  # (T, D)
+        seg_data = []
+        for i in range(n_segments):
+            s, e = i * SEGMENT_SIZE, (i + 1) * SEGMENT_SIZE
+            if e <= layer_data.shape[0]:
+                seg_data.append(layer_data[s:e].mean(axis=0))
+        if len(seg_data) < 3:
+            continue
+        seg_data = np.array(seg_data)
+        key = f"{mod}_L{l}"
+        if seg_data.std() < 1e-10:
+            ai_rdms[key] = np.zeros((len(seg_data), len(seg_data)))
+        else:
+            ai_rdms[key] = squareform(pdist(seg_data, metric="correlation"))
+
+# Compare AI RDMs to brain RDMs using Spearman correlation
+log.info("Computing RSA (Spearman correlation between RDMs)...")
+rsa_matrix = np.zeros((len(ai_rdms), len(brain_rdms)))
+ai_rdm_keys = sorted(ai_rdms.keys())
+brain_rdm_keys = sorted(brain_rdms.keys())
+
+for i, ak in enumerate(ai_rdm_keys):
+    ai_vec = squareform(ai_rdms[ak])  # upper triangle
+    if len(ai_vec) == 0:
+        continue
+    for j, bk in enumerate(brain_rdm_keys):
+        brain_vec = squareform(brain_rdms[bk])
+        min_len = min(len(ai_vec), len(brain_vec))
+        if min_len < 3:
+            continue
+        rho, _ = stats.spearmanr(ai_vec[:min_len], brain_vec[:min_len])
+        rsa_matrix[i, j] = rho if not np.isnan(rho) else 0
+
+log.info(f"RSA matrix: {rsa_matrix.shape}")
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# PHASE 6: Divergence identification
+# ═══════════════════════════════════════════════════════════════════════════════
+
+log.info("PHASE 6: Identifying AI-brain divergences")
+
+results = []
+for idx_j, lid in enumerate(brain_rdm_keys):
+    rinfo = regions[lid]
+    imp = region_mod_importance.get(lid, {})
+
+    # Best RSA alignment across all AI layers
+    rsa_col = rsa_matrix[:, idx_j]
+    best_rsa = np.max(np.abs(rsa_col)) if len(rsa_col) > 0 else 0
+    best_rsa_layer = ai_rdm_keys[np.argmax(np.abs(rsa_col))] if len(rsa_col) > 0 else "none"
+
+    # Best encoding correlation across all AI layers
+    region_corrs = corr_matrix[:, all_region_ids.index(lid)] if lid in all_region_ids else np.zeros(1)
+    best_encoding = np.max(np.abs(region_corrs))
+    best_enc_layer = all_layer_keys[np.argmax(np.abs(region_corrs))] if len(region_corrs) > 0 else "none"
+
+    # Modality entropy
+    probs = np.array(list(imp.values())) + 1e-10
+    probs = probs / probs.sum()
+    entropy = -np.sum(probs * np.log2(probs))
+
+    # Temporal dynamics
+    temporal_var = float(brain_preds[:, rinfo["mask"]].mean(axis=1).var())
+
+    # Divergence score: high temporal dynamics but poor AI alignment
+    divergence = temporal_var * (1 - best_rsa) * entropy
+
+    results.append({
+        "region_id": lid,
+        "region": rinfo["name"],
+        "hemisphere": "LH" if lid < 100 else "RH",  # approximate
+        "n_vertices": rinfo["n_vertices"],
+        "temporal_variance": temporal_var,
+        "best_rsa_alignment": best_rsa,
+        "best_rsa_layer": best_rsa_layer,
+        "best_encoding_corr": best_encoding,
+        "best_encoding_layer": best_enc_layer,
+        "modality_entropy": entropy,
+        "divergence_score": divergence,
+        **{f"imp_{k}": v for k, v in imp.items()},
+    })
+
+df = pd.DataFrame(results)
+df = df.sort_values("divergence_score", ascending=False)
+df.to_csv(OUT / "full_analysis.csv", index=False)
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# PHASE 7: Visualization and reporting
+# ═══════════════════════════════════════════════════════════════════════════════
+
+log.info("PHASE 7: Generating visualizations and report")
+
+# ── Plot 1: Layer-brain correlation heatmap ──
+fig, ax = plt.subplots(figsize=(20, max(8, len(all_layer_keys) * 0.3)))
+im = ax.imshow(corr_matrix, aspect="auto", cmap="RdBu_r", vmin=-1, vmax=1)
+ax.set_yticks(range(len(all_layer_keys)))
+ax.set_yticklabels(all_layer_keys, fontsize=6)
+ax.set_xticks(range(len(region_names_list)))
+ax.set_xticklabels(region_names_list, fontsize=4, rotation=90)
+ax.set_title("Layer-wise Encoding: AI Layer ↔ Brain Region Correlation", fontsize=14)
+ax.set_ylabel("AI Encoder Layer")
+ax.set_xlabel("Brain Region (Destrieux)")
+plt.colorbar(im, ax=ax, label="Pearson r")
+fig.tight_layout()
+fig.savefig(OUT / "01_layer_brain_correlation.png", dpi=200)
+plt.close(fig)
+log.info("Saved 01_layer_brain_correlation.png")
+
+# ── Plot 2: RSA heatmap ──
+fig, ax = plt.subplots(figsize=(20, max(8, len(ai_rdm_keys) * 0.3)))
+im = ax.imshow(rsa_matrix, aspect="auto", cmap="RdBu_r", vmin=-0.5, vmax=0.5)
+ax.set_yticks(range(len(ai_rdm_keys)))
+ax.set_yticklabels(ai_rdm_keys, fontsize=6)
+brain_region_names_rsa = [regions[lid]["name"] for lid in brain_rdm_keys]
+ax.set_xticks(range(len(brain_region_names_rsa)))
+ax.set_xticklabels(brain_region_names_rsa, fontsize=4, rotation=90)
+ax.set_title("RSA: AI Layer ↔ Brain Region Representational Similarity", fontsize=14)
+ax.set_ylabel("AI Encoder Layer")
+ax.set_xlabel("Brain Region (Destrieux)")
+plt.colorbar(im, ax=ax, label="Spearman ρ")
+fig.tight_layout()
+fig.savefig(OUT / "02_rsa_heatmap.png", dpi=200)
+plt.close(fig)
+log.info("Saved 02_rsa_heatmap.png")
+
+# ── Plot 3: Divergence scatter ──
+fig, ax = plt.subplots(figsize=(12, 10))
+sc = ax.scatter(
+    df["best_rsa_alignment"],
+    df["temporal_variance"],
+    s=df["n_vertices"] * 0.5,
+    c=df["divergence_score"],
+    cmap="YlOrRd",
+    alpha=0.7,
+    edgecolors="k",
+    linewidths=0.3,
+)
+# Label top divergent
+for _, row in df.head(12).iterrows():
+    ax.annotate(
+        row["region"],
+        (row["best_rsa_alignment"], row["temporal_variance"]),
+        fontsize=6, alpha=0.9,
+        arrowprops=dict(arrowstyle="-", alpha=0.3),
+        textcoords="offset points", xytext=(5, 5),
+    )
+ax.set_xlabel("Best RSA Alignment (max |Spearman ρ| across AI layers)", fontsize=11)
+ax.set_ylabel("Temporal Variance (brain dynamics)", fontsize=11)
+ax.set_title("AI-Brain Divergence Map\nBottom-right = active brain regions poorly captured by AI", fontsize=13)
+plt.colorbar(sc, label="Divergence Score")
+fig.tight_layout()
+fig.savefig(OUT / "03_divergence_scatter.png", dpi=200)
+plt.close(fig)
+log.info("Saved 03_divergence_scatter.png")
+
+# ── Plot 4: Modality importance per region (stacked bar) ──
+imp_cols = [c for c in df.columns if c.startswith("imp_")]
+if imp_cols:
+    top_30 = df.nlargest(30, "temporal_variance")
+    fig, ax = plt.subplots(figsize=(14, 8))
+    bottom = np.zeros(len(top_30))
+    colors = plt.cm.Set2(np.linspace(0, 1, len(imp_cols)))
+    for ci, col in enumerate(imp_cols):
+        vals = top_30[col].values
+        ax.barh(range(len(top_30)), vals, left=bottom, color=colors[ci],
+                label=col.replace("imp_", "").upper())
+        bottom += vals
+    ax.set_yticks(range(len(top_30)))
+    ax.set_yticklabels(top_30["region"].values, fontsize=7)
+    ax.set_xlabel("Relative Modality Importance")
+    ax.set_title("Modality Contribution per Brain Region (top 30 by dynamics)", fontsize=13)
+    ax.legend(loc="lower right")
+    fig.tight_layout()
+    fig.savefig(OUT / "04_modality_importance.png", dpi=200)
+    plt.close(fig)
+    log.info("Saved 04_modality_importance.png")
+
+# ── Plot 5: Brain surface divergence map ──
+try:
+    from nilearn.plotting import plot_surf_stat_map
+
+    vertex_divergence = np.zeros(N_VERT * 2)
+    for _, row in df.iterrows():
+        lid = row["region_id"]
+        if lid in regions:
+            vertex_divergence[regions[lid]["mask"]] = row["divergence_score"]
+
+    fig = plt.figure(figsize=(16, 12))
+    for idx, (hemi, view) in enumerate([
+        ("left", "lateral"), ("left", "medial"),
+        ("right", "lateral"), ("right", "medial")
+    ]):
+        ax = fig.add_subplot(2, 2, idx + 1, projection="3d")
+        if hemi == "left":
+            data = vertex_divergence[:N_VERT]
+        else:
+            data = vertex_divergence[N_VERT:]
+        plot_surf_stat_map(
+            fsaverage5[f"pial_{hemi}"],
+            data,
+            hemi=hemi,
+            view=view,
+            cmap="YlOrRd",
+            title=f"Divergence ({hemi} {view})",
+            figure=fig,
+            axes=ax,
+        )
+    fig.suptitle("Brain Surface: AI-Brain Divergence Scores", fontsize=14, y=1.02)
+    fig.tight_layout()
+    fig.savefig(OUT / "05_brain_surface_divergence.png", dpi=200, bbox_inches="tight")
+    plt.close(fig)
+    log.info("Saved 05_brain_surface_divergence.png")
+except Exception as e:
+    log.warning(f"Brain surface plot failed: {e}")
+
+# ── Plot 6: Per-modality best layer alignment profile ──
+for mod in all_features:
+    mod_keys = [k for k in all_layer_keys if k.startswith(f"{mod}_")]
+    if not mod_keys:
+        continue
+    mod_indices = [all_layer_keys.index(k) for k in mod_keys]
+    mod_corr = corr_matrix[mod_indices, :]  # (n_layers_mod, n_regions)
+
+    fig, ax = plt.subplots(figsize=(14, 6))
+    im = ax.imshow(mod_corr, aspect="auto", cmap="RdBu_r", vmin=-1, vmax=1)
+    ax.set_yticks(range(len(mod_keys)))
+    ax.set_yticklabels([f"Layer {i}" for i in range(len(mod_keys))], fontsize=8)
+    ax.set_xticks(range(len(region_names_list)))
+    ax.set_xticklabels(region_names_list, fontsize=4, rotation=90)
+    ax.set_title(f"{mod.upper()} Encoder: Layer-wise Brain Alignment", fontsize=13)
+    ax.set_ylabel("Encoder Layer (early → late)")
+    ax.set_xlabel("Brain Region")
+    plt.colorbar(im, ax=ax, label="Pearson r")
+    fig.tight_layout()
+    fig.savefig(OUT / f"06_{mod}_layer_alignment.png", dpi=200)
+    plt.close(fig)
+    log.info(f"Saved 06_{mod}_layer_alignment.png")
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# PHASE 8: Final report
+# ═══════════════════════════════════════════════════════════════════════════════
+
+report = []
+report.append("=" * 100)
+report.append("CLOSING THE AI-BRAIN LOOP: Analysis Report")
+report.append(f"Generated: {pd.Timestamp.now()}")
+report.append("=" * 100)
+
+report.append("\n\n--- DATASET ---")
+report.append(f"Stimulus: {VIDEO}")
+report.append(f"Total time points: {T_total}")
+report.append(f"Brain vertices: {V} (fsaverage5)")
+report.append(f"Brain regions analyzed: {len(regions)} (Destrieux atlas)")
+report.append(f"AI modalities: {modalities}")
+for mod, feats in all_features.items():
+    report.append(f"  {mod}: {feats.shape[1]} layers, {feats.shape[2]}-dim features")
+
+report.append("\n\n--- TOP 15 DIVERGENCE REGIONS ---")
+report.append("(Brain regions with high dynamics but poor AI alignment)")
+report.append("")
+cols = ["region", "temporal_variance", "best_rsa_alignment", "best_rsa_layer",
+        "modality_entropy", "divergence_score"]
+cols_present = [c for c in cols if c in df.columns]
+report.append(df[cols_present].head(15).to_string(index=False, float_format="%.4f"))
+
+report.append("\n\n--- TOP 15 WELL-ALIGNED REGIONS ---")
+report.append("(Brain regions where AI encoders match brain representations well)")
+well_aligned = df.nlargest(15, "best_rsa_alignment")
+report.append(well_aligned[cols_present].to_string(index=False, float_format="%.4f"))
+
+report.append("\n\n--- MODALITY DOMINANCE ---")
+for mod in modalities:
+    col = f"imp_{mod}"
+    if col not in df.columns:
+        continue
+    report.append(f"\n{mod.upper()}-dominated regions (top 5):")
+    top = df.nlargest(5, col)
+    for _, row in top.iterrows():
+        report.append(f"  {row['region']:45s} importance={row[col]:.4f}  rsa={row['best_rsa_alignment']:.4f}")
+
+report.append("\n\n--- ARCHITECTURAL IMPLICATIONS ---")
+report.append("""
+Based on the divergence analysis, here are the gaps in current AI architectures
+and proposed solutions:
+
+1. HIGH-ENTROPY DIVERGENCE REGIONS (multiple modalities contribute equally,
+   but overall alignment is poor):
+   → The brain performs CROSS-MODAL INTEGRATION that the concatenation-based
+     fusion in TRIBE v2 (and most multimodal AI) doesn't capture.
+   → Proposed fix: EARLY FUSION with cross-attention between modality streams
+     at intermediate layers, not just late concatenation.
+
+2. HIGH TEMPORAL VARIANCE + LOW ALIGNMENT:
+   → The brain has strong TEMPORAL DYNAMICS (prediction, memory, feedback loops)
+     that feedforward AI encoders miss entirely.
+   → Proposed fix: Add RECURRENT connections or PREDICTIVE CODING layers that
+     generate top-down predictions and propagate prediction errors.
+
+3. REGIONS WHERE NO SINGLE LAYER ALIGNS WELL:
+   → The brain's computation in these areas may involve representations that
+     DON'T EXIST in any layer of V-JEPA2, LLaMA, or Wav2Vec-BERT.
+   → Proposed fix: Train a NEW encoder objective that explicitly optimizes for
+     brain alignment in these gap regions (brain-guided contrastive learning).
+
+4. LAYER DEPTH PATTERNS:
+   → If early AI layers align with sensory cortex and late layers align with
+     association cortex, this confirms the HIERARCHICAL CORRESPONDENCE between
+     DNNs and the cortical hierarchy.
+   → Where this breaks (e.g., late layers don't align with prefrontal cortex),
+     it suggests the model lacks EXECUTIVE/ABSTRACT processing.
+""")
+
+report.append("\n--- FILES ---")
+report.append(f"Full CSV: {OUT / 'full_analysis.csv'}")
+report.append(f"Plots: {OUT / '*.png'}")
+report.append("=" * 100)
+
+report_text = "\n".join(report)
+print(report_text)
+
+with open(OUT / "report.txt", "w") as f:
+    f.write(report_text)
+
+log.info(f"All results saved to {OUT}")
+log.info("Done.")

generate_dataset.py

@@ -0,0 +1,160 @@
+#!/usr/bin/env python3
+"""Generate a diverse character dataset using Flux 1 + PuLID for LoRA training."""
+
+import json
+import urllib.request
+import time
+import os
+
+COMFYUI_URL = "http://127.0.0.1:80"
+DATASET_DIR = "/home/azureuser/ai-toolkit/character_dataset"
+TRIGGER_WORD = "ohwx"
+
+# Diverse prompts covering different angles, scenes, lighting, outfits
+PROMPTS = [
+    # Close-ups / headshots
+    (f"close up portrait photo of {TRIGGER_WORD} woman, natural lighting, soft smile, looking at camera, shallow depth of field", "closeup_front"),
+    (f"close up portrait of {TRIGGER_WORD} woman, side profile view, golden hour lighting, outdoors", "closeup_profile"),
+    (f"close up portrait of {TRIGGER_WORD} woman, three quarter view, studio lighting, neutral expression", "closeup_34"),
+    (f"headshot of {TRIGGER_WORD} woman, looking slightly up, soft natural light, gentle smile", "closeup_up"),
+    (f"close up of {TRIGGER_WORD} woman, looking down at something in her hands, natural indoor lighting", "closeup_down"),
+
+    # Half body shots
+    (f"half body photo of {TRIGGER_WORD} woman in a white blouse, sitting at a cafe table, warm ambient light, candid", "half_cafe"),
+    (f"half body photo of {TRIGGER_WORD} woman wearing a leather jacket, urban street background, overcast day", "half_street"),
+    (f"half body photo of {TRIGGER_WORD} woman in athletic wear, gym setting, bright lighting", "half_gym"),
+    (f"half body photo of {TRIGGER_WORD} woman in a sundress, garden background, dappled sunlight", "half_garden"),
+    (f"half body photo of {TRIGGER_WORD} woman in business attire, modern office, professional lighting", "half_office"),
+
+    # Full body shots
+    (f"full body photo of {TRIGGER_WORD} woman walking on the beach, sunset, casual summer outfit, warm tones", "full_beach"),
+    (f"full body photo of {TRIGGER_WORD} woman standing in a city street, winter coat, evening city lights", "full_city"),
+    (f"full body photo of {TRIGGER_WORD} woman hiking on a mountain trail, athletic outfit, golden hour", "full_hike"),
+    (f"full body photo of {TRIGGER_WORD} woman leaning against a wall, casual jeans and t-shirt, natural daylight", "full_casual"),
+
+    # Different lighting conditions
+    (f"portrait of {TRIGGER_WORD} woman, dramatic side lighting, dark moody background, artistic photo", "light_dramatic"),
+    (f"portrait of {TRIGGER_WORD} woman, bright overcast natural light, white background, clean look", "light_bright"),
+    (f"portrait of {TRIGGER_WORD} woman, warm golden hour backlight, hair glowing, outdoor", "light_golden"),
+    (f"portrait of {TRIGGER_WORD} woman, soft window light, sitting on a couch, cozy indoor setting", "light_window"),
+
+    # Different expressions
+    (f"photo of {TRIGGER_WORD} woman laughing genuinely, candid moment, natural setting", "expr_laugh"),
+    (f"photo of {TRIGGER_WORD} woman with a serious contemplative expression, looking into distance", "expr_serious"),
+]
+
+# Corresponding captions (without trigger word for variety, will be added by trainer)
+CAPTIONS = {
+    "closeup_front": f"close up portrait photo of {TRIGGER_WORD} woman, natural lighting, soft smile, looking at camera",
+    "closeup_profile": f"close up portrait of {TRIGGER_WORD} woman, side profile view, golden hour lighting, outdoors",
+    "closeup_34": f"close up portrait of {TRIGGER_WORD} woman, three quarter view, studio lighting, neutral expression",
+    "closeup_up": f"headshot of {TRIGGER_WORD} woman, looking slightly up, soft natural light, gentle smile",
+    "closeup_down": f"close up of {TRIGGER_WORD} woman, looking down, natural indoor lighting",
+    "half_cafe": f"half body photo of {TRIGGER_WORD} woman in a white blouse, sitting at a cafe, warm ambient light",
+    "half_street": f"half body photo of {TRIGGER_WORD} woman wearing a leather jacket, urban street, overcast day",
+    "half_gym": f"half body photo of {TRIGGER_WORD} woman in athletic wear, gym setting, bright lighting",
+    "half_garden": f"half body photo of {TRIGGER_WORD} woman in a sundress, garden, dappled sunlight",
+    "half_office": f"half body photo of {TRIGGER_WORD} woman in business attire, modern office",
+    "full_beach": f"full body photo of {TRIGGER_WORD} woman walking on the beach, sunset, casual summer outfit",
+    "full_city": f"full body photo of {TRIGGER_WORD} woman standing in city street, winter coat, evening lights",
+    "full_hike": f"full body photo of {TRIGGER_WORD} woman hiking on mountain trail, athletic outfit, golden hour",
+    "full_casual": f"full body photo of {TRIGGER_WORD} woman leaning against wall, jeans and t-shirt, daylight",
+    "light_dramatic": f"portrait of {TRIGGER_WORD} woman, dramatic side lighting, dark moody background",
+    "light_bright": f"portrait of {TRIGGER_WORD} woman, bright overcast light, white background, clean",
+    "light_golden": f"portrait of {TRIGGER_WORD} woman, warm golden hour backlight, hair glowing, outdoor",
+    "light_window": f"portrait of {TRIGGER_WORD} woman, soft window light, sitting on couch, cozy indoor",
+    "expr_laugh": f"photo of {TRIGGER_WORD} woman laughing genuinely, candid moment, natural setting",
+    "expr_serious": f"photo of {TRIGGER_WORD} woman with serious contemplative expression, looking into distance",
+}
+
+
+def queue_prompt(prompt_text, filename, seed):
+    workflow = {
+        "1": {"class_type": "UNETLoader", "inputs": {"unet_name": "flux1-dev.safetensors", "weight_dtype": "default"}},
+        "2": {"class_type": "DualCLIPLoader", "inputs": {"clip_name1": "t5xxl_fp16.safetensors", "clip_name2": "clip_l.safetensors", "type": "flux"}},
+        "3": {"class_type": "VAELoader", "inputs": {"vae_name": "ae.safetensors"}},
+        "4": {"class_type": "PulidFluxModelLoader", "inputs": {"pulid_file": "pulid_flux_v0.9.1.safetensors"}},
+        "5": {"class_type": "PulidFluxInsightFaceLoader", "inputs": {"provider": "CUDA"}},
+        "6": {"class_type": "PulidFluxEvaClipLoader", "inputs": {}},
+        "7": {"class_type": "LoadImage", "inputs": {"image": "reference_face.png"}},
+        "8": {"class_type": "ApplyPulidFlux", "inputs": {
+            "model": ["1", 0], "pulid_flux": ["4", 0], "eva_clip": ["6", 0],
+            "face_analysis": ["5", 0], "image": ["7", 0],
+            "weight": 0.85, "start_at": 0.0, "end_at": 1.0
+        }},
+        "9": {"class_type": "CLIPTextEncodeFlux", "inputs": {
+            "clip": ["2", 0],
+            "clip_l": prompt_text[:77],
+            "t5xxl": prompt_text,
+            "guidance": 3.5
+        }},
+        "10": {"class_type": "EmptySD3LatentImage", "inputs": {"width": 1024, "height": 1024, "batch_size": 1}},
+        "11": {"class_type": "KSampler", "inputs": {
+            "model": ["8", 0], "positive": ["9", 0], "negative": ["9", 0],
+            "latent_image": ["10", 0], "seed": seed,
+            "control_after_generate": "fixed", "steps": 20, "cfg": 1.0,
+            "sampler_name": "euler", "scheduler": "simple", "denoise": 1.0
+        }},
+        "12": {"class_type": "VAEDecode", "inputs": {"samples": ["11", 0], "vae": ["3", 0]}},
+        "13": {"class_type": "SaveImage", "inputs": {"images": ["12", 0], "filename_prefix": f"dataset_{filename}"}}
+    }
+
+    data = json.dumps({"prompt": workflow}).encode()
+    req = urllib.request.Request(f'{COMFYUI_URL}/prompt', data=data, headers={'Content-Type': 'application/json'})
+    resp = urllib.request.urlopen(req)
+    return json.loads(resp.read())['prompt_id']
+
+
+def wait_for_completion(prompt_id, timeout=600):
+    start = time.time()
+    while time.time() - start < timeout:
+        req = urllib.request.Request(f'{COMFYUI_URL}/history/{prompt_id}')
+        resp = urllib.request.urlopen(req)
+        history = json.loads(resp.read())
+        if prompt_id in history:
+            h = history[prompt_id]
+            status = h.get('status', {}).get('status_str', '')
+            if status == 'success':
+                for nid, out in h['outputs'].items():
+                    if 'images' in out:
+                        return out['images'][0]['filename']
+            elif status == 'error':
+                msgs = h.get('status', {}).get('messages', [])
+                for m in msgs:
+                    if m[0] == 'execution_error':
+                        print(f"  ERROR: {m[1].get('exception_message', 'unknown')[:200]}")
+                return None
+        time.sleep(2)
+    return None
+
+
+def main():
+    print(f"Generating {len(PROMPTS)} character images for LoRA training...\n")
+
+    for i, (prompt, name) in enumerate(PROMPTS):
+        seed = 10000 + i * 1337
+        print(f"[{i+1}/{len(PROMPTS)}] {name} (seed={seed})")
+        prompt_id = queue_prompt(prompt, name, seed)
+        filename = wait_for_completion(prompt_id)
+
+        if filename:
+            # Copy to dataset folder
+            src = f"/home/azureuser/ComfyUI/output/{filename}"
+            dst = os.path.join(DATASET_DIR, f"{name}.png")
+            os.system(f"cp '{src}' '{dst}'")
+
+            # Write caption
+            caption = CAPTIONS[name]
+            with open(os.path.join(DATASET_DIR, f"{name}.txt"), 'w') as f:
+                f.write(caption)
+
+            print(f"  -> saved {name}.png + caption")
+        else:
+            print(f"  -> FAILED")
+
+    print(f"\nDone! Dataset at: {DATASET_DIR}")
+    print(f"Files: {len(os.listdir(DATASET_DIR))}")
+
+
+if __name__ == "__main__":
+    main()

generate_dataset_flux2.py

@@ -0,0 +1,132 @@
+#!/usr/bin/env python3
+"""Generate diverse Indian woman character dataset using Flux 2 Dev for LoRA training."""
+
+import json
+import urllib.request
+import time
+import os
+
+COMFYUI_URL = "http://127.0.0.1:80"
+DATASET_DIR = "/home/azureuser/ai-toolkit/character_dataset"
+TRIGGER = "ohwx"
+
+# Consistent identity description used in EVERY prompt
+IDENTITY = f"a young 21 year old Indian woman called {TRIGGER}, long straight black hair, brown eyes, light brown skin, soft facial features, natural beauty"
+
+# (prompt_suffix, filename, caption)
+SHOTS = [
+    # === ANGLES / HEAD POSITIONS ===
+    ("front facing portrait, looking directly at camera, neutral expression, white background, studio lighting, photorealistic", "angle_front", "front facing portrait, looking at camera, neutral expression, studio lighting"),
+    ("side profile portrait, left side of face visible, looking left, clean background, studio lighting, photorealistic", "angle_left_profile", "side profile, left side visible, looking left, studio lighting"),
+    ("side profile portrait, right side of face visible, looking right, clean background, studio lighting, photorealistic", "angle_right_profile", "side profile, right side visible, looking right, studio lighting"),
+    ("three quarter view portrait, slightly turned to the left, soft lighting, clean background, photorealistic", "angle_34_left", "three quarter view, slightly turned left, soft lighting"),
+    ("three quarter view portrait, slightly turned to the right, soft lighting, clean background, photorealistic", "angle_34_right", "three quarter view, slightly turned right, soft lighting"),
+    ("looking up, chin slightly raised, soft lighting from above, clean background, photorealistic portrait", "angle_looking_up", "looking up, chin raised, soft overhead lighting"),
+    ("looking down, reading a book, natural indoor lighting, soft focus background, candid photo", "angle_looking_down", "looking down reading a book, natural indoor lighting, candid"),
+    ("back view, looking over shoulder at camera, long black hair visible, outdoor setting, natural light", "angle_back", "back view looking over shoulder, long black hair visible, outdoor natural light"),
+    ("back of head view, long straight black hair falling down, standing outdoors, natural daylight", "angle_back_full", "back of head, long straight black hair, standing outdoors"),
+    ("tilted head portrait, slight head tilt to the right, gentle smile, soft lighting, photorealistic", "angle_tilt", "tilted head portrait, slight tilt right, gentle smile"),
+
+    # === EMOTIONS ===
+    ("smiling brightly, genuine happy smile showing teeth, warm expression, natural light, photorealistic portrait", "emo_happy", "smiling brightly, genuine happy smile, warm expression, natural light"),
+    ("sad expression, looking down slightly, melancholic mood, soft moody lighting, photorealistic portrait", "emo_sad", "sad expression, looking down, melancholic, soft moody lighting"),
+    ("laughing candidly, eyes slightly closed, natural joyful moment, outdoor sunlight, candid photo", "emo_laugh", "laughing candidly, joyful moment, outdoor sunlight, candid"),
+    ("serious contemplative expression, deep in thought, looking into distance, dramatic lighting, photorealistic", "emo_serious", "serious contemplative expression, deep in thought, dramatic lighting"),
+    ("surprised expression, eyes wide, mouth slightly open, bright lighting, photorealistic portrait", "emo_surprised", "surprised expression, eyes wide, bright lighting"),
+    ("shy smile, looking slightly away from camera, subtle smile, soft warm lighting, photorealistic", "emo_shy", "shy smile, looking slightly away, subtle smile, soft warm lighting"),
+    ("confident expression, strong gaze at camera, slight smirk, professional lighting, photorealistic", "emo_confident", "confident expression, strong gaze, slight smirk, professional lighting"),
+    ("peaceful serene expression, eyes closed, meditating, soft natural light, photorealistic portrait", "emo_peaceful", "peaceful serene expression, eyes closed, meditating, soft light"),
+
+    # === HALF BODY / DIFFERENT OUTFITS ===
+    ("half body photo, wearing traditional Indian salwar kameez, standing in a garden, natural sunlight, photorealistic", "outfit_salwar", "half body, wearing traditional Indian salwar kameez, garden, sunlight"),
+    ("half body photo, wearing modern casual jeans and white t-shirt, urban street background, daylight, photorealistic", "outfit_casual", "half body, casual jeans and white t-shirt, urban street, daylight"),
+    ("half body photo, wearing formal black blazer and white shirt, office setting, professional lighting, photorealistic", "outfit_formal", "half body, formal black blazer, office setting, professional lighting"),
+    ("half body photo, wearing a red saree with gold border, elegant pose, warm indoor lighting, photorealistic", "outfit_saree", "half body, wearing red saree with gold border, elegant pose, warm indoor light"),
+    ("half body photo, wearing athletic sportswear, gym background, bright fluorescent lighting, photorealistic", "outfit_sport", "half body, athletic sportswear, gym, bright lighting"),
+    ("half body photo, wearing a cozy sweater, sitting by window, rainy day outside, soft natural light, photorealistic", "outfit_cozy", "half body, cozy sweater, sitting by window, rainy day, soft light"),
+
+    # === FULL BODY / DIFFERENT SCENES ===
+    ("full body photo, walking on a beach at sunset, wearing summer dress, golden hour light, barefoot on sand, photorealistic", "full_beach", "full body, walking on beach at sunset, summer dress, golden hour"),
+    ("full body photo, standing in a city street, wearing winter jacket, evening city lights, urban photography, photorealistic", "full_city", "full body, city street, winter jacket, evening city lights"),
+    ("full body photo, sitting cross legged on grass in a park, casual outfit, dappled sunlight through trees, photorealistic", "full_park", "full body, sitting on grass in park, casual outfit, dappled sunlight"),
+    ("full body photo, dancing in the rain, white dress, joyful expression, dramatic wet look, photorealistic", "full_rain", "full body, dancing in rain, white dress, joyful, dramatic"),
+    ("full body photo, standing at a temple entrance, traditional Indian outfit, warm morning light, photorealistic", "full_temple", "full body, standing at temple entrance, traditional outfit, morning light"),
+
+    # === DIFFERENT LIGHTING ===
+    ("portrait, dramatic chiaroscuro lighting, half face in shadow, artistic, high contrast, photorealistic", "light_dramatic", "portrait, dramatic chiaroscuro lighting, half face in shadow, high contrast"),
+    ("portrait, soft golden hour backlight, hair glowing, warm tones, outdoor, lens flare, photorealistic", "light_golden", "portrait, golden hour backlight, hair glowing, warm tones, outdoor"),
+    ("portrait, bright overcast daylight, flat even lighting, outdoor, clean natural look, photorealistic", "light_overcast", "portrait, bright overcast daylight, even lighting, outdoor, natural"),
+    ("portrait, neon colored lights reflecting on face, nighttime, urban setting, cinematic, photorealistic", "light_neon", "portrait, neon lights on face, nighttime, urban, cinematic"),
+    ("portrait, soft warm candlelight, intimate setting, warm orange tones, photorealistic", "light_candle", "portrait, soft candlelight, intimate setting, warm orange tones"),
+]
+
+
+def queue_prompt(prompt_text, filename, seed):
+    workflow = {
+        "1": {"class_type": "UNETLoader", "inputs": {"unet_name": "flux2_dev_fp8mixed.safetensors", "weight_dtype": "default"}},
+        "2": {"class_type": "CLIPLoader", "inputs": {"clip_name": "mistral_3_small_flux2_bf16.safetensors", "type": "flux2"}},
+        "3": {"class_type": "VAELoader", "inputs": {"vae_name": "flux2-vae.safetensors"}},
+        "9": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["2", 0], "text": prompt_text}},
+        "10": {"class_type": "EmptyFlux2LatentImage", "inputs": {"width": 1024, "height": 1024, "batch_size": 1}},
+        "11": {"class_type": "KSampler", "inputs": {
+            "model": ["1", 0], "positive": ["9", 0], "negative": ["9", 0],
+            "latent_image": ["10", 0], "seed": seed,
+            "control_after_generate": "fixed", "steps": 20, "cfg": 1.0,
+            "sampler_name": "euler", "scheduler": "simple", "denoise": 1.0
+        }},
+        "12": {"class_type": "VAEDecode", "inputs": {"samples": ["11", 0], "vae": ["3", 0]}},
+        "13": {"class_type": "SaveImage", "inputs": {"images": ["12", 0], "filename_prefix": f"ds_{filename}"}}
+    }
+    data = json.dumps({"prompt": workflow}).encode()
+    req = urllib.request.Request(f'{COMFYUI_URL}/prompt', data=data, headers={'Content-Type': 'application/json'})
+    resp = urllib.request.urlopen(req)
+    return json.loads(resp.read())['prompt_id']
+
+
+def wait_for_completion(prompt_id, timeout=600):
+    start = time.time()
+    while time.time() - start < timeout:
+        req = urllib.request.Request(f'{COMFYUI_URL}/history/{prompt_id}')
+        resp = urllib.request.urlopen(req)
+        history = json.loads(resp.read())
+        if prompt_id in history:
+            h = history[prompt_id]
+            s = h.get('status', {}).get('status_str', '')
+            if s == 'success':
+                for out in h['outputs'].values():
+                    if 'images' in out:
+                        return out['images'][0]['filename']
+            elif s == 'error':
+                return None
+        time.sleep(2)
+    return None
+
+
+def main():
+    print(f"Generating {len(SHOTS)} images with Flux 2 Dev...\n")
+
+    for i, (suffix, name, caption) in enumerate(SHOTS):
+        seed = 42424 + i * 997
+        full_prompt = f"{IDENTITY}, {suffix}"
+        full_caption = f"photo of {IDENTITY}, {caption}"
+
+        print(f"[{i+1}/{len(SHOTS)}] {name}")
+        prompt_id = queue_prompt(full_prompt, name, seed)
+        filename = wait_for_completion(prompt_id)
+
+        if filename:
+            src = f"/home/azureuser/ComfyUI/output/{filename}"
+            dst = os.path.join(DATASET_DIR, f"{name}.png")
+            os.system(f"cp '{src}' '{dst}'")
+            with open(os.path.join(DATASET_DIR, f"{name}.txt"), 'w') as f:
+                f.write(full_caption)
+            print(f"  -> OK")
+        else:
+            print(f"  -> FAILED")
+
+    total = len([f for f in os.listdir(DATASET_DIR) if f.endswith('.png')])
+    print(f"\nDone! {total} images in {DATASET_DIR}")
+
+
+if __name__ == "__main__":
+    main()

gsoc_proposal_content.md

@@ -0,0 +1,277 @@
+# GSoC 2026 Proposal: LLM Support for 7B Models (OLMo) in DeepChem
+# CONTENT REFERENCE — REWRITE IN YOUR OWN WORDS BEFORE SUBMITTING
+
+---
+
+## 1. INTRODUCTION
+
+### What this project is about
+DeepChem's HuggingFaceModel wrapper currently supports encoder-only models
+(ChemBERTa, MoLFormer) through masked language modeling. There is no support
+for decoder-only causal language models. This project adds OLMo-2
+(Allen AI's open language model) to DeepChem, enabling:
+- Continued pretraining on molecular data (SMILES)
+- Fine-tuning for classification and regression
+- Autoregressive molecular generation
+
+### Why it matters
+Encoder models (ChemBERTa) can classify and predict properties, but they
+CANNOT generate molecules. A causal LM like OLMo opens up:
+- De novo molecular generation (drug discovery)
+- Text-molecule bridging (OLMo understands English AND can learn SMILES)
+- In-context few-shot learning without fine-tuning
+- Transfer learning from scientific literature
+
+### Why OLMo specifically
+- Fully open (weights, data, training code) — unlike LLaMA/GPT
+- OLMo-2 is natively supported in HuggingFace transformers (no custom code)
+- 1B and 7B variants available for different compute budgets
+- Trained on Dolma corpus which includes scientific papers
+
+### What I've already done (reference your PR and experiments)
+- Found and fixed a transformers 5.x compatibility bug in Chemberta (PR #4913)
+- Filed issue #4912 documenting broader transformers 5.x compat gap
+- Built a working OLMo wrapper prototype (locally) with:
+  - Olmo2ForSequenceClassification (doesn't exist in transformers)
+  - Causal LM pretraining on SMILES
+  - All 8 unit tests passing
+- Ran experiments on real MoleculeNet data:
+  - BBBP classification: ROC-AUC 0.67 (random init, tiny model)
+  - ESOL regression: R^2 = 0.37
+  - SMILES generation: 0% validity (expected — proves pretraining is essential)
+
+---
+
+## 2. RELEVANT EXPERIENCE & INTEREST
+
+### Technical background
+- Parameter Golf (OpenAI competition, March 2026): Trained language models
+  from scratch under 16MB constraint. Custom SentencePiece tokenizers,
+  GPTQ-lite quantization, flash attention, architecture design (11L 512d
+  transformer). This is directly relevant — I understand transformer
+  training at a low level.
+- GSPO-DeepSeek-R1-Distill-Qwen-1.5B (15 GitHub stars): Fine-tuning and
+  distillation of large language models.
+- wingman-AI (29 GitHub stars): Production AI assistant system.
+- Open source contributions: PRs to HuggingFace transformers, Unsloth,
+  Anthropic SDK, OpenAI SDK, Karpathy's nanochat.
+
+### Why I want to work on this
+[WRITE THIS YOURSELF — what genuinely interests you about molecular ML?
+Why DeepChem? Be specific and honest. Don't say "I'm passionate about
+open source" — say what specific thing drew you to this project.]
+
+### Links
+- GitHub: https://github.com/vivekvar-dl
+- PR #4913: https://github.com/deepchem/deepchem/pull/4913
+- Issue #4912: https://github.com/deepchem/deepchem/issues/4912
+
+---
+
+## 3. WORK PLAN
+
+### 3.1 Design
+
+The implementation has four components:
+
+**Component A: Base class changes to HuggingFaceModel**
+- Add `causal_lm` task support (DataCollatorForLanguageModeling with mlm=False)
+- Add `AutoModelForCausalLM` branch in load_from_pretrained()
+- Add `generate()` method for autoregressive text generation
+- Add causal LM batch preparation in _prepare_batch()
+
+Note: PR #4907 by another contributor adds a similar generate() method.
+My work is complementary — I'm adding a full model wrapper, not just
+generation plumbing.
+
+**Component B: Olmo2ForSequenceClassification**
+This class DOES NOT EXIST in HuggingFace transformers. OLMo only has
+OlmoForCausalLM — no classification head. I built one:
+- Extends Olmo2PreTrainedModel
+- Uses last-token pooling (last non-padded token's hidden state)
+- Linear projection head for classification/regression
+- Supports single-label, multi-label, and regression via problem_type config
+- Computes CrossEntropyLoss / BCEWithLogitsLoss / MSELoss based on task
+
+This follows the same pattern as LlamaForSequenceClassification.
+
+**Component C: OLMo wrapper class**
+```
+OLMo(HuggingFaceModel)
+  __init__(task, model_name, n_tasks, config)
+    - task: causal_lm | regression | classification | mtr
+    - Loads tokenizer from HuggingFace Hub
+    - Sets pad_token = eos_token (decoder models don't have pad by default)
+    - Syncs vocab_size between config and tokenizer
+    - Creates appropriate model class based on task
+
+  _prepare_batch(batch)
+    - causal_lm: labels = input_ids (model shifts internally)
+    - regression/classification: labels from dataset, proper dtype casting
+    - Multi-task classification: float labels for BCEWithLogitsLoss
+```
+
+**Component D: Tokenization strategy**
+Phase 1 (GSoC): Use OLMo's pretrained tokenizer as-is on SMILES.
+  - OLMo's 100K BPE vocab actually tokenizes SMILES more efficiently
+    than ChemBERTa's 600-token vocab (0.9x token ratio in my analysis)
+  - BUT it fragments chemical semantics: [C@@H] -> [C, @@, H, ]
+  - ChemBERTa learns chemistry-aware merges: (=O), ccccc, COc
+
+Phase 2 (stretch): Extend tokenizer with SMILES-specific tokens.
+  - Add special tokens for stereochemistry: [C@@H], [C@H], [nH]
+  - Add aromatic ring tokens: c1ccccc1
+  - Retrain BPE on mixed English + SMILES corpus
+
+### 3.2 Pseudocode
+
+Olmo2ForSequenceClassification.forward():
+```
+  hidden_states = self.model(input_ids, attention_mask)
+  # Pool: use last non-padded token
+  seq_lengths = (input_ids != pad_token_id).sum(-1) - 1
+  pooled = hidden_states[batch_range, seq_lengths]
+  logits = self.score(pooled)  # Linear(hidden_size, num_labels)
+  if labels:
+    if regression: loss = MSELoss(logits, labels)
+    if single_class: loss = CrossEntropy(logits, labels)
+    if multi_label: loss = BCEWithLogits(logits, labels)
+  return {loss, logits}
+```
+
+OLMo._prepare_batch() for causal_lm:
+```
+  tokens = tokenizer(smiles_list, padding=True)
+  input_ids = tokens.input_ids.to(device)
+  labels = input_ids.clone()  # next-token prediction
+  return {input_ids, attention_mask, labels}
+```
+
+HuggingFaceModel.generate():
+```
+  tokens = tokenizer(inputs, padding=True)
+  output_ids = model.generate(**tokens, max_new_tokens=N, **kwargs)
+  return tokenizer.batch_decode(output_ids)
+```
+
+### 3.3 Testing Plan
+
+8 unit tests (all passing in my prototype):
+
+| Test | What it validates |
+|------|-------------------|
+| test_olmo_causal_lm_pretraining | Causal LM trains, loss > 0 |
+| test_olmo_regression_finetuning | Regression trains, predictions match shape, MAE computable |
+| test_olmo_classification | Classification on binary labels, loss > 0 |
+| test_olmo_multitask_regression | MTR with 2 tasks, predictions shape matches |
+| test_olmo_save_and_restore | Checkpoint save/load, weights match exactly |
+| test_olmo_load_from_pretrained | Pretrain causal LM -> load into regression model |
+| test_olmo_generate | Single and batch generation returns strings |
+| test_olmo_invalid_task | ValueError on bad task name |
+
+All tests use a tiny config (64 hidden, 2 layers, 2 heads) — no model
+download needed, runs in ~27 seconds on CPU.
+
+Integration tests (to add during GSoC):
+- MoleculeNet benchmarks: BBBP, ESOL, FreeSolv, Lipophilicity
+- SMILES generation validity (RDKit validation)
+- Continued pretraining convergence on ZINC/PubChem subsets
+
+### 3.4 Sources of Risk
+
+| Risk | Likelihood | Mitigation |
+|------|-----------|------------|
+| OLMo-7B requires ~14GB VRAM for inference | High | Use OLMo-1B for CI/demos. Test with tiny configs. Document GPU requirements. |
+| SMILES generation validity low without extensive pretraining | High | This IS the core problem. Budget 3 weeks for pretraining experiments. Use ZINC-250K as training corpus. Target >50% validity. |
+| Olmo2ForSequenceClassification not upstream | Medium | Our implementation follows HF patterns exactly. If HF adds it later, we swap to theirs. |
+| Tokenizer fragments chemical semantics | Medium | Phase 1: works as-is (my experiments show learning happens). Phase 2: extend vocabulary. |
+| transformers version compatibility | Low | Already found and fixed one issue (PR #4913). Use top-level imports throughout. |
+
+### 3.5 Milestones & Timeline
+
+Assuming Medium size (175 hours, ~12 weeks):
+
+**Milestone 1: Core wrapper (Weeks 1-3)**
+- PR: Base class changes to HuggingFaceModel (causal_lm task, generate())
+  - Coordinate with PR #4907 to avoid duplication
+- PR: Olmo2ForSequenceClassification
+- PR: OLMo wrapper class with all task modes
+- PR: Unit tests (8 tests)
+- Deliverable: `from deepchem.models import OLMo` works for all tasks
+
+**Milestone 2: Continued pretraining (Weeks 4-6)**
+- PR: Pretraining pipeline on molecular data (ZINC-250K)
+- PR: Data loading utilities for SMILES corpora
+- PR: Pretraining tutorial notebook
+- Deliverable: Pretrained OLMo checkpoint on molecular data
+
+**Milestone 3: Fine-tuning & benchmarks (Weeks 7-9)**
+- PR: Classification tutorial (BBBP, Tox21)
+- PR: Regression tutorial (ESOL, FreeSolv, Lipophilicity)
+- PR: Benchmark results table vs ChemBERTa
+- Deliverable: Published benchmark comparing OLMo vs ChemBERTa on MoleculeNet
+
+**Milestone 4: Generation & polish (Weeks 10-12)**
+- PR: SMILES generation tutorial with RDKit validity checking
+- PR: Documentation (numpydoc, API reference, user guide)
+- PR: Tokenizer extension experiments (stretch goal)
+- Deliverable: Complete documentation and tutorials
+
+Each milestone = 1 evaluation checkpoint. PRs are <50 lines where possible,
+following DeepChem's contribution guidelines.
+
+### 3.6 Pull Request Plan
+
+I will follow DeepChem's guidelines: small PRs (<50 lines for initial ones),
+with tests and numpydoc documentation. Expected ~8-12 PRs total:
+
+1. HuggingFaceModel causal_lm support (~40 lines)
+2. generate() method (~50 lines)
+3. Olmo2ForSequenceClassification (~100 lines — larger, will discuss with mentor)
+4. OLMo wrapper class (~80 lines)
+5. Unit tests (~180 lines)
+6. Pretraining pipeline
+7. Data utilities
+8. Tutorial notebooks (3-4 notebooks)
+9. Documentation updates
+10. Benchmark scripts
+
+---
+
+## 4. COMMUNITY ENGAGEMENT
+
+- Already contributing: PR #4913 (bug fix), Issue #4912 (compat report)
+- Will attend office hours MWF 9am PST
+- Will join Discord for async discussion
+- Will write weekly progress updates
+- Happy to review other contributors' HuggingFace-related PRs
+
+---
+
+## 5. RESOURCES REQUIRED
+
+- GPU: I have access to 1x H100 NVL 96GB (Azure) for development
+- For CI: tiny model configs, no GPU needed
+- For pretraining experiments: my H100 is sufficient for OLMo-1B
+- OLMo-7B experiments: may need multi-GPU setup (discuss with mentor)
+
+---
+
+## 6. BIBLIOGRAPHY
+
+1. Groeneveld et al. (2024). "OLMo: Accelerating the Science of Language Models." arXiv:2402.00838
+2. Chithrananda et al. (2020). "ChemBERTa: Large-Scale Self-Supervised Pretraining for Molecular Property Prediction." arXiv:2010.09885
+3. Ross et al. (2022). "Large-Scale Chemical Language Representations Capture Molecular Structure and Properties." Nature Machine Intelligence.
+4. Weininger (1988). "SMILES, a chemical language and information system." J. Chem. Inf. Comput. Sci.
+5. Wu et al. (2018). "MoleculeNet: A Benchmark for Molecular Machine Learning." Chemical Science.
+
+---
+
+## KEY NUMBERS FROM YOUR EXPERIMENTS (reference these in proposal)
+
+- Tokenization: OLMo uses 0.9x tokens vs ChemBERTa on drug molecules
+- BBBP classification: ROC-AUC 0.67 (random init, 12.9M param model, 200 samples, 3 epochs)
+- ESOL regression: R^2 = 0.37, MAE = 1.27 (same conditions)
+- SMILES generation: 0% validity from random init (proves pretraining is the core challenge)
+- Test suite: 8/8 tests pass in 27 seconds on CPU
+- Stereochemistry fragmentation: [C@@H] splits into 4 tokens in OLMo vs 7 in ChemBERTa

gsoc_proposal_final.md

@@ -0,0 +1,200 @@
+# LLM Support for 7B Models (OLMo) in DeepChem
+
+Vivek Varikuti
+github.com/vivekvar-dl
+
+
+## Introduction
+
+DeepChems HuggingFaceModel wrapper works well for encoder models. ChemBERTa, MoLFormer, you can do MLM pretraining and finetune for classification or regression. But it has no support for decoder-only models. Like at all. You cant use any GPT style causal LM, and theres no way to do text generation.
+
+I want to add OLMo-2 (Allen AIs language model) to DeepChem. The idea is to make it actually useful for molecular work, not just wrap an API. That means you should be able to pretrain on SMILES data, finetune for property prediction, and generate new molecules.
+
+The reason I think this matters is that ChemBERTa can predict properties of molecules but it fundamentally cannot generate new ones. Its an encoder. OLMo is a causal LM so it can actually produce novel SMILES. For drug discovery thats a big deal. Also OLMo was trained on the Dolma corpus which has a bunch of scientific papers in it, so theres already some chemistry knowledge baked in before you even start finetuning.
+
+Why OLMo and not some other model? Mainly because its actually fully open. Weights, data, code, everything is public. No weird license stuff like LLaMA. And OLMo-2 works natively in HuggingFace transformers without needing custom packages (the older OLMo-7B needs hf_olmo installed which is annoying). It also has a 1B version which is great for testing.
+
+
+### What I already did
+
+I didnt want to just write a proposal without touching the code so I cloned DeepChem and started building.
+
+First thing that happened was ChemBERTa wouldnt even import. Turns out `transformers.models.roberta.tokenization_roberta_fast` got removed in transformers 5.x. Nobody had reported it. I fixed it (PR #4913) and filed a bigger issue about transformers 5.x compatibility (#4912) because theres more stuff broken beyond just the import.
+
+After that I started building an OLMo wrapper. Ran into an interesting problem pretty quickly. HuggingFace has OlmoForCausalLM but theres no OlmoForSequenceClassification. It doesnt exist. So if you want to do regression or classification with OLMo you have to build the classification head yourself. I wrote one using last-token pooling, basically the same thing LlamaForSequenceClassification does internally.
+
+Got everything working and ran some quick experiments on MoleculeNet:
+
+BBBP (blood brain barrier classification): ROC-AUC 0.67. This was with a tiny model, random init, 200 training samples, 3 epochs. Not great but its above 0.5 so the architecture is clearly learning something.
+
+ESOL (solubility regression): R squared 0.37. Same deal, tiny model from scratch.
+
+SMILES generation: 0% valid molecules. Every single generated SMILES was broken.
+
+That generation result is honestly the most useful thing I found. It tells you exactly where the hard problem is. The wrapper code works, the training loop works, all the plumbing is fine. But without real pretraining on a molecular corpus the model just outputs garbage. Thats what this GSoC project needs to solve.
+
+I also compared how OLMo and ChemBERTa tokenize drug molecules. Tested on aspirin, caffeine, penicillin, paclitaxel etc. OLMo actually uses fewer tokens overall (100K vocab vs 600) but it breaks up chemistry in weird ways. Like [C@@H] is a single concept (stereocenter) but OLMo splits it into four tokens. ChemBERTas tokenizer learned chemical groupings like (=O) and ccccc that make more sense.
+
+Wrote 8 unit tests, all pass in 27 sec on CPU.
+
+
+## About Me
+
+Im Vivek. I build AI stuff.
+
+The thing thats most relevant here is probably the Parameter Golf competition by OpenAI (happening right now, March 2026). You have to train the best language model that fits in 16MB total. I built custom SentencePiece tokenizers, did GPTQ quantization, designed the transformer architecture from scratch (11 layers, 512 dim, 8 heads). Dealt with flash attention compatibility issues across different hardware. Point is I actually understand how transformers work at a low level, not just how to call .fit() on them.
+
+Other stuff: GSPO-DeepSeek-R1-Distill-Qwen-1.5B (15 stars) where I did LLM distillation, wingman-AI (29 stars) which is a production AI assistant.
+
+Ive submitted PRs to a bunch of repos. HuggingFace transformers, Unsloth (aarch64 support), Anthropic SDK (fixed a streaming bottleneck), OpenAI python SDK, Karpathys nanochat (NaN loss bug in SFT). Mix of bug fixes and features.
+
+For DeepChem I found a real bug on day one (PR #4913), reported the broader compat issue (#4912), and built the OLMo prototype with passing tests.
+
+I have GPU access through Azure cloud for development.
+
+GitHub: https://github.com/vivekvar-dl
+Bug fix: https://github.com/deepchem/deepchem/pull/4913
+Issue: https://github.com/deepchem/deepchem/issues/4912
+
+
+## Work Plan
+
+### Design
+
+Four things need to be built.
+
+**1. Base class changes (HuggingFaceModel)**
+
+Right now HuggingFaceModel doesnt know causal LMs exist. Need to add:
+- "causal_lm" as a task type, using DataCollatorForLanguageModeling(mlm=False)
+- AutoModelForCausalLM branch in load_from_pretrained()
+- A generate() method wrapping HFs model.generate()
+- _prepare_batch() handling for causal LM where labels = input_ids
+
+Theres already PR #4907 from someone else that adds generation. My stuff is different, Im building a whole model wrapper not just generation support. But will coordinate so we dont duplicate work.
+
+**2. Olmo2ForSequenceClassification**
+
+This class doesnt exist in HuggingFace. Had to write it.
+
+How it works: run input through Olmo2Model, take the last non-padded tokens hidden state, project through a linear layer. Loss depends on the problem type — MSELoss for regression, CrossEntropyLoss for classification, BCEWithLogitsLoss for multi label. About 100 lines total, same pattern as LlamaForSequenceClassification.
+
+**3. OLMo wrapper class**
+
+User facing class extending HuggingFaceModel. Same structure as ChemBERTa/MoLFormer.
+
+```
+OLMo(HuggingFaceModel)
+  __init__(task, model_name, n_tasks, config)
+    task: causal_lm | regression | classification | mtr
+    Loads tokenizer, sets pad_token = eos_token
+    Syncs vocab_size with tokenizer
+    Picks model class based on task
+
+  _prepare_batch(batch)
+    causal_lm: labels = input_ids clone
+    regression/classification: labels from dataset, right dtype
+```
+
+**4. Tokenization**
+
+For now, use OLMos tokenizer as is on SMILES. My experiments show it works well enough to learn (0.67 AUC, 0.37 R2 from random init even). Not perfect with stereocenters but functional.
+
+Stretch goal: extend the vocab with chemistry tokens like [C@@H], (=O), aromatic rings. Retrain BPE on English+SMILES mix.
+
+### Pseudocode
+
+Sequence classification forward:
+```
+hidden = base_model(input_ids, attention_mask)
+seq_lengths = (input_ids != pad_id).sum(-1) - 1
+pooled = hidden[range(batch_size), seq_lengths]
+logits = linear_head(pooled)
+loss = compute_loss(logits, labels, problem_type)
+```
+
+Causal LM batch:
+```
+tokens = tokenizer(smiles_list, padding=True)
+inputs = {input_ids, attention_mask, labels: input_ids.clone()}
+```
+
+Generate:
+```
+encoded = tokenizer(prompts, padding=True)
+output_ids = model.generate(**encoded, max_new_tokens=N)
+return tokenizer.batch_decode(output_ids)
+```
+
+### Testing
+
+8 tests written and passing:
+
+1. Causal LM pretraining (loss > 0)
+2. Regression finetuning (correct prediction shape)
+3. Classification (binary labels)
+4. Multitask regression (2 targets)
+5. Save and restore checkpoint (weights match)
+6. Load pretrained into regression model
+7. Generation (single + batch)
+8. Invalid task raises error
+
+Tiny config, no downloads, 27 sec on CPU.
+
+Will add during GSoC:
+- MoleculeNet benchmarks (BBBP, ESOL, FreeSolv, Lipophilicity)
+- Generation validity checking with RDKit
+- Pretraining convergence on ZINC
+
+### Risks
+
+Generation quality is the big one. 0% validity from random init is expected obviously but getting to 50%+ valid SMILES needs real pretraining on a decent corpus. Im allocating 3 weeks for this using ZINC-250K.
+
+OLMo-7B is ~14GB just for inference. CI uses tiny configs so no GPU needed there. OLMo-1B for demos. 7B for real benchmarks, might need multi GPU, will figure that out with mentor.
+
+Olmo2ForSequenceClassification isnt upstream. If HF adds one later we swap ours out.
+
+Transformers compat — already found one issue, using top level imports everywhere going forward.
+
+### Timeline
+
+12 weeks (Medium, 175 hours):
+
+**Weeks 1-3**
+Get the core wrapper merged. Small PRs, follow DeepChems contribution guidelines. Base class changes, Olmo2ForSequenceClassification, OLMo wrapper, tests. By week 3 you should be able to do `from deepchem.models import OLMo`.
+
+**Weeks 4-6**
+Pretraining pipeline. Load SMILES from ZINC-250K, causal LM training, checkpointing. Tutorial notebook for pretraining on custom molecular data.
+
+**Weeks 7-9**
+Finetune pretrained model on MoleculeNet. BBBP and Tox21 classification, ESOL and FreeSolv and Lipophilicity regression. Benchmark table vs ChemBERTa.
+
+**Weeks 10-12**
+Generation experiments with RDKit validity checking. Tutorial notebooks. Docs. If time allows, tokenizer extension experiments.
+
+### PRs
+
+Small PRs especially at first. Bigger ones (Olmo2ForSequenceClassification ~100 lines) I will discuss with mentor before submitting. Expecting 8-12 PRs across the summer.
+
+
+## Community
+
+Already in:
+- PR #4913 (bug fix)
+- Issue #4912 (compat report)
+
+Going to do office hours MWF 9am PST, Discord for regular questions, weekly updates.
+
+
+## Resources
+
+GPU access through Azure cloud. Tiny configs for CI. OLMo-7B training setup to be discussed with mentor depending on whats needed.
+
+
+## References
+
+1. Groeneveld et al. (2024). OLMo: Accelerating the Science of Language Models. arXiv:2402.00838
+2. Chithrananda et al. (2020). ChemBERTa: Large-Scale Self-Supervised Pretraining for Molecular Property Prediction. arXiv:2010.09885
+3. Ross et al. (2022). Large-Scale Chemical Language Representations Capture Molecular Structure and Properties. Nature Machine Intelligence
+4. Weininger (1988). SMILES, a chemical language and information system. J Chem Inf Comput Sci
+5. Wu et al. (2018). MoleculeNet: A Benchmark for Molecular Machine Learning. Chemical Science

h100_training.log

@@ -0,0 +1,98 @@
+WARNING ⚠️ user config directory '/home/azureuser/.config/Ultralytics' is not writable, using '/tmp/Ultralytics'. Set YOLO_CONFIG_DIR to override.
+Creating new Ultralytics Settings v0.0.6 file ✅ 
+View Ultralytics Settings with 'yolo settings' or at '/tmp/Ultralytics/settings.json'
+Update Settings with 'yolo settings key=value', i.e. 'yolo settings runs_dir=path/to/dir'. For help see https://docs.ultralytics.com/quickstart/#ultralytics-settings.
+CUDA initialization: The NVIDIA driver on your system is too old (found version 12080). Please update your GPU driver by downloading and installing a new version from the URL: http://www.nvidia.com/Download/index.aspx Alternatively, go to: https://pytorch.org to install a PyTorch version that has been compiled with your version of the CUDA driver. (Triggered internally at /pytorch/c10/cuda/CUDAFunctions.cpp:119.)
+======================================================================
+FINAL TRAINING ON H100 - BALANCED DATASET
+======================================================================
+
+GPU Available: False
+
+======================================================================
+STEP 1: Downloading Datasets from Roboflow
+======================================================================
+
+Dataset 1: New helmet images (212)...
+
+
+
+
+
+
+Dataset 2: No-helmet images (499)...
+
+
+
+
+
+
+Dataset 3: With-helmet images (300)...
+
+
+
+
+
+
+Dataset 4: Triple-riding (626)...
+
+
+
+
+
+
+✅ All datasets downloaded!
+
+======================================================================
+STEP 2: Merging ALL Datasets
+======================================================================
+
+Unified classes (8): ['Helmet', 'Motorcycle', 'Rider', 'Triple Riding', 'helmet', 'more-than-2-person-on-2-wheeler', 'no helmet', 'with helmet']
+
+Copying datasets...
+  helmet212: 206 images
+  nohelmet499: 496 images
+  withhelmet300: 242 images
+  triple626: 626 images
+
+Final merged dataset:
+  train: 1335 images
+  valid: 186 images
+  test: 116 images
+
+Config saved: /home/azureuser/final_merged_h100/data.yaml
+
+======================================================================
+STEP 3: TRAINING ON H100 (96GB VRAM!)
+======================================================================
+
+
+Training config:
+  Model: YOLO26m
+  Epochs: 150 (faster with H100)
+  Batch: -1 (auto - H100 can handle 64-128!)
+  Image size: 640
+  Classes: 8
+
+Starting training...
+Ultralytics 8.4.37 🚀 Python-3.12.3 torch-2.11.0+cu130 
+Traceback (most recent call last):
+  File "/home/azureuser/train_h100_final.py", line 182, in <module>
+    results = model.train(
+              ^^^^^^^^^^^^
+  File "/home/azureuser/yolo_h100_env/lib/python3.12/site-packages/ultralytics/engine/model.py", line 781, in train
+    self.trainer = (trainer or self._smart_load("trainer"))(overrides=args, _callbacks=self.callbacks)
+                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/azureuser/yolo_h100_env/lib/python3.12/site-packages/ultralytics/models/yolo/detect/train.py", line 63, in __init__
+    super().__init__(cfg, overrides, _callbacks)
+  File "/home/azureuser/yolo_h100_env/lib/python3.12/site-packages/ultralytics/engine/trainer.py", line 128, in __init__
+    self.device = select_device(self.args.device)
+                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/azureuser/yolo_h100_env/lib/python3.12/site-packages/ultralytics/utils/torch_utils.py", line 230, in select_device
+    raise ValueError(
+ValueError: Invalid CUDA 'device=0' requested. Use 'device=cpu' or pass valid CUDA device(s) if available, i.e. 'device=0' or 'device=0,1,2,3' for Multi-GPU.
+
+torch.cuda.is_available(): False
+torch.cuda.device_count(): 1
+os.environ['CUDA_VISIBLE_DEVICES']: None
+

make_proposal_doc.py

@@ -0,0 +1,444 @@
+from docx import Document
+from docx.shared import Pt, Inches
+from docx.enum.text import WD_ALIGN_PARAGRAPH
+
+doc = Document()
+
+style = doc.styles['Normal']
+font = style.font
+font.name = 'Arial'
+font.size = Pt(11)
+
+# Title
+title = doc.add_heading('LLM Support for 7B Models (OLMo) in DeepChem', level=1)
+title.alignment = WD_ALIGN_PARAGRAPH.CENTER
+
+p = doc.add_paragraph()
+p.alignment = WD_ALIGN_PARAGRAPH.CENTER
+run = p.add_run('V SaraVivek (Vivek Varikuti)')
+run.bold = True
+run.font.size = Pt(12)
+p.add_run('\nvivekvarikuti22@gmail.com | github.com/vivekvar-dl | vivekvari.dev')
+
+doc.add_paragraph('')
+
+# ============ INTRODUCTION ============
+doc.add_heading('Introduction', level=2)
+
+doc.add_paragraph(
+    'DeepChems HuggingFaceModel wrapper does a solid job with encoder models. '
+    'ChemBERTa, MoLFormer, masked language modeling, classification, regression — all of that works. '
+    'But theres a pretty fundamental gap. It has zero support for decoder-only causal language models. '
+    'No GPT-style models, no text generation, nothing autoregressive.'
+)
+
+doc.add_paragraph(
+    'This project adds OLMo-2 (Allen AIs open language model) to DeepChem. '
+    'The goal isnt just to wrap another HuggingFace model though. '
+    'Its to make a causal LM genuinely useful for molecular science. '
+    'That means continued pretraining on SMILES strings, finetuning for property prediction tasks, '
+    'and most importantly — generating new molecules.'
+)
+
+doc.add_paragraph(
+    'Why does this matter for the community? ChemBERTa can tell you about a molecule but it cant make new ones. '
+    'Its an encoder, thats not what encoders do. A causal LM like OLMo can actually produce novel SMILES strings. '
+    'For drug discovery and molecular design thats a really big deal. '
+    'And because OLMo was pretrained on the Dolma corpus (which includes tons of scientific papers), '
+    'it already has some chemistry knowledge baked in before you even start finetuning on molecular data.'
+)
+
+doc.add_paragraph(
+    'I picked OLMo specifically because its fully open — weights, training data, code, all of it. '
+    'No license restrictions like LLaMA. OLMo-2 is natively supported in HuggingFace transformers so you '
+    'dont need custom packages. And it has 1B and 7B variants which is nice for development vs production.'
+)
+
+doc.add_heading('What I already built', level=3)
+
+doc.add_paragraph(
+    'I cloned DeepChem and started building before writing this proposal. '
+    'First thing that happened — ChemBERTa wouldnt import. The module '
+    'transformers.models.roberta.tokenization_roberta_fast got removed in transformers 5.x '
+    'and nobody had caught it. I fixed it in PR #4913 and filed issue #4912 about the broader '
+    'transformers 5.x compatibility problems. Thats how I got into the codebase.'
+)
+
+doc.add_paragraph(
+    'Then I started on the OLMo wrapper. Hit an interesting problem right away — '
+    'HuggingFace doesnt have OlmoForSequenceClassification. The class just doesnt exist. '
+    'So you cant do regression or classification with OLMo out of the box. '
+    'I wrote one from scratch using last-token pooling (same approach as LlamaForSequenceClassification). '
+    'Also added causal_lm as a task type in HuggingFaceModel and built a generate() method.'
+)
+
+doc.add_paragraph(
+    'Ran experiments on MoleculeNet to see if the thing actually works:'
+)
+
+doc.add_paragraph(
+    'BBBP (blood brain barrier): ROC-AUC 0.67 — tiny random init model, 200 samples, 3 epochs. '
+    'Not amazing but clearly above random chance so the architecture learns.', style='List Bullet'
+)
+doc.add_paragraph(
+    'ESOL (solubility): R squared 0.37 — same conditions.', style='List Bullet'
+)
+doc.add_paragraph(
+    'SMILES generation: 0% valid molecules. Everything it generated was broken SMILES.', style='List Bullet'
+)
+
+doc.add_paragraph(
+    'That 0% generation result is actually the most important finding. '
+    'It tells you the wrapper works fine, training works fine, but without serious pretraining on '
+    'a molecular corpus the model just outputs nonsense. Thats the core problem this project needs to solve.'
+)
+
+doc.add_paragraph(
+    'I also compared tokenization — OLMo vs ChemBERTa on real drugs (aspirin, caffeine, penicillin, paclitaxel). '
+    'OLMo uses fewer tokens overall (100K vocab vs 600) but fragments chemical concepts. '
+    '[C@@H] which is one thing in chemistry (a stereocenter) gets split into 4 tokens. '
+    'ChemBERTa learned better groupings like (=O) for carbonyl. '
+    'Something to address in the stretch goals.'
+)
+
+doc.add_paragraph('8 unit tests written, all passing in 27 seconds on CPU.')
+
+# ============ RELEVANT EXPERIENCE ============
+doc.add_heading('Relevant Experience and Interest', level=2)
+
+doc.add_paragraph(
+    'Im Vivek, I just finished my B.Tech in AI & ML from Usha Rama College of Engineering (2021-2025). '
+    'Ive been working as an AI Engineer for the past year building production systems.'
+)
+
+doc.add_heading('Work experience', level=3)
+
+doc.add_paragraph(
+    'Right now Im working with the Andhra Pradesh Police on a government AI initiative (AI4AP). '
+    'I built an AI legal compliance system for POCSO cases — RAG pipeline over 1000+ legal documents, '
+    'FastAPI backend, vector embeddings for citation-backed responses. Its a real production system that '
+    'investigating officers actually use every day. This gave me solid experience with large scale NLP, '
+    'document processing, and building things that need to work reliably.'
+)
+
+doc.add_paragraph(
+    'Before that I did an ML internship at GGS Information Services where I worked on '
+    '3D model compression (custom GANs, STEP file processing) and optimized inference pipelines '
+    'using CUDA kernels and model quantization — 30% latency reduction, 50% memory decrease. '
+    'The CUDA and quantization experience is directly relevant to working with large language models.'
+)
+
+doc.add_heading('Relevant projects', level=3)
+
+doc.add_paragraph(
+    'Parameter Golf (OpenAI, March 2026): Train the best LM that fits in 16MB. '
+    'I designed the transformer architecture from scratch — 11 layers, 512 dim, 8 heads. '
+    'Custom SentencePiece tokenizers, GPTQ quantization, flash attention. '
+    'This isnt tutorial-level stuff, I understand how transformers train from the ground up.'
+)
+
+doc.add_paragraph(
+    'GSPO-DeepSeek-R1-Distill-Qwen-1.5B (15 GitHub stars): Implemented the GSPO algorithm from '
+    'the Qwen team. Got 60% accuracy on ZebraLogic, 75.8% on math. '
+    'Beat PPO and GRPO baselines.'
+)
+
+doc.add_paragraph(
+    'Dial 112 AI: Built a production system for AP Police processing 1000+ emergency calls daily. '
+    'Speech-to-text, sentiment analysis, priority classification, geospatial dispatch optimization.'
+)
+
+doc.add_paragraph(
+    'Published a paper on pose-guided image generation (PPAG) — '
+    'progressive pose attention for identity-preserving synthesis, 92% identity preservation score.'
+)
+
+doc.add_heading('Open source', level=3)
+
+doc.add_paragraph(
+    'PRs to HuggingFace transformers, Unsloth (aarch64 support), Anthropic SDK (streaming perf fix), '
+    'OpenAI SDK, Karpathys nanochat (NaN loss fix for SFT training).'
+)
+
+doc.add_paragraph(
+    'DeepChem contributions: PR #4913 (ChemBERTa import fix for transformers 5.x), '
+    'Issue #4912 (transformers compatibility report).'
+)
+
+doc.add_heading('Why this project', level=3)
+
+doc.add_paragraph(
+    '[FILL THIS IN YOURSELF — what got you interested in molecular ML? '
+    'Why DeepChem specifically? Be honest here, mentors can tell when someone is just saying '
+    'what they think you want to hear. Talk about what actually excites you.]'
+)
+
+p = doc.add_paragraph()
+p.add_run('GitHub: ').bold = True
+p.add_run('https://github.com/vivekvar-dl')
+p = doc.add_paragraph()
+p.add_run('LinkedIn: ').bold = True
+p.add_run('https://linkedin.com/in/vivekvar')
+p = doc.add_paragraph()
+p.add_run('Bug fix PR: ').bold = True
+p.add_run('https://github.com/deepchem/deepchem/pull/4913')
+p = doc.add_paragraph()
+p.add_run('Compat issue: ').bold = True
+p.add_run('https://github.com/deepchem/deepchem/issues/4912')
+
+# ============ WORK PLAN ============
+doc.add_heading('Work Plan', level=2)
+
+doc.add_paragraph(
+    'The project breaks down into four components that build on each other. '
+    'First the base class needs to learn about causal LMs, then we build the OLMo specific stuff on top, '
+    'then pretraining and benchmarks, then generation and docs.'
+)
+
+# Design
+doc.add_heading('Design and Pseudocode', level=3)
+
+doc.add_paragraph().add_run('Component 1: HuggingFaceModel base class changes').bold = True
+doc.add_paragraph(
+    'The wrapper currently doesnt know causal LMs exist. Need to add "causal_lm" as a task type '
+    'with DataCollatorForLanguageModeling(mlm=False), an AutoModelForCausalLM branch in '
+    'load_from_pretrained(), a generate() method wrapping HFs model.generate() API, and '
+    '_prepare_batch() handling where labels = input_ids for next token prediction.'
+)
+doc.add_paragraph(
+    'Note: PR #4907 from another contributor adds a generate() method too. '
+    'My work is different — building a complete model wrapper, not just generation. '
+    'Will coordinate to avoid overlap.'
+)
+
+doc.add_paragraph().add_run('Component 2: Olmo2ForSequenceClassification').bold = True
+doc.add_paragraph(
+    'This class doesnt exist in HuggingFace. I already built it. '
+    'Takes Olmo2Model output, grabs the last non-padded tokens hidden state (last-token pooling), '
+    'projects through a linear layer. Loss computation depends on problem_type — '
+    'MSELoss for regression, CrossEntropyLoss for single label, BCEWithLogitsLoss for multi label. '
+    'About 100 lines, follows same pattern as LlamaForSequenceClassification.'
+)
+
+doc.add_paragraph('Forward pass pseudocode:')
+doc.add_paragraph(
+    'hidden = base_model(input_ids, attention_mask)\n'
+    'seq_lengths = (input_ids != pad_id).sum(-1) - 1\n'
+    'pooled = hidden[range(batch_size), seq_lengths]\n'
+    'logits = linear_head(pooled)\n'
+    'loss = compute_loss(logits, labels, problem_type)',
+    style='No Spacing'
+)
+
+doc.add_paragraph().add_run('Component 3: OLMo wrapper class').bold = True
+doc.add_paragraph(
+    'User-facing class extending HuggingFaceModel. Same pattern as ChemBERTa and MoLFormer.'
+)
+doc.add_paragraph(
+    'OLMo.__init__(task, model_name, n_tasks, config)\n'
+    '  task: causal_lm | regression | classification | mtr\n'
+    '  Loads tokenizer, sets pad_token = eos_token\n'
+    '  Syncs vocab_size with tokenizer\n'
+    '  Creates right model class based on task\n\n'
+    'OLMo._prepare_batch(batch)\n'
+    '  causal_lm: labels = input_ids clone\n'
+    '  regression: float labels\n'
+    '  classification: long (single) or float (multi)',
+    style='No Spacing'
+)
+
+doc.add_paragraph().add_run('Component 4: Tokenization strategy').bold = True
+doc.add_paragraph(
+    'Phase 1: Use OLMos tokenizer directly on SMILES. Works — '
+    'my experiments show 0.67 AUC and 0.37 R2 from random init which means learning is happening. '
+    'Not perfect with stereocenters but functional.'
+)
+doc.add_paragraph(
+    'Stretch: Extend vocab with chemistry tokens — [C@@H], [nH], (=O), aromatic ring patterns. '
+    'Retrain BPE on mixed English + SMILES corpus. Could help generation quality a lot.'
+)
+
+# Testing
+doc.add_heading('Testing Plan', level=3)
+
+doc.add_paragraph('I already have 8 unit tests written and passing:')
+
+tests = [
+    ('test_olmo_causal_lm_pretraining', 'Trains with causal LM objective, loss is positive'),
+    ('test_olmo_regression_finetuning', 'Regression training, predictions have correct shape, MAE score works'),
+    ('test_olmo_classification', 'Binary classification on random labels'),
+    ('test_olmo_multitask_regression', '2-task regression, output shape matches'),
+    ('test_olmo_save_and_restore', 'Save checkpoint, load into new model, verify all weights match'),
+    ('test_olmo_load_from_pretrained', 'Pretrain causal LM then load weights into regression model'),
+    ('test_olmo_generate', 'Single string and batch generation, returns valid strings'),
+    ('test_olmo_invalid_task', 'Bad task name raises ValueError'),
+]
+for name, desc in tests:
+    doc.add_paragraph(f'{name} — {desc}', style='List Bullet')
+
+doc.add_paragraph(
+    'All tests use a tiny config (64 hidden, 2 layers, 2 heads). No model downloads. '
+    'Runs in 27 seconds on CPU.'
+)
+
+doc.add_paragraph('Integration tests to add during GSoC:')
+doc.add_paragraph('MoleculeNet benchmarks — BBBP, ESOL, FreeSolv, Lipophilicity', style='List Bullet')
+doc.add_paragraph('SMILES generation validity checking with RDKit', style='List Bullet')
+doc.add_paragraph('Pretraining convergence curves on ZINC subsets', style='List Bullet')
+
+# Risks
+doc.add_heading('Sources of Risk', level=3)
+
+doc.add_paragraph().add_run('Generation quality:').bold = True
+doc.add_paragraph(
+    'This is the biggest one. 0% validity from random init is expected but getting to '
+    'something useful like 50%+ valid SMILES needs real pretraining on a molecular corpus. '
+    'Im allocating 3 full weeks for pretraining experiments using ZINC-250K as training data. '
+    'If ZINC isnt enough, backup plan is PubChem subsets or combining multiple datasets.'
+)
+
+doc.add_paragraph().add_run('GPU memory for 7B model:').bold = True
+doc.add_paragraph(
+    'OLMo-7B needs ~14GB VRAM just for inference. '
+    'CI tests use tiny configs so no GPU needed. Demos use OLMo-1B. '
+    '7B experiments may need multi-GPU, will figure out with mentor.'
+)
+
+doc.add_paragraph().add_run('Olmo2ForSequenceClassification not upstream:').bold = True
+doc.add_paragraph(
+    'I wrote this class myself since HuggingFace doesnt have it. '
+    'If they add one later we just swap ours out. Follows their patterns exactly so low risk.'
+)
+
+doc.add_paragraph().add_run('Transformers version compatibility:').bold = True
+doc.add_paragraph(
+    'Already found and fixed one issue (PR #4913). '
+    'Using top-level imports everywhere going forward. Will test against both 4.x and 5.x.'
+)
+
+# Milestones
+doc.add_heading('Milestones', level=3)
+
+doc.add_paragraph().add_run('Milestone 1 (end of week 3): Core wrapper working').bold = True
+doc.add_paragraph(
+    'from deepchem.models import OLMo works. All task modes functional. '
+    'Tests passing. Base class changes merged.'
+)
+
+doc.add_paragraph().add_run('Milestone 2 (end of week 6): Pretraining pipeline done').bold = True
+doc.add_paragraph(
+    'Can load SMILES data, pretrain OLMo with causal LM objective, save checkpoints. '
+    'Tutorial notebook showing how to pretrain on custom data. '
+    'First generation results with validity numbers.'
+)
+
+doc.add_paragraph().add_run('Milestone 3 (end of week 12): Everything shipped').bold = True
+doc.add_paragraph(
+    'MoleculeNet benchmark results published. Generation tutorial with RDKit validation. '
+    'Full documentation. All PRs merged.'
+)
+
+# Timeline
+doc.add_heading('Timeline', level=3)
+
+doc.add_paragraph('12 weeks, Medium size (175 hours):')
+
+weeks = [
+    ('Week 1', 'Set up dev environment, submit base class PR (causal_lm task support, ~40 lines). Start on generate() method.'),
+    ('Week 2', 'Submit generate() PR (~50 lines). Start Olmo2ForSequenceClassification. Discuss size with mentor since its ~100 lines.'),
+    ('Week 3', 'Submit OLMo wrapper PR and unit test PR. Get through review. Target: all core PRs merged by end of week.'),
+    ('Week 4', 'Start pretraining pipeline. Data loading for ZINC-250K. Figure out training hyperparameters on OLMo-1B.'),
+    ('Week 5', 'Run pretraining experiments. Monitor convergence. First generation attempts — check validity with RDKit.'),
+    ('Week 6', 'Finish pretraining PR and tutorial notebook. Submit for review. Save best checkpoint.'),
+    ('Week 7', 'Start finetuning experiments. BBBP and Tox21 classification with pretrained model.'),
+    ('Week 8', 'ESOL, FreeSolv, Lipophilicity regression. Build benchmark comparison table vs ChemBERTa.'),
+    ('Week 9', 'Submit benchmark PR and finetuning tutorials. Respond to review feedback.'),
+    ('Week 10', 'SMILES generation experiments. Validity rate analysis. Different sampling strategies (temperature, top-k, nucleus).'),
+    ('Week 11', 'Generation tutorial notebook. Documentation — numpydoc for all classes/methods, API reference updates.'),
+    ('Week 12', 'Final review rounds. Clean up any open PRs. Stretch: tokenizer extension experiments if time allows.'),
+]
+for week, desc in weeks:
+    p = doc.add_paragraph()
+    p.add_run(f'{week}: ').bold = True
+    p.add_run(desc)
+
+# Pull Requests
+doc.add_heading('Pull Requests', level=3)
+
+doc.add_paragraph(
+    'Following DeepChems guidelines for new contributors — small PRs, especially at the start. '
+    'Heres the planned breakdown:'
+)
+
+prs = [
+    ('PR 1 (Week 1)', 'HuggingFaceModel causal_lm support', '~40 lines'),
+    ('PR 2 (Week 2)', 'generate() method', '~50 lines'),
+    ('PR 3 (Week 2-3)', 'Olmo2ForSequenceClassification', '~100 lines (will discuss with mentor)'),
+    ('PR 4 (Week 3)', 'OLMo wrapper class', '~80 lines'),
+    ('PR 5 (Week 3)', 'Unit tests', '~180 lines'),
+    ('PR 6 (Week 6)', 'Pretraining pipeline + data utils', 'TBD'),
+    ('PR 7 (Week 6)', 'Pretraining tutorial notebook', 'notebook'),
+    ('PR 8 (Week 9)', 'Finetuning tutorials (classification + regression)', 'notebooks'),
+    ('PR 9 (Week 9)', 'Benchmark results', 'TBD'),
+    ('PR 10 (Week 11)', 'Generation tutorial', 'notebook'),
+    ('PR 11 (Week 12)', 'Documentation updates', 'TBD'),
+]
+for pr, desc, size in prs:
+    p = doc.add_paragraph()
+    p.add_run(f'{pr}: ').bold = True
+    p.add_run(f'{desc} ({size})')
+
+doc.add_paragraph(
+    'Each PR goes through review at office hours. Bigger ones might need 2-3 rounds. '
+    'I have buffer built into the timeline for this — not assuming everything merges on first try.'
+)
+
+# ============ COMMUNITY ============
+doc.add_heading('Community', level=2)
+
+doc.add_paragraph('What Ive done so far:')
+doc.add_paragraph('PR #4913 — fixed ChemBERTa import crash for transformers 5.x', style='List Bullet')
+doc.add_paragraph('Issue #4912 — reported broader transformers 5.x compatibility problems', style='List Bullet')
+doc.add_paragraph('Built and tested OLMo prototype locally against DeepChems codebase', style='List Bullet')
+
+doc.add_paragraph(
+    'I can commit to attending at least 2 office hour sessions per week (MWF 9am PST). '
+    'Will also be active on Discord for async discussion and will do weekly progress updates.'
+)
+
+doc.add_paragraph(
+    '[NOTE: mention which mentors you have talked to once you connect on Discord. '
+    'Riya and Harindhar are listed as mentors for this project.]'
+)
+
+# ============ RESOURCES ============
+doc.add_heading('Resources Required', level=2)
+
+doc.add_paragraph(
+    'I have GPU access through Azure cloud which should handle OLMo-1B training and OLMo-7B inference. '
+    'For CI and unit tests everything runs on CPU with tiny model configs so no special compute needed there.'
+)
+
+doc.add_paragraph(
+    'For OLMo-7B full training we might need a multi-GPU setup. '
+    'Would be good to discuss with the mentor what compute DeepChem can provide '
+    'or if Colab Pro / cloud credits would work. '
+    'The pretraining experiments on OLMo-1B should be doable on my current setup.'
+)
+
+# ============ REFERENCES ============
+doc.add_heading('References', level=2)
+
+refs = [
+    'Groeneveld et al. (2024). OLMo: Accelerating the Science of Language Models. arXiv:2402.00838',
+    'Chithrananda et al. (2020). ChemBERTa: Large-Scale Self-Supervised Pretraining for Molecular Property Prediction. arXiv:2010.09885',
+    'Ross et al. (2022). Large-Scale Chemical Language Representations Capture Molecular Structure and Properties. Nature Machine Intelligence',
+    'Weininger (1988). SMILES, a chemical language and information system. J Chem Inf Comput Sci',
+    'Wu et al. (2018). MoleculeNet: A Benchmark for Molecular Machine Learning. Chemical Science',
+]
+for i, ref in enumerate(refs, 1):
+    doc.add_paragraph(f'{i}. {ref}')
+
+# Save
+doc.save('/home/azureuser/GSoC_2026_Proposal_Vivek_OLMo.docx')
+print('Done: ~/GSoC_2026_Proposal_Vivek_OLMo.docx')

merge_v3.py

@@ -0,0 +1,88 @@
+#!/usr/bin/env python3
+# v3: use clean_merged_data (v2) as base, add Khadatkar + Learning to train
+import os, shutil, glob
+from pathlib import Path
+
+HOME = os.path.expanduser('~')
+OUT  = f'{HOME}/merged_v3'
+BASE = f'{HOME}/clean_merged_data'
+
+# Fresh output
+if os.path.exists(OUT): shutil.rmtree(OUT)
+for s in ('train','valid','test'):
+    os.makedirs(f'{OUT}/{s}/images', exist_ok=True)
+    os.makedirs(f'{OUT}/{s}/labels', exist_ok=True)
+
+stats = {s: {c: 0 for c in (0,1,2)} for s in ('train','valid','test')}
+imgcount = {s: 0 for s in ('train','valid','test')}
+
+def copy_split(src_img_dir, src_lbl_dir, target_split, cmap, name_suffix):
+    n = 0
+    for lbl_path in glob.glob(f'{src_lbl_dir}/*.txt'):
+        stem = Path(lbl_path).stem
+        img_path = None
+        for ext in ('.jpg','.jpeg','.png','.JPG','.PNG'):
+            p = f'{src_img_dir}/{stem}{ext}'
+            if os.path.exists(p): img_path = p; break
+        if img_path is None: continue
+        lines = []
+        with open(lbl_path) as f:
+            for line in f:
+                parts = line.strip().split()
+                if not parts: continue
+                cid = int(parts[0])
+                if cid not in cmap: continue
+                lines.append(' '.join([str(cmap[cid])] + parts[1:]))
+        if not lines: continue
+        new_stem = f'{stem}{name_suffix}'
+        ext = Path(img_path).suffix
+        dst_img = f'{OUT}/{target_split}/images/{new_stem}{ext}'
+        dst_lbl = f'{OUT}/{target_split}/labels/{new_stem}.txt'
+        if not os.path.exists(dst_img):
+            try: os.link(img_path, dst_img)
+            except: shutil.copy(img_path, dst_img)
+        with open(dst_lbl, 'w') as f:
+            f.write('\n'.join(lines) + '\n')
+        for ln in lines:
+            stats[target_split][int(ln.split()[0])] += 1
+        imgcount[target_split] += 1
+        n += 1
+    return n
+
+# 1) Copy clean_merged_data AS-IS (identity mapping for 0,1,2), no extra suffix
+#    Images already have _cctv_dataset / _helmet_dataset / _yolo_project suffixes
+print('--- base v2 data ---')
+for s in ('train','valid','test'):
+    n = copy_split(f'{BASE}/{s}/images', f'{BASE}/{s}/labels', s, {0:0,1:1,2:2}, '')
+    print(f'  base -> {s}: {n}')
+
+# 2) Add Khadatkar + Learning ONLY to train split
+EXTRAS = [
+    ('khadatkar', f'{HOME}/extra_khadatkar', {0:1, 1:0}),  # 0=With Helmet->1, 1=Without Helmet->0, drop 2=licence
+    ('learning',  f'{HOME}/extra_learning',  {0:1, 1:0}),  # 0=With Helmet->1, 1=Without Helmet->0
+]
+print('--- extras -> train ---')
+for name, root, cmap in EXTRAS:
+    for src_split in ('train','valid','test'):
+        img_dir = f'{root}/{src_split}/images'
+        lbl_dir = f'{root}/{src_split}/labels'
+        if not os.path.isdir(lbl_dir): continue
+        n = copy_split(img_dir, lbl_dir, 'train', cmap, f'_{name}_{src_split}')
+        print(f'  {name} {src_split} -> train: {n}')
+
+yaml = f'''path: {OUT}
+train: train/images
+val: valid/images
+test: test/images
+nc: 3
+names:
+  0: no-helmet
+  1: with-helmet
+  2: triple-riding
+'''
+with open(f'{OUT}/data.yaml','w') as f: f.write(yaml)
+
+print('\n=== V3 MERGE COMPLETE ===')
+for s in ('train','valid','test'):
+    tot = sum(stats[s].values())
+    print(f'  {s:6s} images={imgcount[s]:5d} | no-helmet={stats[s][0]:5d} with-helmet={stats[s][1]:5d} triple={stats[s][2]:4d} | instances={tot}')

pull_extra.py

@@ -0,0 +1,27 @@
+#!/usr/bin/env python3
+# Pull 3 extra helmet/rider datasets from Roboflow Universe
+from roboflow import Roboflow
+
+# API key reused from previous script
+rf = Roboflow(api_key='qeQs9chVa3kU0XnpTZsd')
+
+targets = [
+    ('gw-khadatkar-and-sv-wasule', 'helmet-and-no-helmet-rider-detection', '~/extra_khadatkar'),
+    ('nckh-2023',                  'helmet-detection-project',              '~/extra_nckh'),
+    ('learning-evidence',          'helmet-detection_yolov8',               '~/extra_learning'),
+]
+
+for ws, proj, loc in targets:
+    print(f'\n=== {ws}/{proj} ===')
+    try:
+        p = rf.workspace(ws).project(proj)
+        # try latest version
+        vs = p.versions()
+        if not vs:
+            print('  NO VERSIONS'); continue
+        v = vs[0]
+        print(f'  version={v.version}')
+        ds = v.download('yolov8', location=loc)
+        print(f'  DOWNLOADED to {loc}')
+    except Exception as e:
+        print(f'  FAIL: {e}')

telugu_voice_clone.py

@@ -0,0 +1,113 @@
+"""
+Telugu Voice Cloning with IndicF5
+
+Usage:
+  1. Place your reference audio as 'reference.wav' (10-15 seconds, clean Telugu speech)
+  2. Edit REF_TEXT with the exact Telugu transcript of your reference audio
+  3. Edit GEN_TEXT with the Telugu text you want to generate
+  4. Run: source ~/indicf5-env/bin/activate && python telugu_voice_clone.py
+"""
+
+import torch
+import numpy as np
+import soundfile as sf
+import io
+import time
+from pydub import AudioSegment, silence
+from huggingface_hub import hf_hub_download
+from f5_tts.infer.utils_infer import (
+    infer_process,
+    load_model,
+    load_vocoder,
+    preprocess_ref_audio_text,
+)
+from f5_tts.model import DiT
+
+# === CONFIGURE THESE ===
+
+# Path to your reference voice recording (WAV, 10-15 seconds, Telugu)
+REF_AUDIO = "reference.wav"
+
+# Exact Telugu transcript of your reference audio
+REF_TEXT = "ఇది నా గొంతు నమూనా, నేను తెలుగులో మాట్లాడుతున్నాను."
+
+# Telugu text you want to generate in your cloned voice
+GEN_TEXT = "నమస్కారం, మీరు ఎలా ఉన్నారు? నేను మీతో తెలుగులో మాట్లాడుతున్నాను."
+
+# Output file
+OUTPUT_FILE = "output_telugu.wav"
+
+SPEED = 1.0
+REMOVE_SILENCE = True
+
+# === END CONFIG ===
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
+
+# Load vocoder
+print("Loading vocoder...")
+vocoder = load_vocoder(vocoder_name="vocos", is_local=False, device=device)
+
+# Download vocab and load model
+print("Downloading IndicF5 model...")
+repo_id = "ai4bharat/IndicF5"
+vocab_path = hf_hub_download(repo_id, filename="checkpoints/vocab.txt")
+
+ema_model = load_model(
+    DiT,
+    dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4),
+    mel_spec_type="vocos",
+    vocab_file=vocab_path,
+    device=device,
+)
+print("Model loaded!")
+
+# Preprocess reference audio
+print(f"Reference audio: {REF_AUDIO}")
+ref_audio, ref_text = preprocess_ref_audio_text(REF_AUDIO, REF_TEXT)
+
+# Generate
+print(f"Generating: {GEN_TEXT[:80]}...")
+start = time.time()
+audio, final_sample_rate, _ = infer_process(
+    ref_audio,
+    ref_text,
+    GEN_TEXT,
+    ema_model,
+    vocoder,
+    mel_spec_type="vocos",
+    speed=SPEED,
+    device=device,
+)
+print(f"Generated in {time.time() - start:.1f}s")
+
+# Post-process: remove silence and normalize
+buffer = io.BytesIO()
+sf.write(buffer, audio, samplerate=24000, format="WAV")
+buffer.seek(0)
+audio_segment = AudioSegment.from_file(buffer, format="wav")
+
+if REMOVE_SILENCE:
+    non_silent_segs = silence.split_on_silence(
+        audio_segment,
+        min_silence_len=1000,
+        silence_thresh=-50,
+        keep_silence=500,
+        seek_step=10,
+    )
+    if non_silent_segs:
+        audio_segment = sum(non_silent_segs, AudioSegment.silent(duration=0))
+
+# Normalize loudness
+target_dBFS = -20.0
+change_in_dBFS = target_dBFS - audio_segment.dBFS
+audio_segment = audio_segment.apply_gain(change_in_dBFS)
+
+# Save
+final_audio = np.array(audio_segment.get_array_of_samples())
+if final_audio.dtype == np.int16:
+    final_audio = final_audio.astype(np.float32) / 32768.0
+sf.write(OUTPUT_FILE, final_audio.astype(np.float32), samplerate=24000)
+print(f"Saved to {OUTPUT_FILE}")
+print("Done!")

train_h100_clean.py

@@ -0,0 +1,50 @@
+#!/usr/bin/env python3
+# 3-class clean training on H100 NVL
+# Classes: 0 no-helmet | 1 with-helmet | 2 triple-riding
+from ultralytics import YOLO
+import torch, os
+
+print('GPU:', torch.cuda.get_device_name(0), '|', torch.cuda.get_device_properties(0).total_memory/1e9, 'GB')
+
+# Start from pretrained yolo26m (auto-downloads if missing)
+model = YOLO('yolo26m.pt')
+
+results = model.train(
+    data='/home/azureuser/clean_merged_data/data.yaml',
+    epochs=150,
+    imgsz=640,
+    batch=64,           # H100 NVL has 95GB, can push batch high
+    device=0,
+    workers=8,
+    project='runs_clean',
+    name='h100_3class',
+    exist_ok=True,
+    amp=True,
+    cos_lr=True,
+    close_mosaic=15,
+    # augmentation — important for 10k image dataset
+    mosaic=1.0,
+    mixup=0.15,
+    copy_paste=0.3,     # boost with-helmet via cross-image pasting
+    hsv_h=0.015, hsv_s=0.7, hsv_v=0.4,
+    degrees=5.0,
+    translate=0.1,
+    scale=0.5,
+    fliplr=0.5,
+    # loss
+    cls=1.0,            # classification loss weight (bump if still confused)
+    box=7.5,
+    dfl=1.5,
+    # regularization
+    weight_decay=0.0005,
+    dropout=0.0,
+    # schedule
+    optimizer='auto',
+    lr0=0.01,
+    patience=40,
+    plots=True,
+    verbose=True,
+)
+print('TRAIN DONE — running val on test split')
+m = YOLO('runs_clean/h100_3class/weights/best.pt')
+m.val(data='/home/azureuser/clean_merged_data/data.yaml', split='test', plots=True, save_json=True)

train_h100_final.py

@@ -0,0 +1,230 @@
+#!/usr/bin/env python3
+"""
+Final Training on H100 - 96GB VRAM Beast!
+Merges ALL datasets and trains with maximum performance
+"""
+
+from roboflow import Roboflow
+from ultralytics import YOLO
+import torch
+import os
+import shutil
+import yaml
+import glob
+from pathlib import Path
+
+print("=" * 70)
+print("FINAL TRAINING ON H100 - BALANCED DATASET")
+print("=" * 70)
+
+# Check GPU
+print(f"\nGPU Available: {torch.cuda.is_available()}")
+if torch.cuda.is_available():
+    print(f"GPU: {torch.cuda.get_device_name(0)}")
+    print(f"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.0f} GB")
+
+# Step 1: Download all datasets from Roboflow
+print("\n" + "=" * 70)
+print("STEP 1: Downloading Datasets from Roboflow")
+print("=" * 70)
+
+rf = Roboflow(api_key="cMpZOr1EizWFVrJ0Au4o")
+
+# Dataset 1: New 212 helmet images
+print("\nDataset 1: New helmet images (212)...")
+project1 = rf.workspace("team11s-workspace-man05").project("helmet-detection-ihomd")
+ds1 = project1.version(1).download("yolov8", location="~/helmet_212")
+
+# Dataset 2: Old no-helmet (499) from first account
+print("\nDataset 2: No-helmet images (499)...")
+rf2 = Roboflow(api_key="qeQs9chVa3kU0XnpTZsd")
+project2 = rf2.workspace("nyc-nleyq").project("indian-cctv-traffic-violations")
+ds2 = project2.version(1).download("yolov8", location="~/no_helmet_499")
+
+# Dataset 3: With-helmet (300) from second account
+print("\nDataset 3: With-helmet images (300)...")
+project3 = rf2.workspace("vivekvarikuti").project("withhelmet")
+ds3 = project3.version(1).download("yolov8", location="~/with_helmet_300")
+
+# Dataset 4: Triple-riding from original (626)
+print("\nDataset 4: Triple-riding (626)...")
+project4 = rf2.workspace("triple-ride-rsysj").project("triple-riding-detection-pniom")
+ds4 = project4.version(1).download("yolov8", location="~/triple_riding_626")
+
+print("\n✅ All datasets downloaded!")
+
+# Step 2: Merge all datasets
+print("\n" + "=" * 70)
+print("STEP 2: Merging ALL Datasets")
+print("=" * 70)
+
+MERGED_DIR = os.path.expanduser("~/final_merged_h100")
+
+for split in ['train', 'valid', 'test']:
+    os.makedirs(f"{MERGED_DIR}/{split}/images", exist_ok=True)
+    os.makedirs(f"{MERGED_DIR}/{split}/labels", exist_ok=True)
+
+# Collect all classes
+all_classes = set()
+datasets = [
+    (ds1.location, 'helmet212'),
+    (ds2.location, 'nohelmet499'),
+    (ds3.location, 'withhelmet300'),
+    (ds4.location, 'triple626')
+]
+
+class_configs = {}
+for ds_path, ds_name in datasets:
+    yaml_path = f"{ds_path}/data.yaml"
+    if os.path.exists(yaml_path):
+        with open(yaml_path, 'r') as f:
+            cfg = yaml.safe_load(f)
+            class_configs[ds_name] = cfg
+            if 'names' in cfg:
+                all_classes.update(cfg['names'])
+
+unified_classes = sorted(list(all_classes))
+print(f"\nUnified classes ({len(unified_classes)}): {unified_classes}")
+
+# Create class mappings
+class_maps = {}
+for ds_name, cfg in class_configs.items():
+    class_maps[ds_name] = {}
+    if 'names' in cfg:
+        for i, cls in enumerate(cfg['names']):
+            class_maps[ds_name][i] = unified_classes.index(cls)
+
+# Copy and merge datasets
+def copy_with_remap(src_dir, prefix, class_mapping):
+    total = 0
+    for split in ['train', 'valid', 'test']:
+        src_img = f"{src_dir}/{split}/images"
+        src_lbl = f"{src_dir}/{split}/labels"
+
+        if not os.path.exists(src_img):
+            continue
+
+        imgs = glob.glob(f"{src_img}/*.jpg") + glob.glob(f"{src_img}/*.png")
+
+        for img_path in imgs:
+            img_name = os.path.basename(img_path)
+            lbl_name = Path(img_path).stem + '.txt'
+            lbl_path = f"{src_lbl}/{lbl_name}"
+
+            # Copy image with prefix
+            dst_img = f"{MERGED_DIR}/{split}/images/{prefix}_{img_name}"
+            shutil.copy2(img_path, dst_img)
+
+            # Remap and copy label
+            if os.path.exists(lbl_path):
+                with open(lbl_path, 'r') as f:
+                    lines = f.readlines()
+
+                remapped = []
+                for line in lines:
+                    parts = line.strip().split()
+                    if len(parts) >= 5:
+                        old_cls = int(parts[0])
+                        new_cls = class_mapping.get(old_cls, old_cls)
+                        remapped.append(f"{new_cls} {' '.join(parts[1:])}\n")
+
+                if remapped:
+                    dst_lbl = f"{MERGED_DIR}/{split}/labels/{prefix}_{lbl_name}"
+                    with open(dst_lbl, 'w') as f:
+                        f.writelines(remapped)
+                    total += 1
+
+    return total
+
+print("\nCopying datasets...")
+for (ds_path, ds_name), prefix in zip(datasets, ['h212', 'nh499', 'wh300', 'tr626']):
+    count = copy_with_remap(ds_path, prefix, class_maps.get(ds_name, {}))
+    print(f"  {ds_name}: {count} images")
+
+# Count final
+print("\nFinal merged dataset:")
+for split in ['train', 'valid', 'test']:
+    imgs = glob.glob(f"{MERGED_DIR}/{split}/images/*")
+    print(f"  {split}: {len(imgs)} images")
+
+# Create YAML
+merged_yaml = {
+    'path': MERGED_DIR,
+    'train': 'train/images',
+    'val': 'valid/images',
+    'test': 'test/images',
+    'nc': len(unified_classes),
+    'names': unified_classes
+}
+
+yaml_path = f"{MERGED_DIR}/data.yaml"
+with open(yaml_path, 'w') as f:
+    yaml.dump(merged_yaml, f, default_flow_style=False)
+
+print(f"\nConfig saved: {yaml_path}")
+
+# Step 3: Train on H100 with OPTIMIZED settings
+print("\n" + "=" * 70)
+print("STEP 3: TRAINING ON H100 (96GB VRAM!)")
+print("=" * 70)
+
+model = YOLO('yolo26m.pt')
+
+print(f"\nTraining config:")
+print(f"  Model: YOLO26m")
+print(f"  Epochs: 150 (faster with H100)")
+print(f"  Batch: -1 (auto - H100 can handle 64-128!)")
+print(f"  Image size: 640")
+print(f"  Classes: {len(unified_classes)}")
+
+print("\nStarting training...")
+
+results = model.train(
+    data=yaml_path,
+    epochs=150,  # Fewer epochs needed with large batch on H100
+    imgsz=640,
+    batch=-1,  # Auto batch (H100 will use 64-128!)
+    cache='ram',  # H100 has tons of RAM
+    device=0,
+    workers=8,
+    patience=30,
+    name='h100_final',
+    project='outputs',
+
+    # Augmentation
+    hsv_h=0.015,
+    hsv_s=0.7,
+    hsv_v=0.4,
+    degrees=10,
+    translate=0.1,
+    scale=0.5,
+    fliplr=0.5,
+    mosaic=1.0,
+    mixup=0.1,
+
+    lr0=0.01,
+    lrf=0.01,
+    amp=True,
+    val=True,
+    plots=True,
+)
+
+print("\n" + "=" * 70)
+print("TRAINING COMPLETE!")
+print("=" * 70)
+
+# Validate
+metrics = model.val()
+print(f"\nFinal Metrics:")
+print(f"  mAP50:     {metrics.box.map50:.4f} ({metrics.box.map50*100:.1f}%)")
+print(f"  mAP50-95:  {metrics.box.map:.4f} ({metrics.box.map*100:.1f}%)")
+print(f"  Precision: {metrics.box.mp:.4f} ({metrics.box.mp*100:.1f}%)")
+print(f"  Recall:    {metrics.box.mr:.4f} ({metrics.box.mr*100:.1f}%)")
+
+# Export
+print("\nExporting to ONNX...")
+model.export(format='onnx', dynamic=True, simplify=True)
+
+print("\n" + "=" * 70)
+print("Model saved: outputs/h100_final/weights/best.pt")
+print("=" * 70)

train_v3.py

@@ -0,0 +1,35 @@
+#!/usr/bin/env python3
+# v3 max-throughput: yolo26m on 6673-img dataset, batch=128, cache=ram
+from ultralytics import YOLO
+import torch
+print('GPU:', torch.cuda.get_device_name(0), '|', round(torch.cuda.get_device_properties(0).total_memory/1e9), 'GB')
+
+model = YOLO('yolo26m.pt')
+model.train(
+    data='/home/azureuser/merged_v3/data.yaml',
+    epochs=200,
+    imgsz=640,
+    batch=128,          # 2x v2 — should hit ~70GB VRAM
+    device=0,
+    workers=16,         # feed data faster
+    cache='ram',        # dataset is ~1GB, fits easily
+    project='runs_v3',
+    name='h100_3class_v3',
+    exist_ok=True,
+    amp=True,
+    cos_lr=True,
+    close_mosaic=20,
+    mosaic=1.0, mixup=0.15, copy_paste=0.3,
+    hsv_h=0.015, hsv_s=0.7, hsv_v=0.4,
+    degrees=5.0, translate=0.1, scale=0.5, fliplr=0.5,
+    cls=1.0, box=7.5, dfl=1.5,
+    weight_decay=0.0005,
+    optimizer='auto',
+    patience=60,
+    plots=True, verbose=True,
+)
+print('TRAIN DONE — running val + test')
+m = YOLO('runs_v3/h100_3class_v3/weights/best.pt')
+print('--- VAL ---'); m.val(data='/home/azureuser/merged_v3/data.yaml', split='val')
+print('--- TEST ---'); m.val(data='/home/azureuser/merged_v3/data.yaml', split='test')
+print('--- TEST + TTA ---'); m.val(data='/home/azureuser/merged_v3/data.yaml', split='test', augment=True)

turboquant_case_study.md

@@ -0,0 +1,72 @@
+# TurboQuant: 44-59% KV-Cache Reduction With Zero Quality Loss
+
+## The Problem
+
+LLM inference is memory-bound. As context length grows, the KV-cache eats your VRAM alive. At 8K tokens, a Gemma-2-9B model burns nearly 4 GB just on KV-cache. That's memory you can't use for batching more requests, which means fewer concurrent users per GPU, which means higher cost per query.
+
+Every production team running LLMs hits this wall.
+
+## What TurboQuant Does
+
+TurboQuant applies mixed-precision quantization to the KV-cache during inference. It profiles each layer's activation norms, identifies outlier layers that need full precision, and quantizes the rest — cutting KV-cache memory by 44-59% while maintaining exact prefill fidelity.
+
+No retraining. No fine-tuning. Drop-in replacement.
+
+## Results
+
+Benchmarked across 5 model families on NVIDIA H100 NVL (96 GB):
+
+### Memory Savings at 8K Context
+
+| Model | Default VRAM | TurboQuant VRAM | KV-Cache Saved | Prefill Fidelity |
+|-------|-------------|----------------|----------------|-----------------|
+| Gemma-2-9B | 9.98 GB | 7.71 GB | 2,323 MB (~59%) | Exact |
+| Qwen2.5-32B | 23.16 GB | 21.41 GB | 1,791 MB (~47%) | Exact |
+| Phi-4-14B | 12.28 GB | 10.92 GB | 1,392 MB (~44%) | Exact |
+| LLaMA-3.1-8B | 7.71 GB | 6.84 GB | 890 MB (~44%) | Exact |
+| Qwen2.5-7B | 7.08 GB | 6.71 GB | 380 MB (~44%) | Exact |
+
+### Quality Verification
+
+- **Prefill logit difference: 0.0 across all models** — the quantized KV-cache produces identical logits at the prefill stage
+- **Same top-1 token prediction: 100%** — no drift in the most likely next token
+- **Output coherence: 100%** — both default and TurboQuant outputs are fully coherent across all test prompts
+- **Token match rate: 18-100%** on generation (expected — autoregressive sampling diverges naturally, but both outputs remain equally valid)
+
+### Scaling With Context Length
+
+Memory savings grow linearly with context. LLaMA-3.1-8B example:
+
+| Context Length | Saved |
+|---------------|-------|
+| 1K tokens | 93 MB |
+| 4K tokens | 417 MB |
+| 8K tokens | 890 MB |
+
+At 32K or 128K context (LLaMA-3.1 supports 128K), the savings become massive — potentially 3-14 GB on a single model.
+
+### Outlier-Aware Design
+
+Not all layers are equal. TurboQuant detects outlier layers with abnormal activation norms and keeps them at full precision:
+
+- **Qwen2.5-7B**: layers 0 and 27 flagged as outliers (norms 273.84 and 239.91 vs median 16.86) — kept at BF16
+- **All other models**: uniform norm distributions, all layers quantized
+
+This is why quality stays intact — the layers that matter most keep their precision.
+
+## What This Means For Production
+
+If you're running LLMs in production:
+
+- **2-3x more concurrent users** on the same GPU (freed VRAM = larger batch sizes)
+- **Same quality** — your users won't notice any difference
+- **No model changes** — works with any transformer architecture using standard KV-cache
+- **Tested across**: Qwen2, LLaMA, Gemma2, Phi3 architectures
+
+## About
+
+Built by Vivek Varikuti. I optimize LLM inference for production workloads.
+
+If your GPU bill is too high or your throughput is too low, I can help. Free 1-week proof-of-concept on your setup — you pay nothing if it doesn't beat your current numbers.
+
+Reach me: domainluther1234@gmail.com | GitHub: vivekvar-dl

vivek_complete_profile.md

@@ -0,0 +1,367 @@
+# Vivek Varikuti — Complete Profile & Project Portfolio
+
+## Who I Am
+
+- 22 years old, AI Engineer & Startup Founder
+- GitHub: vivekvar-dl
+- Email: domainluther1234@gmail.com
+- Strong Python/PyTorch/LLM skills, deep transformer training experience
+- Hardware: 1x NVIDIA H100 NVL 96GB on Azure (NC40ads H100 v5)
+- CUDA 12.8, PyTorch 2.7.0+cu128, flash-attn 2.8.3 (FA2)
+- Transformers 5.4.0
+
+---
+
+## Working Style
+
+- No AI fluff. No menus of options. Make the decision and execute.
+- Write like a human — no perfect grammar, no emojis, no "leveraging" or "seamless"
+- Any public text must read like a tired developer typed it at 2am
+- No co-authored-by Claude in git commits — public contributions look fully human
+- Verify before claiming. Test before shipping. Always run the actual code.
+
+---
+
+## Project 1: TurboQuant — KV Cache Compression
+
+**What:** Implementing Google's TurboQuant paper (arXiv 2504.19874, Zandieh et al.) for KV cache compression during LLM inference.
+
+**Why:** Compress KV cache ~4-7x on production LLMs to enable longer contexts and batching on H100 NVL (96GB).
+
+**Location:** /home/azureuser/turboquant/
+
+**Status:** Working prototype. Google hasn't released their code publicly — this is one of the first working implementations.
+
+**Core Method:** Mixed-precision quantization of KV cache. Profile each layer's activation norms, identify outlier layers that need full precision, quantize the rest. No retraining, no fine-tuning — drop-in replacement.
+
+**Key Discovery:** Layer 0 (and sometimes last layer) of Qwen models have anomalously large key norms (~16-50x median). These layers must be kept in BF16 (skip_layers). Auto-calibration function detects outlier layers.
+
+### Benchmark Results (H100 NVL 96GB)
+
+#### Model Architecture Summary
+
+| Model | Architecture | KV Heads | head_dim | Outlier Layers | Prefill Fidelity |
+|-------|-------------|----------|---------|----------------|-----------------|
+| Qwen2.5-7B | 28L, qwen2 | 4 | 128 | layers 0, 27 | exact |
+| Llama-3.1-8B | 32L, llama | 8 | 128 | none | exact |
+| Gemma-2-9B | 42L, gemma2 | 8 | 256 | none | exact |
+| Phi-4-14B | 40L, phi3 | 10 | 128 | none | exact |
+| Qwen2.5-32B | 64L, qwen2 | 8 | 128 | none | exact |
+| Llama-3.3-70B | 80L, llama | 8 | 128 | N/A (disk full) | N/A |
+
+#### Memory Savings at 8K Context
+
+| Model | Default VRAM | TurboQuant VRAM | Saved | KV Cache Reduction |
+|-------|-------------|----------------|-------|-------------------|
+| Gemma-2-9B | 9.98 GB | 7.71 GB | 2,323 MB | ~59% |
+| Qwen2.5-32B | 23.16 GB | 21.41 GB | 1,791 MB | ~47% |
+| Phi-4-14B | 12.28 GB | 10.92 GB | 1,392 MB | ~44% |
+| LLaMA-3.1-8B | 7.71 GB | 6.84 GB | 890 MB | ~44% |
+| Qwen2.5-7B | 7.08 GB | 6.71 GB | 380 MB | ~44% |
+
+#### Memory Savings Scaling (LLaMA-3.1-8B)
+
+| Context Length | Default VRAM | TurboQuant VRAM | Saved |
+|---------------|-------------|----------------|-------|
+| 1K tokens | 6.00 GB | 5.91 GB | 93 MB |
+| 4K tokens | 6.67 GB | 6.27 GB | 417 MB |
+| 8K tokens | 7.71 GB | 6.84 GB | 890 MB |
+
+#### Full Memory Data Per Model
+
+**Qwen2.5-7B (5.45 GB model)**
+- Layer norms: median 16.86, max 273.84 (layer 0), ratio 16.24x
+- Outlier layers: 0 (norm 273.84), 27 (norm 239.91)
+- 1K: 5.76→5.73 GB (37 MB saved)
+- 4K: 6.27→6.10 GB (176 MB saved)
+- 8K: 7.08→6.71 GB (380 MB saved)
+
+**LLaMA-3.1-8B (5.68 GB model)**
+- Layer norms: median 17.90, max 21.05 (layer 7), ratio 1.18x
+- No outlier layers
+- 1K: 6.00→5.91 GB (93 MB saved, output match)
+- 4K: 6.67→6.27 GB (417 MB saved, output match)
+- 8K: 7.71→6.84 GB (890 MB saved, output match)
+
+**Gemma-2-9B (6.08 GB model)**
+- Layer norms: median 17.82, max 21.28 (layer 25), ratio 1.19x
+- No outlier layers
+- 1K: 6.62→6.38 GB (244 MB saved)
+- 4K: 7.96→6.89 GB (1,096 MB saved)
+- 8K: 9.98→7.71 GB (2,323 MB saved)
+
+**Phi-4-14B (9.10 GB model)**
+- Layer norms: median 19.21, max 26.46 (layer 0), ratio 1.38x
+- No outlier layers
+- 1K: 9.75→9.61 GB (146 MB saved)
+- 4K: 10.72→10.09 GB (650 MB saved)
+- 8K: 12.28→10.92 GB (1,392 MB saved)
+
+**Qwen2.5-32B (19.31 GB model)**
+- Layer norms: median 16.09, max 37.82 (layer 0), ratio 2.35x
+- No outlier layers
+- 1K: 19.97→19.79 GB (186 MB saved)
+- 4K: 21.23→20.42 GB (833 MB saved)
+- 8K: 23.16→21.41 GB (1,791 MB saved)
+
+**LLaMA-3.3-70B** — failed with "No space left on device"
+
+#### Quality Verification
+
+All models tested with 3 prompts: "Explain quantum computing", "Write a Python prime checker", "What causes northern lights?"
+
+- Prefill logit difference: 0.0 across ALL models
+- Same top-1 token prediction: 100% across ALL models
+- Output coherence: 100% — both default and TurboQuant outputs fully coherent
+- Token match rate varies (18-100%) due to natural autoregressive sampling divergence — both outputs equally valid
+
+**Detailed quality per model:**
+
+Qwen2.5-7B: token match 39%, 3%, 54% — both coherent all 3 prompts
+LLaMA-3.1-8B: token match 89.1%, 100%, 100% — 2/3 exact match
+Phi-4-14B: token match 100%, 44%, 100% — 2/3 exact match
+Gemma-2-9B: token match 100%, 100%, 18.8% — 2/3 exact match
+Qwen2.5-32B: token match 71%, 25%, 53% — both coherent all 3 prompts
+
+#### Infrastructure Notes
+- Environment: torch 2.7.0+cu128, transformers 5.4.0, H100 NVL CUDA 12.8 (driver 570)
+- PyTorch compiled for CUDA 13.0+ won't work — need cu128 wheel
+- Core quantizer verified (MSE matches paper bounds)
+- Cache integrates with HF Transformers v5.4.0 QuantizedLayer API
+
+---
+
+## Project 2: Parameter Golf Competition (OpenAI)
+
+**What:** OpenAI competition — train the best language model within a 16MB artifact, 10 minutes on 8xH100.
+
+**Metric:** Bits-per-byte (BPB) on FineWeb validation (62M tokens sp1024, 45.5M tokens sp4096)
+
+**Timeline:** March 18 - April 30, 2026
+
+**Current SOTA (merged):** 1.1194 BPP (PR #549, LeakyReLU^2 + TTT + Parallel Muon)
+
+### Our Edge: sp4096 Vocabulary
+
+- sp4096 tokens_per_byte: 0.3063 vs sp1024: 0.4149 → 26.2% fewer tokens
+- Baseline A/B test (400 steps): sp4096 = 1.6208 BPB vs sp1024 = 1.7144 BPB → -5.5%
+- #1 arch A/B test (400 steps, seed 42): sp4096+factored = 1.8693 BPB vs sp1024 = 2.0067 BPB → -6.8%
+- Extrapolated SOTA: 1.1194 × 0.93 ≈ 1.04-1.06 BPB
+
+### Architecture
+
+- 11L, 512d, 8H, 4KV, 3x MLP, LeakyReLU(0.5)^2
+- Factored embeddings: tok_emb(4096x256) + embed_up(256→512) + embed_down(512→256)
+- All tricks from #1 submission: XSA, Partial RoPE, LN Scale, SmearGate, BigramHash, EMA, TTT, GPTQ-lite
+
+### Key Files
+
+- our_submission/train_gpt.py — modified #1 with sp4096 + factored embed + FA2 fallback
+- our_submission/train_gpt_original.py — unmodified #1 with FA2 fallback
+- train_sp4096.py — tokenizer training + data sharding script
+- data/tokenizers/fineweb_4096_bpe.model — trained sp4096 tokenizer
+- data/datasets/fineweb10B_sp4096/ — 80 train shards + 1 val shard
+
+### N-gram Cache: CONFIRMED FAKE
+
+- 256M bucket experiment: collision-free hash tables give 1.11 BPB (no improvement)
+- All sub-1.0 BPB claims are measurement artifacts from hash collisions
+- Valid Dirichlet smoothing gives at most ~0.002-0.005 genuine improvement
+
+### Next Steps
+
+1. Medium fidelity run (10min 1xH100)
+2. Int5 MLP quantization (saves ~1.86MB for artifact budget headroom)
+3. Get 8xH100 access for final submission (compute grant or RunPod)
+4. Temperature scaling, document-isolated TTT for extra gains
+
+### Hardware
+
+- Dev: 1xH100 NVL (Azure NC40ads H100 v5), 96GB VRAM, CUDA 12.8, PyTorch 2.9.1+cu128
+- flash-attn 2.8.3 (FA2, not FA3)
+- Final submission needs 8xH100
+
+---
+
+## Project 3: GSoC 2026 — DeepChem OLMo Wrapper
+
+**What:** Adding OLMo-2 7B LLM support to DeepChem for molecular property prediction and SMILES generation.
+
+**Org:** DeepChem (standalone first time in GSoC 2026)
+**Mentors:** Riya, Harindhar
+**Deadline:** March 31, 2026 18:00 UTC (submitted)
+
+### What Was Built
+
+**PR #4913 (LIVE) — Bug Fix**
+- Fixed ChemBERTa broken import for transformers 5.x
+- `transformers.models.roberta.tokenization_roberta_fast` removed in 5.x
+- 3 additions / 4 deletions
+- https://github.com/deepchem/deepchem/pull/4913
+
+**Issue #4912 (LIVE) — Compat Report**
+- Broader transformers 5.x compatibility issues documented
+- https://github.com/deepchem/deepchem/issues/4912
+
+**OLMo Wrapper (LOCAL ONLY — not pushed)**
+- Files at ~/olmo_draft/olmo.py and ~/olmo_draft/test_olmo.py
+- Olmo2ForSequenceClassification — built from scratch (doesn't exist in HF)
+- OLMo wrapper class extending HuggingFaceModel
+- Added causal_lm task + generate() to base HuggingFaceModel
+- 8/8 tests pass in 27 seconds on CPU
+- Uses OLMo-2 (allenai/OLMo-2-1124-7B)
+
+### Experiments Run
+
+- BBBP classification: ROC-AUC 0.67 (random init, 12.9M params, 200 samples)
+- ESOL regression: R² 0.37, MAE 1.27
+- SMILES generation: 0% validity (proves pretraining is core challenge)
+- Tokenization analysis: OLMo 0.9x tokens vs ChemBERTa, but fragments stereocenters
+
+### Proposal
+
+- ~/gsoc_proposal_final.md — human-written version
+- ~/gsoc_proposal_content.md — raw technical reference
+
+### Key Context
+
+- PR #4907 by Aditya-ad48 also adds causal LM generation — complementary not competing
+- DeepChem wants small PRs (<50 lines) for new contributors
+- rbharath is the main reviewer/maintainer
+- Office hours MWF 9am PST
+- Discord: https://discord.gg/RYTrUY8Ssn
+
+---
+
+## Project 4: Genesis — Artificial Life Simulation
+
+**What:** Virtual world where blank GRU neural net agents evolve survival behaviors from scratch — foraging, water-seeking, communication — on H100 GPU using JAX.
+
+**Location:** /home/azureuser/genesis/ (venv at ~/genesis_env/)
+
+### World Setup
+
+- 512x512 grid with Perlin noise terrain
+- Food regrowth, water sources, day/night cycles, seasons
+- 1000 agents with GRU brains (~82K params each)
+- Tournament selection + Gaussian mutation (self-adaptive sigma)
+- Agents start with zero knowledge — must learn to survive
+
+### Status (2026-04-01)
+
+Phases 1-3 complete. 500K step run finished successfully:
+- 86 generations evolved
+- Agents sustain avg age 3,742 steps, energy 0.98, hydration 0.79
+- Signal entropy dropping (4.28→3.58) — indicating early communication structure
+- Simulation runs at ~1000 steps/s on H100 (JAX jit-compiled)
+
+### Key Fix
+
+food_growth_rate bumped from 0.005→0.02, food_eat_amount 0.05→0.03 to prevent ecological collapse at high generations.
+
+### Architecture
+
+- World: grid.py, resources.py, environment.py, physics.py, observations.py, spatial.py
+- Agent: body.py (metabolism), brain.py (GRU + vmap batched), actions.py
+- Evolution: fitness.py, selection.py, mutation.py (self-adaptive sigma), population.py
+- Communication: signals.py (8-channel, spatial attenuation, top-4 reception)
+- Analysis: emergence.py (signal entropy, magnitude, R², diversity, clustering)
+- Visualization: renderer.py (dashboard, world map, zoom views)
+
+### Run Data
+
+~/genesis/runs/run_20260401_111309/ — metrics.csv (500 rows), emergence.csv (100 rows), 50 viz frames, 10 checkpoints + FINAL, config.json
+
+### Next Steps
+
+- Phase 4: TRIBE v2 integration — compare evolved GRU representations to human brain activity via RSA
+- Phase 5: Scale to 5K+ agents, longer runs for 500+ generations
+- Checkpoints at 50K intervals allow comparing brain representations across evolutionary time
+
+---
+
+## Project 5: TRIBE v2 — AI-Brain Loop
+
+**What:** Closing the AI-brain comparison loop using Meta's TRIBE v2 — comparing AI encoder representations to predicted brain activity to find architectural gaps.
+
+**Location:** /home/azureuser/tribev2 (venv at /home/azureuser/tribev2_env)
+
+### What's Built
+
+- Full analysis script: /home/azureuser/tribev2/close_the_loop_v2.py
+- 8 phases: load model → extract per-layer features → brain parcellation → layer-wise encoding → modality ablation → RSA → divergence mapping → visualization
+- Multimodal stimulus: /home/azureuser/multimodal_stimulus.mp4
+- Results: /home/azureuser/loop_results_v2/
+- Runs with video (V-JEPA2) + audio (Wav2Vec-BERT) + text (LLaMA 3.2-3B)
+
+### Status
+
+LLaMA 3.2 access granted. Full 3-modality analysis pipeline complete. Brain-guided ViT training attempted 5 times — all failed.
+
+### Why Attempts Failed
+
+- Never had real brain targets — routed ViT-Small features through TRIBE v2's projector (trained for V-JEPA2), producing random outputs
+- Evaluated on wrong metric (classification accuracy instead of robustness)
+- Literature shows brain-guided training helps ROBUSTNESS (+3-8%), not classification accuracy
+
+### What Would Actually Work (from RESEARCH_BRIEF.md)
+
+1. Pre-compute real brain targets using TRIBE v2's full pipeline
+2. Train student with classification + per-vertex Pearson correlation brain loss
+3. Evaluate on corruption/adversarial robustness, shape bias, brain-score — NOT accuracy
+4. Or: use real fMRI data (Natural Scenes Dataset) instead of TRIBE v2 predictions
+
+### Key Infrastructure
+
+- Training scripts: /home/azureuser/brain_guided/train_*.py
+- UCF-101 dataset: /home/azureuser/brain_guided/data/UCF-101 (13K videos)
+- Results: /home/azureuser/brain_guided/results_final/
+
+---
+
+## Project 6: Instagram Cinema
+
+**What:** AI-generated cinematic videos using LTX-2.3 22B on ComfyUI for Instagram growth.
+
+**Setup:** LTX-2.3 22B dev model running on H100 via ComfyUI, exposed via cloudflared tunnel.
+
+**Format:** Instagram Reels — 9:16 portrait, 544x960
+
+**Goal:** Create viral-quality cinematic content for Instagram Reels.
+
+---
+
+## Money-Making Strategy (April 2026)
+
+### Sellable Assets
+
+1. **TurboQuant** — working implementation nobody else has publicly. Lead magnet for consulting.
+2. **Parameter Golf** — competition result (if top placement) = massive credibility signal
+3. **Fine-tuning expertise** — proven on H100, multiple model families
+4. **Inference optimization consulting** — directly from TurboQuant benchmarks
+
+### Immediate Plan
+
+- Path to 10L: Freelancing/consulting — fine-tuning + inference optimization
+- Path to 1Cr: Productized consulting at scale or AI startup
+- Channel: X (Twitter) for distribution, direct DMs to founders for sales
+
+### X (Twitter) Growth Strategy
+
+- Account: 10 followers currently, Premium purchased (213.50/month with 50% off)
+- Strategy: 70% replies (to bigger accounts), 30% original posts
+- Target: 15 strategic replies/day to accounts with 100-5000 followers
+- Post timing: 6:30 PM IST (9:00 AM EST) on Tue/Wed/Thu
+- Pinned thread: TurboQuant benchmarks
+- Goal: 500 followers in 4 weeks, first paid client in 2-4 weeks
+
+### Cold Outreach Template
+
+"I noticed you're using [X model]. I can cut your inference cost by 40%. Free 1-week proof. Interested?"
+
+### Target Clients
+
+- Indian startups using LLMs in production (inc42 AI list)
+- US startups from YC directory (AI/ML category, S24/W25 batches)
+- Anyone on Twitter complaining about GPU costs / inference scaling
+- Companies with >$10K/month GPU spend

yc_scrape.py

@@ -0,0 +1,68 @@
+import requests, json, re, sys
+from openpyxl import Workbook
+
+APP = "45BWZJ1SGC"
+
+def get_key():
+    r = requests.get("https://www.ycombinator.com/companies",
+                     headers={"User-Agent": "Mozilla/5.0"}, timeout=30)
+    m = re.search(r'AlgoliaOpts\s*=\s*(\{[^}]*\})', r.text)
+    return json.loads(m.group(1))["key"]
+
+KEY = get_key()
+URL = f"https://{APP.lower()}-dsn.algolia.net/1/indexes/YCCompany_production/query"
+HDR = {"X-Algolia-Application-Id": APP, "X-Algolia-API-Key": KEY,
+       "Content-Type": "application/json"}
+
+BATCHES = ["Fall 2025", "Winter 2026", "Spring 2026", "Summer 2026"]
+
+def fetch_batch(batch):
+    hits = []
+    page = 0
+    while True:
+        body = {"query": "", "facetFilters": [[f"batch:{batch}"]],
+                "hitsPerPage": 1000, "page": page}
+        r = requests.post(URL, headers=HDR, data=json.dumps(body), timeout=30)
+        d = r.json()
+        hits.extend(d.get("hits", []))
+        if page + 1 >= d.get("nbPages", 0):
+            break
+        page += 1
+    return hits
+
+wb = Workbook()
+ws = wb.active
+ws.title = "YC Startups"
+ws.append(["Name", "Batch", "Website", "One-liner", "Location",
+           "Industry", "Team Size", "Status", "Hiring", "Tags", "YC Page"])
+
+totals = {}
+all_hits = []
+for b in BATCHES:
+    hits = fetch_batch(b)
+    totals[b] = len(hits)
+    print(f"{b}: {len(hits)}", flush=True)
+    all_hits.extend(hits)
+
+for h in all_hits:
+    ws.append([
+        h.get("name", ""),
+        h.get("batch", ""),
+        h.get("website", ""),
+        h.get("one_liner", ""),
+        h.get("all_locations", ""),
+        h.get("industry", ""),
+        h.get("team_size", ""),
+        h.get("status", ""),
+        "Yes" if h.get("isHiring") else "No",
+        ", ".join(h.get("tags", []) or []),
+        f"https://www.ycombinator.com/companies/{h.get('slug','')}",
+    ])
+
+for col, width in enumerate([22, 14, 40, 55, 28, 18, 11, 10, 8, 40, 50], start=1):
+    ws.column_dimensions[chr(64 + col)].width = width
+
+out = "/home/azureuser/yc_companies.xlsx"
+wb.save(out)
+print(f"\nTOTAL: {sum(totals.values())} companies")
+print(f"Saved: {out}")