Upload 27 files

app.py

@@ -0,0 +1,389 @@
+import copy
+import json
+import os
+from datetime import datetime
+
+import gradio as gr
+import torch
+
+from config import CONFIG
+from inference import load_model, run_inference, _build_tokenizers, _resolve_device
+
+
+RESULTS_DIR = "generated_results"
+os.makedirs(RESULTS_DIR, exist_ok=True)
+
+
+def discover_checkpoints():
+    found = []
+    for root in ("ablation_results", "results7", "results"):
+        if not os.path.isdir(root):
+            continue
+        for entry in sorted(os.listdir(root)):
+            ckpt = os.path.join(root, entry, "best_model.pt")
+            if not os.path.exists(ckpt):
+                continue
+            found.append({
+                "label": f"{entry}  [{root}]",
+                "path": ckpt,
+                "experiment": entry,
+                "root": root,
+            })
+    return found
+
+
+def default_checkpoint_label():
+    checkpoints = discover_checkpoints()
+    if not checkpoints:
+        return None
+    for item in checkpoints:
+        if item["path"].endswith("ablation_results/T4/best_model.pt"):
+            return item["label"]
+    return checkpoints[0]["label"]
+
+
+def checkpoint_map():
+    return {item["label"]: item for item in discover_checkpoints()}
+
+
+def infer_model_type(experiment_name: str, root: str = "") -> str:
+    if root == "ablation_results":
+        return "d3pm_cross_attention"
+    if experiment_name.startswith("d3pm_cross_attention"):
+        return "d3pm_cross_attention"
+    if experiment_name.startswith("d3pm_encoder_decoder"):
+        return "d3pm_encoder_decoder"
+    if experiment_name.startswith("baseline_cross_attention"):
+        return "baseline_cross_attention"
+    if experiment_name.startswith("baseline_encoder_decoder"):
+        return "baseline_encoder_decoder"
+    return CONFIG["model_type"]
+
+
+def infer_include_negative(experiment_name: str, root: str = "") -> bool:
+    if root == "ablation_results":
+        return False
+    if "_neg_True" in experiment_name:
+        return True
+    if "_neg_False" in experiment_name:
+        return False
+    return CONFIG["data"]["include_negative_examples"]
+
+
+def build_runtime_cfg(ckpt_path: str):
+    experiment = os.path.basename(os.path.dirname(ckpt_path))
+    root = os.path.basename(os.path.dirname(os.path.dirname(ckpt_path)))
+    cfg = copy.deepcopy(CONFIG)
+    cfg["model_type"] = infer_model_type(experiment, root=root)
+    cfg["data"]["include_negative_examples"] = infer_include_negative(experiment, root=root)
+    if root == "ablation_results" and experiment.startswith("T") and experiment[1:].isdigit():
+        t_val = int(experiment[1:])
+        cfg["model"]["diffusion_steps"] = t_val
+        cfg["inference"]["num_steps"] = t_val
+    device = _resolve_device(cfg)
+    return cfg, device, experiment
+
+
+def load_selected_model(checkpoint_label):
+    mapping = checkpoint_map()
+    if checkpoint_label not in mapping:
+        raise gr.Error("Selected checkpoint was not found. Refresh the dropdown.")
+
+    ckpt_path = mapping[checkpoint_label]["path"]
+    cfg, device, experiment = build_runtime_cfg(ckpt_path)
+    model, cfg = load_model(ckpt_path, cfg, device)
+    src_tok, tgt_tok = _build_tokenizers(cfg)
+
+    bundle = {
+        "ckpt_path": ckpt_path,
+        "experiment": experiment,
+        "device": str(device),
+        "cfg": cfg,
+        "model": model,
+        "src_tok": src_tok,
+        "tgt_tok": tgt_tok,
+    }
+
+    model_info = {
+        "checkpoint": ckpt_path,
+        "experiment": experiment,
+        "model_type": cfg["model_type"],
+        "include_negatives": cfg["data"]["include_negative_examples"],
+        "device": str(device),
+        "max_seq_len": cfg["model"]["max_seq_len"],
+        "diffusion_steps": cfg["model"]["diffusion_steps"],
+        "d_model": cfg["model"]["d_model"],
+        "n_layers": cfg["model"]["n_layers"],
+        "n_heads": cfg["model"]["n_heads"],
+    }
+    status = f"Loaded `{experiment}` on `{device}`."
+    return bundle, status, model_info, cfg["inference"]["num_steps"]
+
+
+def apply_preset(preset_name):
+    presets = {
+        "Manual": (0.70, 40, 1.20, 0.0, 64),
+        "Literal": (0.60, 20, 1.25, 0.0, 64),
+        "Balanced": (0.70, 40, 1.20, 0.0, 64),
+        "Creative": (0.85, 80, 1.20, 0.2, 64),
+    }
+    return presets.get(preset_name, presets["Balanced"])
+
+
+def task_notes_md():
+    return """
+### Task Notes
+
+**Task 1: KV Cache**
+- Benchmark encoder caching vs standard generation.
+- Best for engineering evaluation, not language quality evaluation.
+
+**Task 2: Attention + Drift**
+- Shows internal attention maps and output stabilization over diffusion steps.
+- Useful for diagnostics and mentor discussion of model behavior.
+
+**Task 3: Concept Vectors**
+- Experimental PCA steering over decoder hidden states.
+- Current outputs are exploratory, not strong semantic evidence yet.
+
+**Task 4: Step Ablation**
+- Requires retraining separate checkpoints for each diffusion step count.
+- Use this UI for generation only; ablation analysis runs from `analysis/run_analysis.py`.
+
+**Task 5: Quality Guidance**
+- Advanced experimental feature in the analysis pipeline.
+- Not exposed in this UI because the current evidence is still under validation.
+"""
+
+
+def save_generation(experiment, record):
+    ts = datetime.now().strftime("%Y%m%d")
+    path = os.path.join(RESULTS_DIR, f"{experiment}_ui_{ts}.json")
+    existing = []
+    if os.path.exists(path):
+        with open(path, "r", encoding="utf-8") as f:
+            existing = json.load(f)
+    existing.append(record)
+    with open(path, "w", encoding="utf-8") as f:
+        json.dump(existing, f, ensure_ascii=False, indent=2)
+    return path
+
+
+def clean_generated_text(text: str, max_consecutive: int = 2, max_occurrence_ratio: float = 0.15) -> str:
+    """
+    Lightweight cleanup for repetitive diffusion outputs.
+    Keeps Sanskrit tokens but trims pathological token loops.
+    """
+    text = " ".join(text.split())
+    if not text:
+        return text
+
+    tokens = text.split()
+    cleaned = []
+
+    # 1) Limit consecutive token repetitions.
+    prev = None
+    run = 0
+    for tok in tokens:
+        if tok == prev:
+            run += 1
+        else:
+            prev = tok
+            run = 1
+        if run <= max_consecutive:
+            cleaned.append(tok)
+
+    # 2) Limit global over-dominant tokens (common in collapse cases).
+    if cleaned:
+        max_occ = max(3, int(len(cleaned) * max_occurrence_ratio))
+        counts = {}
+        filtered = []
+        for tok in cleaned:
+            c = counts.get(tok, 0) + 1
+            counts[tok] = c
+            if c <= max_occ:
+                filtered.append(tok)
+        cleaned = filtered
+
+    out = " ".join(cleaned)
+    out = out.replace(" ।", "।").replace(" ॥", "॥")
+    out = " ".join(out.split())
+    return out
+
+
+def generate_from_ui(
+    model_bundle,
+    input_text,
+    temperature,
+    top_k,
+    repetition_penalty,
+    diversity_penalty,
+    num_steps,
+    clean_output,
+):
+    if not model_bundle:
+        raise gr.Error("Load a model first.")
+    if not input_text.strip():
+        raise gr.Error("Enter input text first.")
+
+    cfg = copy.deepcopy(model_bundle["cfg"])
+    cfg["inference"]["temperature"] = float(temperature)
+    cfg["inference"]["top_k"] = int(top_k)
+    cfg["inference"]["repetition_penalty"] = float(repetition_penalty)
+    cfg["inference"]["diversity_penalty"] = float(diversity_penalty)
+    cfg["inference"]["num_steps"] = int(num_steps)
+
+    src_tok = model_bundle["src_tok"]
+    tgt_tok = model_bundle["tgt_tok"]
+    device = torch.device(model_bundle["device"])
+
+    input_ids = torch.tensor(
+        [src_tok.encode(input_text.strip())],
+        dtype=torch.long,
+        device=device,
+    )
+    out = run_inference(model_bundle["model"], input_ids, cfg)
+    clean = [x for x in out[0].tolist() if x > 4]
+    raw_output_text = tgt_tok.decode(clean).strip()
+    output_text = clean_generated_text(raw_output_text) if clean_output else raw_output_text
+    if not output_text:
+        output_text = "(empty output)"
+
+    record = {
+        "timestamp": datetime.now().isoformat(timespec="seconds"),
+        "experiment": model_bundle["experiment"],
+        "checkpoint": model_bundle["ckpt_path"],
+        "input_text": input_text,
+        "raw_output_text": raw_output_text,
+        "output_text": output_text,
+        "clean_output": bool(clean_output),
+        "temperature": float(temperature),
+        "top_k": int(top_k),
+        "repetition_penalty": float(repetition_penalty),
+        "diversity_penalty": float(diversity_penalty),
+        "num_steps": int(num_steps),
+    }
+    log_path = save_generation(model_bundle["experiment"], record)
+    status = f"Generated with `{model_bundle['experiment']}`. Saved to `{log_path}`."
+    return output_text, status, record
+
+
+with gr.Blocks(title="Sanskrit D3PM Studio") as demo:
+    model_state = gr.State(None)
+
+    gr.Markdown(
+        """
+# Sanskrit D3PM Studio
+
+Load any available checkpoint, generate Devanagari output from Roman/IAST Sanskrit,
+and inspect the settings used for evaluation or demos.
+"""
+    )
+
+    with gr.Row():
+        with gr.Column(scale=2):
+            checkpoint_dropdown = gr.Dropdown(
+                label="Available Checkpoints",
+                choices=list(checkpoint_map().keys()),
+                value=default_checkpoint_label(),
+                interactive=True,
+            )
+        with gr.Column(scale=1):
+            refresh_btn = gr.Button("Refresh List")
+            load_btn = gr.Button("Load Model", variant="primary")
+
+    load_status = gr.Markdown("Select a checkpoint and load it.")
+    model_info = gr.JSON(label="Loaded Model Info")
+
+    with gr.Row():
+        with gr.Column(scale=2):
+            input_text = gr.Textbox(
+                label="Input Text (Roman / IAST Sanskrit)",
+                placeholder="dharmo rakṣati rakṣitaḥ",
+                lines=4,
+            )
+            output_text = gr.Textbox(
+                label="Generated Output (Devanagari)",
+                lines=6,
+                interactive=False,
+            )
+            generate_btn = gr.Button("Generate", variant="primary")
+        with gr.Column(scale=1):
+            preset = gr.Radio(
+                ["Manual", "Literal", "Balanced", "Creative"],
+                value="Balanced",
+                label="Inference Preset",
+            )
+            temperature = gr.Slider(0.4, 1.2, value=0.70, step=0.05, label="Temperature")
+            top_k = gr.Slider(5, 100, value=40, step=1, label="Top-K")
+            repetition_penalty = gr.Slider(1.0, 3.0, value=1.20, step=0.05, label="Repetition Penalty")
+            diversity_penalty = gr.Slider(0.0, 1.0, value=0.0, step=0.05, label="Diversity Penalty")
+            num_steps = gr.Slider(1, 128, value=64, step=1, label="Inference Steps")
+            clean_output = gr.Checkbox(value=True, label="Clean Output (dedupe loops)")
+
+    run_status = gr.Markdown("")
+    run_record = gr.JSON(label="Last Generation Metadata")
+
+    with gr.Accordion("Task Details and Evaluation Notes", open=False):
+        task_notes = gr.Markdown(task_notes_md())
+
+    gr.Examples(
+        examples=[
+            ["dharmo rakṣati rakṣitaḥ"],
+            ["satyameva jayate"],
+            ["ahaṃ brahmāsmi"],
+            ["yatra nāryastu pūjyante"],
+        ],
+        inputs=[input_text],
+        label="Quick Examples",
+    )
+
+    def refresh_checkpoints():
+        choices = list(checkpoint_map().keys())
+        value = choices[0] if choices else None
+        return gr.Dropdown(choices=choices, value=value)
+
+    refresh_btn.click(fn=refresh_checkpoints, outputs=[checkpoint_dropdown])
+    load_btn.click(
+        fn=load_selected_model,
+        inputs=[checkpoint_dropdown],
+        outputs=[model_state, load_status, model_info, num_steps],
+    )
+    preset.change(
+        fn=apply_preset,
+        inputs=[preset],
+        outputs=[temperature, top_k, repetition_penalty, diversity_penalty, num_steps],
+    )
+    generate_btn.click(
+        fn=generate_from_ui,
+        inputs=[
+            model_state,
+            input_text,
+            temperature,
+            top_k,
+            repetition_penalty,
+            diversity_penalty,
+            num_steps,
+            clean_output,
+        ],
+        outputs=[output_text, run_status, run_record],
+    )
+    input_text.submit(
+        fn=generate_from_ui,
+        inputs=[
+            model_state,
+            input_text,
+            temperature,
+            top_k,
+            repetition_penalty,
+            diversity_penalty,
+            num_steps,
+            clean_output,
+        ],
+        outputs=[output_text, run_status, run_record],
+    )
+
+
+if __name__ == "__main__":
+    port = int(os.environ["GRADIO_SERVER_PORT"]) if "GRADIO_SERVER_PORT" in os.environ else None
+    demo.launch(server_name="127.0.0.1", server_port=port, share=False)

attention_viz.py

@@ -0,0 +1,379 @@
+"""
+analysis/attention_viz.py
+==========================
+Task 2: Attention weight capture and visualization across diffusion steps.
+
+How it works (no retraining needed):
+  MultiHeadAttention now has two attributes:
+    - capture_weights: bool  — set True to start storing weights
+    - last_attn_weights: Tensor — [B, n_heads, Lq, Lk], updated each forward call
+
+  AttentionCapture:
+    - Sets capture_weights=True on all cross-attention layers
+    - Hooks into generate_cached() to record weights at every diffusion step
+    - Returns a dict: {t_val: [layer_0_weights, layer_1_weights, ...]}
+
+  Visualization:
+    - plot_attn_heatmap(): shows src→tgt alignment at a single step
+    - plot_attn_evolution(): shows how one src→tgt pair evolves over T steps
+    - plot_all_layers(): grid of heatmaps per layer at a given step
+
+Usage:
+    from analysis.attention_viz import AttentionCapture, plot_attn_heatmap
+
+    capturer = AttentionCapture(model)
+    weights  = capturer.capture(src_ids, src_tokens, tgt_tokens)
+    plot_attn_heatmap(weights, step=0, layer=0, src_tokens=..., tgt_tokens=...)
+"""
+
+import torch
+import numpy as np
+import os
+from typing import List, Dict, Optional
+
+
+# ── Attention capture ─────────────────────────────────────────────────
+
+class AttentionCapture:
+    """
+    Captures cross-attention weights from all decoder layers at every
+    diffusion step during generate_cached().
+
+    Works by:
+      1. Setting capture_weights=True on each DecoderBlock.cross_attn
+      2. Running generate_cached() (encoder runs once via KV cache)
+      3. After each denoising step, reading last_attn_weights from each layer
+      4. Storing as {t_val: list_of_layer_weights}
+
+    Zero retraining required — uses the flag added to MultiHeadAttention.
+    """
+
+    def __init__(self, model):
+        """
+        Args:
+            model : SanskritModel wrapper (must be D3PMCrossAttention)
+        """
+        self.model       = model
+        self.inner       = model.model   # D3PMCrossAttention
+        self._cross_attns = []
+
+        # Collect all cross-attention modules from decoder blocks
+        if hasattr(self.inner, 'decoder_blocks'):
+            for block in self.inner.decoder_blocks:
+                if hasattr(block, 'cross_attn'):
+                    self._cross_attns.append(block.cross_attn)
+
+        if not self._cross_attns:
+            raise ValueError(
+                "No cross-attention layers found. "
+                "AttentionCapture only works with D3PMCrossAttention."
+            )
+
+        print(f"AttentionCapture: found {len(self._cross_attns)} cross-attention layers.")
+
+    def _enable(self):
+        """Turn on weight capture for all cross-attention layers."""
+        for ca in self._cross_attns:
+            ca.capture_weights = True
+
+    def _disable(self):
+        """Turn off weight capture (restores zero overhead)."""
+        for ca in self._cross_attns:
+            ca.capture_weights = False
+            ca.last_attn_weights = None
+
+    def _read_weights(self) -> List[np.ndarray]:
+        """
+        Read current last_attn_weights from all layers.
+        Returns list of [B, n_heads, Lq, Lk] arrays — one per layer.
+        Averages over heads to produce [B, Lq, Lk].
+        """
+        weights = []
+        for ca in self._cross_attns:
+            if ca.last_attn_weights is not None:
+                # Average over attention heads → [B, Lq, Lk]
+                w = ca.last_attn_weights.float().mean(dim=1)
+                weights.append(w.numpy())
+        return weights
+
+    @torch.no_grad()
+    def capture(
+        self,
+        src:        torch.Tensor,
+        capture_every: int = 10,
+    ) -> Dict[int, List[np.ndarray]]:
+        """
+        Run full generation while capturing attention at every `capture_every` steps.
+
+        Args:
+            src           : [1, src_len] or [B, src_len] IAST token ids
+            capture_every : capture weights every N steps (default 10)
+                            Use 1 to capture every step (slow, high memory).
+
+        Returns:
+            step_weights : dict mapping t_val → list of [B, Lq, Lk] arrays
+                           one array per decoder layer
+                           keys are t values: T-1, T-1-N, ..., 0
+
+        Example:
+            weights = capturer.capture(src_ids, capture_every=10)
+            # weights[127] = layer weights at t=127 (heavy noise)
+            # weights[0]   = layer weights at t=0   (clean output)
+        """
+        if src.dim() == 1:
+            src = src.unsqueeze(0)
+
+        inner  = self.inner
+        T      = inner.scheduler.num_timesteps
+        device = src.device
+
+        # KV cache: encode source once
+        memory, src_pad_mask = inner.encode_source(src)
+
+        B       = src.shape[0]
+        tgt_len = inner.max_seq_len
+        mask_id = inner.mask_token_id
+
+        x0_est = torch.full((B, tgt_len), mask_id, dtype=torch.long, device=device)
+        hint   = None
+
+        step_weights: Dict[int, List[np.ndarray]] = {}
+
+        self._enable()
+        try:
+            inner.eval()
+            for t_val in range(T - 1, -1, -1):
+                t       = torch.full((B,), t_val, dtype=torch.long, device=device)
+                is_last = (t_val == 0)
+
+                logits, _ = inner.forward_cached(
+                    memory, src_pad_mask, x0_est, t,
+                    x0_hint=hint, inference_mode=True,
+                )
+
+                # Capture at this step if scheduled or it's the last step
+                if (T - 1 - t_val) % capture_every == 0 or is_last:
+                    step_weights[t_val] = self._read_weights()
+
+                import torch.nn.functional as F
+                probs  = F.softmax(logits / 0.8, dim=-1)
+                x0_est = torch.argmax(probs, dim=-1) if is_last else \
+                         _multinomial_sample(probs)
+                hint   = x0_est
+
+        finally:
+            self._disable()   # always restore — even if exception raised
+
+        print(f"Captured attention at {len(step_weights)} steps "
+              f"({len(self._cross_attns)} layers each).")
+        return step_weights
+
+
+def _multinomial_sample(probs: torch.Tensor) -> torch.Tensor:
+    B, L, V = probs.shape
+    flat    = probs.view(B * L, V).clamp(min=1e-9)
+    flat    = flat / flat.sum(dim=-1, keepdim=True)
+    return torch.multinomial(flat, 1).squeeze(-1).view(B, L)
+
+
+# ── Visualization ─────────────────────────────────────────────────────
+
+def plot_attn_heatmap(
+    step_weights:  Dict[int, List[np.ndarray]],
+    t_val:         int,
+    layer:         int,
+    src_tokens:    List[str],
+    tgt_tokens:    List[str],
+    sample_idx:    int  = 0,
+    save_path:     Optional[str] = None,
+    title:         Optional[str] = None,
+):
+    """
+    Plot cross-attention heatmap for a single step and layer.
+
+    X-axis = source (IAST) tokens
+    Y-axis = target (Devanagari) positions
+    Color  = attention weight (brighter = stronger attention)
+
+    Args:
+        step_weights : output of AttentionCapture.capture()
+        t_val        : which diffusion step to visualize
+        layer        : which decoder layer (0 = first, -1 = last)
+        src_tokens   : list of IAST token strings for x-axis labels
+        tgt_tokens   : list of Devanagari token strings for y-axis labels
+        sample_idx   : which batch item to visualize (default 0)
+        save_path    : if given, save figure to this path
+        title        : custom plot title
+    """
+    try:
+        import matplotlib.pyplot as plt
+        import matplotlib.ticker as ticker
+    except ImportError:
+        print("pip install matplotlib to use visualization functions.")
+        return
+
+    if t_val not in step_weights:
+        available = sorted(step_weights.keys())
+        raise ValueError(
+            f"t_val={t_val} not in captured steps. "
+            f"Available: {available[:5]}...{available[-5:]}"
+        )
+
+    layers  = step_weights[t_val]
+    weights = layers[layer][sample_idx]   # [Lq, Lk]
+
+    # Trim to actual token lengths
+    n_src = min(len(src_tokens), weights.shape[1])
+    n_tgt = min(len(tgt_tokens), weights.shape[0])
+    weights = weights[:n_tgt, :n_src]
+
+    fig, ax = plt.subplots(figsize=(max(8, n_src * 0.4), max(6, n_tgt * 0.35)))
+    im = ax.imshow(weights, aspect='auto', cmap='YlOrRd', interpolation='nearest')
+
+    ax.set_xticks(range(n_src))
+    ax.set_xticklabels(src_tokens[:n_src], rotation=45, ha='right', fontsize=9)
+    ax.set_yticks(range(n_tgt))
+    ax.set_yticklabels(tgt_tokens[:n_tgt], fontsize=9)
+
+    ax.set_xlabel("Source (IAST)", fontsize=11)
+    ax.set_ylabel("Target position (Devanagari)", fontsize=11)
+
+    plot_title = title or f"Cross-Attention  |  t={t_val}  |  Layer {layer}"
+    ax.set_title(plot_title, fontsize=12, pad=10)
+
+    plt.colorbar(im, ax=ax, label="Attention weight")
+    plt.tight_layout()
+
+    if save_path:
+        os.makedirs(os.path.dirname(save_path) or ".", exist_ok=True)
+        plt.savefig(save_path, dpi=150, bbox_inches='tight')
+        print(f"Saved: {save_path}")
+    else:
+        plt.show()
+    plt.close()
+
+
+def plot_attn_evolution(
+    step_weights:  Dict[int, List[np.ndarray]],
+    src_token_idx: int,
+    tgt_token_idx: int,
+    layer:         int = -1,
+    sample_idx:    int = 0,
+    src_token_str: str = "",
+    tgt_token_str: str = "",
+    save_path:     Optional[str] = None,
+):
+    """
+    Plot how attention between one specific src↔tgt token pair evolves
+    across all captured diffusion steps (T → 0).
+
+    Reveals whether a token pair is 'locked' (stable from early steps)
+    or 'flexible' (weight fluctuates until final steps).
+
+    Args:
+        step_weights  : output of AttentionCapture.capture()
+        src_token_idx : index of source token to track
+        tgt_token_idx : index of target position to track
+        layer         : decoder layer index
+        sample_idx    : batch item
+        src_token_str : string label for the source token (for plot title)
+        tgt_token_str : string label for the target token (for plot title)
+        save_path     : if given, save figure to this path
+    """
+    try:
+        import matplotlib.pyplot as plt
+    except ImportError:
+        print("pip install matplotlib to use visualization functions.")
+        return
+
+    t_vals  = sorted(step_weights.keys(), reverse=True)  # T-1 → 0
+    weights = []
+
+    for t_val in t_vals:
+        layers = step_weights[t_val]
+        w      = layers[layer][sample_idx]   # [Lq, Lk]
+        if tgt_token_idx < w.shape[0] and src_token_idx < w.shape[1]:
+            weights.append(w[tgt_token_idx, src_token_idx])
+        else:
+            weights.append(0.0)
+
+    fig, ax = plt.subplots(figsize=(12, 4))
+    ax.plot(range(len(t_vals)), weights, linewidth=1.5, color='steelblue')
+    ax.fill_between(range(len(t_vals)), weights, alpha=0.2, color='steelblue')
+
+    # Mark every 10th step on x-axis
+    step_labels = [str(t) if i % max(1, len(t_vals)//10) == 0 else ""
+                   for i, t in enumerate(t_vals)]
+    ax.set_xticks(range(len(t_vals)))
+    ax.set_xticklabels(step_labels, fontsize=8)
+    ax.set_xlabel("Diffusion step (T → 0)", fontsize=11)
+    ax.set_ylabel("Attention weight", fontsize=11)
+
+    pair_str = f"src[{src_token_idx}]={src_token_str!r} → tgt[{tgt_token_idx}]={tgt_token_str!r}"
+    ax.set_title(f"Attention evolution  |  {pair_str}  |  Layer {layer}", fontsize=11)
+    ax.set_xlim(0, len(t_vals) - 1)
+    ax.set_ylim(0, None)
+    plt.tight_layout()
+
+    if save_path:
+        os.makedirs(os.path.dirname(save_path) or ".", exist_ok=True)
+        plt.savefig(save_path, dpi=150, bbox_inches='tight')
+        print(f"Saved: {save_path}")
+    else:
+        plt.show()
+    plt.close()
+
+
+def plot_all_layers(
+    step_weights: Dict[int, List[np.ndarray]],
+    t_val:        int,
+    src_tokens:   List[str],
+    tgt_tokens:   List[str],
+    sample_idx:   int          = 0,
+    save_path:    Optional[str] = None,
+):
+    """
+    Plot attention heatmaps for ALL decoder layers at a single diffusion step.
+    Shows how different layers specialize their attention patterns.
+    """
+    try:
+        import matplotlib.pyplot as plt
+    except ImportError:
+        print("pip install matplotlib to use visualization functions.")
+        return
+
+    layers  = step_weights[t_val]
+    n_layers = len(layers)
+    n_cols   = min(4, n_layers)
+    n_rows   = (n_layers + n_cols - 1) // n_cols
+
+    fig, axes = plt.subplots(n_rows, n_cols,
+                              figsize=(n_cols * 5, n_rows * 4))
+    axes = np.array(axes).flatten() if n_layers > 1 else [axes]
+
+    n_src = min(len(src_tokens), layers[0][sample_idx].shape[1])
+    n_tgt = min(len(tgt_tokens), layers[0][sample_idx].shape[0])
+
+    for i, (ax, layer_w) in enumerate(zip(axes, layers)):
+        w  = layer_w[sample_idx][:n_tgt, :n_src]
+        im = ax.imshow(w, aspect='auto', cmap='YlOrRd', interpolation='nearest',
+                       vmin=0, vmax=w.max())
+        ax.set_title(f"Layer {i}", fontsize=10)
+        ax.set_xticks(range(n_src))
+        ax.set_xticklabels(src_tokens[:n_src], rotation=45, ha='right', fontsize=7)
+        ax.set_yticks(range(n_tgt))
+        ax.set_yticklabels(tgt_tokens[:n_tgt], fontsize=7)
+
+    for ax in axes[n_layers:]:
+        ax.set_visible(False)
+
+    fig.suptitle(f"All layers at t={t_val}", fontsize=13, y=1.02)
+    plt.tight_layout()
+
+    if save_path:
+        os.makedirs(os.path.dirname(save_path) or ".", exist_ok=True)
+        plt.savefig(save_path, dpi=150, bbox_inches='tight')
+        print(f"Saved: {save_path}")
+    else:
+        plt.show()
+    plt.close()

best_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b1baa03568c2bed42621da115f6e6971411b59cc9dec6b58cf8f2ed87ba2e770
+size 1077681643

concept_vectors.py

@@ -0,0 +1,368 @@
+"""
+analysis/concept_vectors.py
+============================
+Task 3: Concept Vector Extraction + Controlled Paraphrase Diversity
+
+No retraining required. Uses decoder hidden states already computed
+during generate_cached() — stored in model.model._last_hidden after
+each forward_cached() call.
+
+Steps:
+  1. Collect hidden states from N examples at a fixed diffusion step
+  2. Pool sequence dimension → [N, d_model] representation per example
+  3. PCA → find principal directions in concept space
+  4. Identify "diversity direction" (PC that best separates short/long outputs)
+  5. Steer: at inference, shift hidden states along diversity direction
+     before the output head projection
+  6. Generate at 5 points along the direction, measure output diversity
+
+Key insight: the diversity direction is found purely from model outputs
+(no human annotation needed). We use output length as a proxy:
+  short output  → low diversity (model collapsed to simple token)
+  long output   → high diversity (model exploring more of the space)
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from typing import List, Dict, Optional, Tuple
+
+
+# ── Hidden state collection ───────────────────────────────────────────
+
+@torch.no_grad()
+def collect_hidden_states(
+    model,
+    src_list:    List[torch.Tensor],
+    t_capture:   int   = 0,
+    temperature: float = 0.8,
+    top_k:       int   = 40,
+    max_samples: int   = 1000,
+) -> Tuple[np.ndarray, List[str]]:
+    """
+    Run generate_cached() on a list of source tensors, collecting the
+    decoder hidden state at timestep t_capture for each sample.
+
+    Args:
+        model      : SanskritModel (D3PMCrossAttention)
+        src_list   : list of [1, src_len] tensors, one per sample
+        t_capture  : which diffusion step to capture hidden states at
+                     0 = final (clean), T-1 = noisy start
+        temperature: sampling temperature
+        top_k      : top-k filter
+        max_samples: cap at this many samples
+
+    Returns:
+        hidden_matrix : np.ndarray [N, d_model] — pooled hidden states
+        output_texts  : list of N decoded output strings (for diversity analysis)
+    """
+    inner   = model.model
+    T       = inner.scheduler.num_timesteps
+    device  = next(inner.parameters()).device
+
+    hidden_list  = []
+    output_list  = []
+
+    n = min(len(src_list), max_samples)
+    print(f"Collecting hidden states from {n} examples at t={t_capture}...")
+
+    for i, src in enumerate(src_list[:n]):
+        if i % 100 == 0:
+            print(f"  {i}/{n}")
+
+        if src.dim() == 1:
+            src = src.unsqueeze(0)
+        src = src.to(device)
+
+        B       = src.shape[0]
+        tgt_len = inner.max_seq_len
+        mask_id = inner.mask_token_id
+
+        # KV cache
+        memory, src_pad_mask = inner.encode_source(src)
+
+        x0_est = torch.full((B, tgt_len), mask_id, dtype=torch.long, device=device)
+        hint   = None
+        captured_hidden = None
+
+        for t_val in range(T - 1, -1, -1):
+            t       = torch.full((B,), t_val, dtype=torch.long, device=device)
+            is_last = (t_val == 0)
+
+            logits, _ = inner.forward_cached(
+                memory, src_pad_mask, x0_est, t,
+                x0_hint=hint, inference_mode=True,
+            )
+
+            # Capture hidden state at target step
+            if t_val == t_capture and hasattr(inner, '_last_hidden'):
+                captured_hidden = inner._last_hidden.detach().cpu()
+
+            logits = logits / max(temperature, 1e-8)
+            if top_k > 0:
+                V = logits.shape[-1]
+                if top_k < V:
+                    vals, _ = torch.topk(logits, top_k, dim=-1)
+                    logits  = logits.masked_fill(logits < vals[..., -1:], float('-inf'))
+
+            probs  = F.softmax(logits, dim=-1)
+            x0_est = torch.argmax(probs, dim=-1) if is_last else _sample(probs)
+            hint   = x0_est
+
+        # Pool hidden state over non-PAD positions → [d_model]
+        if captured_hidden is not None:
+            non_pad = (x0_est[0] > 1).cpu()   # [tgt_len] bool
+            if non_pad.sum() > 0:
+                h = captured_hidden[0][non_pad].mean(dim=0)   # [d_model]
+            else:
+                h = captured_hidden[0].mean(dim=0)
+            hidden_list.append(h.numpy())
+
+        # Decode output
+        ids  = [x for x in x0_est[0].tolist() if x > 4]
+
+    print(f"Collected {len(hidden_list)} hidden states.")
+    return np.stack(hidden_list), output_list
+
+
+# ── PCA on hidden states ──────────────────────────────────────────────
+
+def fit_pca(
+    hidden_matrix: np.ndarray,
+    n_components:  int = 50,
+) -> object:
+    """
+    Fit PCA on hidden state matrix.
+
+    Args:
+        hidden_matrix : [N, d_model]
+        n_components  : number of PCA components to retain
+
+    Returns:
+        fitted sklearn PCA object
+    """
+    from sklearn.decomposition import PCA
+    n_comp = min(n_components, hidden_matrix.shape[0] - 1, hidden_matrix.shape[1])
+    pca    = PCA(n_components=n_comp)
+    pca.fit(hidden_matrix)
+    print(f"PCA fit: {n_comp} components explain "
+          f"{pca.explained_variance_ratio_.sum()*100:.1f}% of variance.")
+    return pca
+
+
+def find_diversity_direction(
+    hidden_matrix: np.ndarray,
+    output_lengths: List[int],
+    pca:           object,
+) -> np.ndarray:
+    """
+    Find the PCA direction that best correlates with output diversity
+    (measured by output length as proxy).
+
+    Projects hidden states into PCA space, then finds the PC whose
+    scores have highest Spearman correlation with output lengths.
+
+    Returns:
+        direction : np.ndarray [d_model] — diversity direction in original space
+    """
+    from scipy.stats import spearmanr
+
+    projected = pca.transform(hidden_matrix)   # [N, n_components]
+    lengths   = np.array(output_lengths)
+
+    correlations = []
+    for pc_idx in range(projected.shape[1]):
+        r, _ = spearmanr(projected[:, pc_idx], lengths)
+        correlations.append(abs(r))
+
+    best_pc = int(np.argmax(correlations))
+    print(f"Diversity direction: PC {best_pc}  "
+          f"(|r|={correlations[best_pc]:.3f} with output length)")
+
+    # Map back to original d_model space
+    direction = pca.components_[best_pc]   # [d_model]
+    direction = direction / (np.linalg.norm(direction) + 1e-8)
+    return direction, best_pc, correlations[best_pc]
+
+
+# ── Steered generation ────────────────────────────────────────────────
+
+@torch.no_grad()
+def generate_steered(
+    model,
+    src:       torch.Tensor,
+    direction: np.ndarray,
+    alpha:     float = 0.0,
+    temperature: float = 0.8,
+    top_k:     int   = 40,
+) -> torch.Tensor:
+    """
+    Generate output while steering hidden states along diversity direction.
+
+    At each diffusion step, after the decoder runs, we shift the hidden state
+    by alpha * direction before projecting to logits.
+
+    alpha > 0 → push toward high-diversity output
+    alpha < 0 → push toward low-diversity output
+    alpha = 0 → standard generation (no steering)
+
+    Args:
+        model     : SanskritModel (D3PMCrossAttention)
+        src       : [1, src_len] IAST token ids
+        direction : [d_model] diversity direction from find_diversity_direction()
+        alpha     : steering strength
+        temperature / top_k: sampling params
+
+    Returns:
+        x0_est : [1, tgt_len] generated token ids
+    """
+    inner   = model.model
+    T       = inner.scheduler.num_timesteps
+    device  = next(inner.parameters()).device
+
+    if src.dim() == 1:
+        src = src.unsqueeze(0)
+    src = src.to(device)
+
+    B       = src.shape[0]
+    tgt_len = inner.max_seq_len
+    mask_id = inner.mask_token_id
+
+    dir_tensor = torch.tensor(direction, dtype=torch.float32, device=device)
+
+    memory, src_pad_mask = inner.encode_source(src)
+    x0_est = torch.full((B, tgt_len), mask_id, dtype=torch.long, device=device)
+    hint   = None
+
+    inner.eval()
+    for t_val in range(T - 1, -1, -1):
+        t       = torch.full((B,), t_val, dtype=torch.long, device=device)
+        is_last = (t_val == 0)
+
+        # Standard forward_cached but we intercept hidden states
+        PAD = 1
+        tgt_pad_mask = None  # inference_mode
+
+        _, x_t_ids = inner.forward_process.q_sample(x0_est, t) if t_val > 0 else \
+                     (None, x0_est)
+        x      = inner.tgt_embed(x_t_ids)
+        t_norm = t.float() / inner.scheduler.num_timesteps
+        t_emb  = inner.time_mlp(t_norm.unsqueeze(-1))
+        x      = x + t_emb.unsqueeze(1)
+
+        if hint is not None:
+            hint_emb = inner.tgt_embed(hint)
+            gate     = inner.hint_gate(x)
+            x        = x + gate * hint_emb
+
+        for block in inner.decoder_blocks:
+            x = block(x, memory, tgt_pad_mask=tgt_pad_mask, src_pad_mask=src_pad_mask)
+
+        # ── STEERING: shift hidden states along diversity direction ───
+        if alpha != 0.0:
+            x = x + alpha * dir_tensor.unsqueeze(0).unsqueeze(0)
+
+        # Project to logits using the head
+        logits = inner.head(x)
+
+        logits = logits / max(temperature, 1e-8)
+        if top_k > 0:
+            V = logits.shape[-1]
+            if top_k < V:
+                vals, _ = torch.topk(logits, top_k, dim=-1)
+                logits  = logits.masked_fill(logits < vals[..., -1:], float('-inf'))
+
+        probs  = F.softmax(logits, dim=-1)
+        x0_est = torch.argmax(probs, dim=-1) if is_last else _sample(probs)
+        hint   = x0_est
+
+    return x0_est
+
+
+def generate_diversity_spectrum(
+    model,
+    src:           torch.Tensor,
+    direction:     np.ndarray,
+    tgt_tokenizer,
+    alphas:        List[float] = [-2.0, -1.0, 0.0, 1.0, 2.0],
+    temperature:   float       = 0.8,
+    top_k:         int         = 40,
+) -> Dict[float, str]:
+    """
+    Generate outputs at 5 points along the diversity direction.
+
+    Args:
+        alphas : steering strengths (negative = low diversity, positive = high)
+
+    Returns:
+        dict mapping alpha → decoded Devanagari string
+    """
+    results = {}
+    for alpha in alphas:
+        out_ids  = generate_steered(model, src, direction, alpha, temperature, top_k)
+        ids      = [x for x in out_ids[0].tolist() if x > 4]
+        text     = tgt_tokenizer.decode(ids).strip()
+        results[alpha] = text
+        print(f"  alpha={alpha:+.1f}  → {text}")
+    return results
+
+
+def plot_pca_space(
+    hidden_matrix:  np.ndarray,
+    output_lengths: List[int],
+    pca:            object,
+    diversity_pc:   int,
+    save_path:      Optional[str] = None,
+):
+    """
+    Scatter plot of examples in PC1 vs PC2 space, coloured by output length.
+    Highlights the diversity direction.
+    """
+    try:
+        import matplotlib.pyplot as plt
+    except ImportError:
+        print("pip install matplotlib.")
+        return
+
+    projected = pca.transform(hidden_matrix)   # [N, n_pc]
+    lengths   = np.array(output_lengths)
+
+    fig, axes = plt.subplots(1, 2, figsize=(14, 5))
+
+    # Left: PC0 vs PC1 coloured by length
+    ax = axes[0]
+    sc = ax.scatter(projected[:, 0], projected[:, 1],
+                    c=lengths, cmap='viridis', alpha=0.6, s=15)
+    plt.colorbar(sc, ax=ax, label="Output length (chars)")
+    ax.set_xlabel(f"PC0 ({pca.explained_variance_ratio_[0]*100:.1f}%)", fontsize=10)
+    ax.set_ylabel(f"PC1 ({pca.explained_variance_ratio_[1]*100:.1f}%)", fontsize=10)
+    ax.set_title("Concept space (PC0 vs PC1)", fontsize=11)
+
+    # Right: explained variance
+    ax2 = axes[1]
+    cumvar = np.cumsum(pca.explained_variance_ratio_) * 100
+    ax2.plot(range(1, len(cumvar)+1), cumvar, linewidth=1.5, color='steelblue')
+    ax2.axvline(diversity_pc, color='coral', linestyle='--', label=f"Diversity PC={diversity_pc}")
+    ax2.set_xlabel("Number of PCs", fontsize=10)
+    ax2.set_ylabel("Cumulative variance (%)", fontsize=10)
+    ax2.set_title("PCA explained variance", fontsize=11)
+    ax2.legend()
+    ax2.set_ylim(0, 102)
+
+    plt.tight_layout()
+    if save_path:
+        import os
+        os.makedirs(os.path.dirname(save_path) or ".", exist_ok=True)
+        plt.savefig(save_path, dpi=150, bbox_inches='tight')
+        print(f"Saved: {save_path}")
+    else:
+        plt.show()
+    plt.close()
+
+
+def _sample(probs):
+    B, L, V = probs.shape
+    flat    = probs.view(B * L, V).clamp(min=1e-9)
+    flat    = flat / flat.sum(dim=-1, keepdim=True)
+    return torch.multinomial(flat, 1).squeeze(-1).view(B, L)

config_T16.py

@@ -0,0 +1,119 @@
+# Ablation config: T=16 diffusion steps
+import os
+import torch
+
+
+def _get_env_int(name, default):
+    value = os.environ.get(name)
+    return int(value) if value is not None else default
+
+
+def _get_env_str(name, default):
+    return os.environ.get(name, default)
+
+
+# 🎛️ BASH-CONTROLLED SWITCHES (Defaults if run manually)
+MODEL = os.environ.get("MODEL_TYPE", "d3pm_encoder_decoder")
+NEGATIVES = os.environ.get("INCLUDE_NEG", "False") == "True"
+DIFFUSION_STEPS = _get_env_int("DIFFUSION_STEPS", 128)
+INFERENCE_STEPS = _get_env_int("INFERENCE_NUM_STEPS", min(64, DIFFUSION_STEPS))
+TRAIN_DEVICE = _get_env_str(
+    "TRAIN_DEVICE",
+    "mps" if torch.backends.mps.is_available() else "cpu",
+)
+
+CONFIG = {
+    "model_type": MODEL,
+
+    "data": {
+        "include_negative_examples": NEGATIVES,
+        "dataset_size": 60000,
+    },
+
+    # "model": {
+    #     "vocab_size": 16000,
+    #     "max_seq_len": 80,
+    #     "diffusion_steps": 10,
+    #     "d_model": 384,
+    #     "n_layers": 6,
+    #     "n_heads": 6,
+    #     "d_ff": 1536,
+    #     "dropout": 0.15
+    # },
+    #
+    # "diffusion": {
+    #     "mask_token_id": 0
+    # },
+    #
+    # "training": {
+    #     "batch_size": 32,
+    #     "epochs": 10,
+    #     "lr": 2e-4,
+    #     "label_smoothing": 0.05,
+    #     "precision": "float32",
+    #     "device": "mps" if torch.backends.mps.is_available() else "cpu",
+    #     "early_stopping_patience": 3
+    # }
+# "model": {
+#         "vocab_size": 16000,
+#         "max_seq_len": 96,       # Optimized for GRETIL slokas
+#         "diffusion_steps": 16,   # Use 16 steps (better than 8)
+#         "d_model": 512,          # Wider model learns faster
+#         "n_layers": 8,
+#         "n_heads": 8,
+#         "d_ff": 2048,
+#         "dropout": 0.1
+#     },
+#
+#     "diffusion": {
+#         "mask_token_id": 0
+#     },
+#
+#     "training": {
+#         "batch_size": 32,
+#         "epochs": 20,            # 20 is enough with these tweaks
+#         "lr": 4e-4,              # Higher LR + Warmup for speed
+#         "label_smoothing": 0.15, # Increased for 16k vocab stability
+#         "precision": "float32",
+#         "device": "mps" if torch.backends.mps.is_available() else "cpu",
+#         "early_stopping_patience": 5
+#     }
+'diffusion': {
+        'mask_token_id': 0,          # [MASK] = ID 0, fixed by tokenizer
+    },
+
+    # ── Model architecture ────────────────────────────────────────────
+    'model': {
+        # 'vocab_size':       16000,
+'src_vocab_size': 16000,   # Roman/IAST BPE vocab
+'tgt_vocab_size': 16000,   # Devanagari BPE vocab
+        'd_model':          1024,#512,     # was 384 — kept same, shared embeds save params
+        'n_heads':          8,       # 384 / 6 = 64 head_dim
+        'd_ff':            4096, #2048, #1536,    # 4 × d_model
+        'n_layers':         8,#4,
+        'dropout':          0.2,
+        'max_seq_len':      80,
+        'diffusion_steps':  DIFFUSION_STEPS,
+    },
+
+    # ── Training ──────────────────────────────────────────────────────
+    'training': {
+        'epochs':           20,       # Target: 0.71→0.83-0.85 in 5 epochs
+        'batch_size':       32,
+        'accum_steps':      2,       # effective batch = 64
+        'lr':               7e-5,#6e-4,    # raised from 3e-4; warmup protects first steps
+        'label_smoothing':  0.1,     # was 0.0; reduces overconfidence (gap 1.7 nats)
+        'patience':         4,       # early stop after 4 non-improving epochs
+        'l1_lambda':        1e-7,    # very light L1
+        'device':           TRAIN_DEVICE,
+    },
+
+    # ── Inference (used during val BERTScore and generate()) ──────────
+    'inference': {
+        'num_steps':          INFERENCE_STEPS,
+        'temperature':        0.7,   # slightly lower = more confident output
+        'top_k':              40,
+        'repetition_penalty': 1.2,
+        'diversity_penalty':  0.5,   # keep off; global-mean penalty is conservative
+    },
+}

config_T32.py

@@ -0,0 +1,119 @@
+# Ablation config: T=32 diffusion steps
+import os
+import torch
+
+
+def _get_env_int(name, default):
+    value = os.environ.get(name)
+    return int(value) if value is not None else default
+
+
+def _get_env_str(name, default):
+    return os.environ.get(name, default)
+
+
+# 🎛️ BASH-CONTROLLED SWITCHES (Defaults if run manually)
+MODEL = os.environ.get("MODEL_TYPE", "d3pm_encoder_decoder")
+NEGATIVES = os.environ.get("INCLUDE_NEG", "False") == "True"
+DIFFUSION_STEPS = _get_env_int("DIFFUSION_STEPS", 128)
+INFERENCE_STEPS = _get_env_int("INFERENCE_NUM_STEPS", min(64, DIFFUSION_STEPS))
+TRAIN_DEVICE = _get_env_str(
+    "TRAIN_DEVICE",
+    "mps" if torch.backends.mps.is_available() else "cpu",
+)
+
+CONFIG = {
+    "model_type": MODEL,
+
+    "data": {
+        "include_negative_examples": NEGATIVES,
+        "dataset_size": 60000,
+    },
+
+    # "model": {
+    #     "vocab_size": 16000,
+    #     "max_seq_len": 80,
+    #     "diffusion_steps": 10,
+    #     "d_model": 384,
+    #     "n_layers": 6,
+    #     "n_heads": 6,
+    #     "d_ff": 1536,
+    #     "dropout": 0.15
+    # },
+    #
+    # "diffusion": {
+    #     "mask_token_id": 0
+    # },
+    #
+    # "training": {
+    #     "batch_size": 32,
+    #     "epochs": 10,
+    #     "lr": 2e-4,
+    #     "label_smoothing": 0.05,
+    #     "precision": "float32",
+    #     "device": "mps" if torch.backends.mps.is_available() else "cpu",
+    #     "early_stopping_patience": 3
+    # }
+# "model": {
+#         "vocab_size": 16000,
+#         "max_seq_len": 96,       # Optimized for GRETIL slokas
+#         "diffusion_steps": 16,   # Use 16 steps (better than 8)
+#         "d_model": 512,          # Wider model learns faster
+#         "n_layers": 8,
+#         "n_heads": 8,
+#         "d_ff": 2048,
+#         "dropout": 0.1
+#     },
+#
+#     "diffusion": {
+#         "mask_token_id": 0
+#     },
+#
+#     "training": {
+#         "batch_size": 32,
+#         "epochs": 20,            # 20 is enough with these tweaks
+#         "lr": 4e-4,              # Higher LR + Warmup for speed
+#         "label_smoothing": 0.15, # Increased for 16k vocab stability
+#         "precision": "float32",
+#         "device": "mps" if torch.backends.mps.is_available() else "cpu",
+#         "early_stopping_patience": 5
+#     }
+'diffusion': {
+        'mask_token_id': 0,          # [MASK] = ID 0, fixed by tokenizer
+    },
+
+    # ── Model architecture ────────────────────────────────────────────
+    'model': {
+        # 'vocab_size':       16000,
+'src_vocab_size': 16000,   # Roman/IAST BPE vocab
+'tgt_vocab_size': 16000,   # Devanagari BPE vocab
+        'd_model':          1024,#512,     # was 384 — kept same, shared embeds save params
+        'n_heads':          8,       # 384 / 6 = 64 head_dim
+        'd_ff':            4096, #2048, #1536,    # 4 × d_model
+        'n_layers':         8,#4,
+        'dropout':          0.2,
+        'max_seq_len':      80,
+        'diffusion_steps':  DIFFUSION_STEPS,
+    },
+
+    # ── Training ──────────────────────────────────────────────────────
+    'training': {
+        'epochs':           20,       # Target: 0.71→0.83-0.85 in 5 epochs
+        'batch_size':       32,
+        'accum_steps':      2,       # effective batch = 64
+        'lr':               7e-5,#6e-4,    # raised from 3e-4; warmup protects first steps
+        'label_smoothing':  0.1,     # was 0.0; reduces overconfidence (gap 1.7 nats)
+        'patience':         4,       # early stop after 4 non-improving epochs
+        'l1_lambda':        1e-7,    # very light L1
+        'device':           TRAIN_DEVICE,
+    },
+
+    # ── Inference (used during val BERTScore and generate()) ──────────
+    'inference': {
+        'num_steps':          INFERENCE_STEPS,
+        'temperature':        0.7,   # slightly lower = more confident output
+        'top_k':              40,
+        'repetition_penalty': 1.2,
+        'diversity_penalty':  0.5,   # keep off; global-mean penalty is conservative
+    },
+}

config_T4.py

@@ -0,0 +1,119 @@
+# Ablation config: T=4 diffusion steps
+import os
+import torch
+
+
+def _get_env_int(name, default):
+    value = os.environ.get(name)
+    return int(value) if value is not None else default
+
+
+def _get_env_str(name, default):
+    return os.environ.get(name, default)
+
+
+# 🎛️ BASH-CONTROLLED SWITCHES (Defaults if run manually)
+MODEL = os.environ.get("MODEL_TYPE", "d3pm_encoder_decoder")
+NEGATIVES = os.environ.get("INCLUDE_NEG", "False") == "True"
+DIFFUSION_STEPS = _get_env_int("DIFFUSION_STEPS", 128)
+INFERENCE_STEPS = _get_env_int("INFERENCE_NUM_STEPS", min(64, DIFFUSION_STEPS))
+TRAIN_DEVICE = _get_env_str(
+    "TRAIN_DEVICE",
+    "mps" if torch.backends.mps.is_available() else "cpu",
+)
+
+CONFIG = {
+    "model_type": MODEL,
+
+    "data": {
+        "include_negative_examples": NEGATIVES,
+        "dataset_size": 60000,
+    },
+
+    # "model": {
+    #     "vocab_size": 16000,
+    #     "max_seq_len": 80,
+    #     "diffusion_steps": 10,
+    #     "d_model": 384,
+    #     "n_layers": 6,
+    #     "n_heads": 6,
+    #     "d_ff": 1536,
+    #     "dropout": 0.15
+    # },
+    #
+    # "diffusion": {
+    #     "mask_token_id": 0
+    # },
+    #
+    # "training": {
+    #     "batch_size": 32,
+    #     "epochs": 10,
+    #     "lr": 2e-4,
+    #     "label_smoothing": 0.05,
+    #     "precision": "float32",
+    #     "device": "mps" if torch.backends.mps.is_available() else "cpu",
+    #     "early_stopping_patience": 3
+    # }
+# "model": {
+#         "vocab_size": 16000,
+#         "max_seq_len": 96,       # Optimized for GRETIL slokas
+#         "diffusion_steps": 16,   # Use 16 steps (better than 8)
+#         "d_model": 512,          # Wider model learns faster
+#         "n_layers": 8,
+#         "n_heads": 8,
+#         "d_ff": 2048,
+#         "dropout": 0.1
+#     },
+#
+#     "diffusion": {
+#         "mask_token_id": 0
+#     },
+#
+#     "training": {
+#         "batch_size": 32,
+#         "epochs": 20,            # 20 is enough with these tweaks
+#         "lr": 4e-4,              # Higher LR + Warmup for speed
+#         "label_smoothing": 0.15, # Increased for 16k vocab stability
+#         "precision": "float32",
+#         "device": "mps" if torch.backends.mps.is_available() else "cpu",
+#         "early_stopping_patience": 5
+#     }
+'diffusion': {
+        'mask_token_id': 0,          # [MASK] = ID 0, fixed by tokenizer
+    },
+
+    # ── Model architecture ────────────────────────────────────────────
+    'model': {
+        # 'vocab_size':       16000,
+'src_vocab_size': 16000,   # Roman/IAST BPE vocab
+'tgt_vocab_size': 16000,   # Devanagari BPE vocab
+        'd_model':          1024,#512,     # was 384 — kept same, shared embeds save params
+        'n_heads':          8,       # 384 / 6 = 64 head_dim
+        'd_ff':            4096, #2048, #1536,    # 4 × d_model
+        'n_layers':         8,#4,
+        'dropout':          0.2,
+        'max_seq_len':      80,
+        'diffusion_steps':  DIFFUSION_STEPS,
+    },
+
+    # ── Training ──────────────────────────────────────────────────────
+    'training': {
+        'epochs':           20,       # Target: 0.71→0.83-0.85 in 5 epochs
+        'batch_size':       32,
+        'accum_steps':      2,       # effective batch = 64
+        'lr':               7e-5,#6e-4,    # raised from 3e-4; warmup protects first steps
+        'label_smoothing':  0.1,     # was 0.0; reduces overconfidence (gap 1.7 nats)
+        'patience':         4,       # early stop after 4 non-improving epochs
+        'l1_lambda':        1e-7,    # very light L1
+        'device':           TRAIN_DEVICE,
+    },
+
+    # ── Inference (used during val BERTScore and generate()) ──────────
+    'inference': {
+        'num_steps':          INFERENCE_STEPS,
+        'temperature':        0.7,   # slightly lower = more confident output
+        'top_k':              40,
+        'repetition_penalty': 1.2,
+        'diversity_penalty':  0.5,   # keep off; global-mean penalty is conservative
+    },
+}

config_T64.py

@@ -0,0 +1,119 @@
+# Ablation config: T=64 diffusion steps
+import os
+import torch
+
+
+def _get_env_int(name, default):
+    value = os.environ.get(name)
+    return int(value) if value is not None else default
+
+
+def _get_env_str(name, default):
+    return os.environ.get(name, default)
+
+
+# 🎛️ BASH-CONTROLLED SWITCHES (Defaults if run manually)
+MODEL = os.environ.get("MODEL_TYPE", "d3pm_encoder_decoder")
+NEGATIVES = os.environ.get("INCLUDE_NEG", "False") == "True"
+DIFFUSION_STEPS = _get_env_int("DIFFUSION_STEPS", 128)
+INFERENCE_STEPS = _get_env_int("INFERENCE_NUM_STEPS", min(64, DIFFUSION_STEPS))
+TRAIN_DEVICE = _get_env_str(
+    "TRAIN_DEVICE",
+    "mps" if torch.backends.mps.is_available() else "cpu",
+)
+
+CONFIG = {
+    "model_type": MODEL,
+
+    "data": {
+        "include_negative_examples": NEGATIVES,
+        "dataset_size": 60000,
+    },
+
+    # "model": {
+    #     "vocab_size": 16000,
+    #     "max_seq_len": 80,
+    #     "diffusion_steps": 10,
+    #     "d_model": 384,
+    #     "n_layers": 6,
+    #     "n_heads": 6,
+    #     "d_ff": 1536,
+    #     "dropout": 0.15
+    # },
+    #
+    # "diffusion": {
+    #     "mask_token_id": 0
+    # },
+    #
+    # "training": {
+    #     "batch_size": 32,
+    #     "epochs": 10,
+    #     "lr": 2e-4,
+    #     "label_smoothing": 0.05,
+    #     "precision": "float32",
+    #     "device": "mps" if torch.backends.mps.is_available() else "cpu",
+    #     "early_stopping_patience": 3
+    # }
+# "model": {
+#         "vocab_size": 16000,
+#         "max_seq_len": 96,       # Optimized for GRETIL slokas
+#         "diffusion_steps": 16,   # Use 16 steps (better than 8)
+#         "d_model": 512,          # Wider model learns faster
+#         "n_layers": 8,
+#         "n_heads": 8,
+#         "d_ff": 2048,
+#         "dropout": 0.1
+#     },
+#
+#     "diffusion": {
+#         "mask_token_id": 0
+#     },
+#
+#     "training": {
+#         "batch_size": 32,
+#         "epochs": 20,            # 20 is enough with these tweaks
+#         "lr": 4e-4,              # Higher LR + Warmup for speed
+#         "label_smoothing": 0.15, # Increased for 16k vocab stability
+#         "precision": "float32",
+#         "device": "mps" if torch.backends.mps.is_available() else "cpu",
+#         "early_stopping_patience": 5
+#     }
+'diffusion': {
+        'mask_token_id': 0,          # [MASK] = ID 0, fixed by tokenizer
+    },
+
+    # ── Model architecture ────────────────────────────────────────────
+    'model': {
+        # 'vocab_size':       16000,
+'src_vocab_size': 16000,   # Roman/IAST BPE vocab
+'tgt_vocab_size': 16000,   # Devanagari BPE vocab
+        'd_model':          1024,#512,     # was 384 — kept same, shared embeds save params
+        'n_heads':          8,       # 384 / 6 = 64 head_dim
+        'd_ff':            4096, #2048, #1536,    # 4 × d_model
+        'n_layers':         8,#4,
+        'dropout':          0.2,
+        'max_seq_len':      80,
+        'diffusion_steps':  DIFFUSION_STEPS,
+    },
+
+    # ── Training ──────────────────────────────────────────────────────
+    'training': {
+        'epochs':           20,       # Target: 0.71→0.83-0.85 in 5 epochs
+        'batch_size':       32,
+        'accum_steps':      2,       # effective batch = 64
+        'lr':               7e-5,#6e-4,    # raised from 3e-4; warmup protects first steps
+        'label_smoothing':  0.1,     # was 0.0; reduces overconfidence (gap 1.7 nats)
+        'patience':         4,       # early stop after 4 non-improving epochs
+        'l1_lambda':        1e-7,    # very light L1
+        'device':           TRAIN_DEVICE,
+    },
+
+    # ── Inference (used during val BERTScore and generate()) ──────────
+    'inference': {
+        'num_steps':          INFERENCE_STEPS,
+        'temperature':        0.7,   # slightly lower = more confident output
+        'top_k':              40,
+        'repetition_penalty': 1.2,
+        'diversity_penalty':  0.5,   # keep off; global-mean penalty is conservative
+    },
+}

config_T8.py

@@ -0,0 +1,119 @@
+# Ablation config: T=8 diffusion steps
+import os
+import torch
+
+
+def _get_env_int(name, default):
+    value = os.environ.get(name)
+    return int(value) if value is not None else default
+
+
+def _get_env_str(name, default):
+    return os.environ.get(name, default)
+
+
+# 🎛️ BASH-CONTROLLED SWITCHES (Defaults if run manually)
+MODEL = os.environ.get("MODEL_TYPE", "d3pm_encoder_decoder")
+NEGATIVES = os.environ.get("INCLUDE_NEG", "False") == "True"
+DIFFUSION_STEPS = _get_env_int("DIFFUSION_STEPS", 128)
+INFERENCE_STEPS = _get_env_int("INFERENCE_NUM_STEPS", min(64, DIFFUSION_STEPS))
+TRAIN_DEVICE = _get_env_str(
+    "TRAIN_DEVICE",
+    "mps" if torch.backends.mps.is_available() else "cpu",
+)
+
+CONFIG = {
+    "model_type": MODEL,
+
+    "data": {
+        "include_negative_examples": NEGATIVES,
+        "dataset_size": 60000,
+    },
+
+    # "model": {
+    #     "vocab_size": 16000,
+    #     "max_seq_len": 80,
+    #     "diffusion_steps": 10,
+    #     "d_model": 384,
+    #     "n_layers": 6,
+    #     "n_heads": 6,
+    #     "d_ff": 1536,
+    #     "dropout": 0.15
+    # },
+    #
+    # "diffusion": {
+    #     "mask_token_id": 0
+    # },
+    #
+    # "training": {
+    #     "batch_size": 32,
+    #     "epochs": 10,
+    #     "lr": 2e-4,
+    #     "label_smoothing": 0.05,
+    #     "precision": "float32",
+    #     "device": "mps" if torch.backends.mps.is_available() else "cpu",
+    #     "early_stopping_patience": 3
+    # }
+# "model": {
+#         "vocab_size": 16000,
+#         "max_seq_len": 96,       # Optimized for GRETIL slokas
+#         "diffusion_steps": 16,   # Use 16 steps (better than 8)
+#         "d_model": 512,          # Wider model learns faster
+#         "n_layers": 8,
+#         "n_heads": 8,
+#         "d_ff": 2048,
+#         "dropout": 0.1
+#     },
+#
+#     "diffusion": {
+#         "mask_token_id": 0
+#     },
+#
+#     "training": {
+#         "batch_size": 32,
+#         "epochs": 20,            # 20 is enough with these tweaks
+#         "lr": 4e-4,              # Higher LR + Warmup for speed
+#         "label_smoothing": 0.15, # Increased for 16k vocab stability
+#         "precision": "float32",
+#         "device": "mps" if torch.backends.mps.is_available() else "cpu",
+#         "early_stopping_patience": 5
+#     }
+'diffusion': {
+        'mask_token_id': 0,          # [MASK] = ID 0, fixed by tokenizer
+    },
+
+    # ── Model architecture ────────────────────────────────────────────
+    'model': {
+        # 'vocab_size':       16000,
+'src_vocab_size': 16000,   # Roman/IAST BPE vocab
+'tgt_vocab_size': 16000,   # Devanagari BPE vocab
+        'd_model':          1024,#512,     # was 384 — kept same, shared embeds save params
+        'n_heads':          8,       # 384 / 6 = 64 head_dim
+        'd_ff':            4096, #2048, #1536,    # 4 × d_model
+        'n_layers':         8,#4,
+        'dropout':          0.2,
+        'max_seq_len':      80,
+        'diffusion_steps':  DIFFUSION_STEPS,
+    },
+
+    # ── Training ──────────────────────────────────────────────────────
+    'training': {
+        'epochs':           20,       # Target: 0.71→0.83-0.85 in 5 epochs
+        'batch_size':       32,
+        'accum_steps':      2,       # effective batch = 64
+        'lr':               7e-5,#6e-4,    # raised from 3e-4; warmup protects first steps
+        'label_smoothing':  0.1,     # was 0.0; reduces overconfidence (gap 1.7 nats)
+        'patience':         4,       # early stop after 4 non-improving epochs
+        'l1_lambda':        1e-7,    # very light L1
+        'device':           TRAIN_DEVICE,
+    },
+
+    # ── Inference (used during val BERTScore and generate()) ──────────
+    'inference': {
+        'num_steps':          INFERENCE_STEPS,
+        'temperature':        0.7,   # slightly lower = more confident output
+        'top_k':              40,
+        'repetition_penalty': 1.2,
+        'diversity_penalty':  0.5,   # keep off; global-mean penalty is conservative
+    },
+}

d3pm_model_cross_attention.py

@@ -0,0 +1,271 @@
+"""
+d3pm_model_cross_attention.py  — Cross-Script + Generation-Fixed
+=================================================================
+INPUT  : quote_text       tokens  (Roman script, src_vocab_size)
+OUTPUT : quote_devanagari tokens  (Devanagari script, tgt_vocab_size)
+
+src_embed  uses src_vocab_size  (Roman BPE)
+tgt_embed  uses tgt_vocab_size  (Devanagari BPE)
+head       outputs tgt_vocab_size  (predict Devanagari tokens)
+Weight tying: head <-> tgt_embed only (NOT src_embed)
+
+Generation bugs fixed:
+  BUG 1 - tgt_pad_mask suppressed during inference
+  BUG 2 - q_sample skipped at t=0
+  BUG 3 - time embedding before hint_gate
+  BUG 4 - diversity penalty uses global mean not var
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffusion.scheduler import OptimizedCosineScheduler
+from diffusion.forward_process import AbsorbingForwardProcess
+
+
+class SinusoidalPositionalEncoding(nn.Module):
+    def __init__(self, d_model, max_len=5000):
+        super().__init__()
+        pe       = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len).unsqueeze(1).float()
+        div_term = torch.exp(
+            torch.arange(0, d_model, 2).float() *
+            (-torch.log(torch.tensor(10000.0)) / d_model)
+        )
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        self.register_buffer("pe", pe.unsqueeze(0))
+
+    def forward(self, x):
+        return x + self.pe[:, :x.size(1), :]
+
+
+class SanskritEmbeddings(nn.Module):
+    def __init__(self, vocab_size, d_model, max_seq_len):
+        super().__init__()
+        self.token_emb       = nn.Embedding(vocab_size, d_model)
+        self.pos_enc         = SinusoidalPositionalEncoding(d_model, max_seq_len)
+        self.token_embedding = self.token_emb
+    def forward(self, tokens):
+        return self.pos_enc(self.token_emb(tokens))
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, d_model, n_heads, dropout=0.1):
+        super().__init__()
+        assert d_model % n_heads == 0
+        self.d_model  = d_model
+        self.n_heads  = n_heads
+        self.head_dim = d_model // n_heads
+        self.q_proj   = nn.Linear(d_model, d_model)
+        self.k_proj   = nn.Linear(d_model, d_model)
+        self.v_proj   = nn.Linear(d_model, d_model)
+        self.out_proj = nn.Linear(d_model, d_model)
+        self.dropout  = nn.Dropout(dropout)
+
+    def forward(self, q, k, v, mask=None):
+        B, Lq, _ = q.size()
+        Lk = k.size(1)
+        Q = self.q_proj(q).view(B, Lq, self.n_heads, self.head_dim).transpose(1, 2)
+        K = self.k_proj(k).view(B, Lk, self.n_heads, self.head_dim).transpose(1, 2)
+        V = self.v_proj(v).view(B, Lk, self.n_heads, self.head_dim).transpose(1, 2)
+        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.head_dim ** 0.5)
+        if mask is not None:
+            scores = scores.masked_fill(mask.unsqueeze(1).unsqueeze(2), float('-inf'))
+        attn = self.dropout(torch.softmax(scores, dim=-1))
+        out  = torch.matmul(attn, V).transpose(1, 2).contiguous().view(B, Lq, self.d_model)
+        return self.out_proj(out)
+
+
+class EncoderBlock(nn.Module):
+    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
+        super().__init__()
+        self.mha   = MultiHeadAttention(d_model, n_heads, dropout)
+        self.ff    = nn.Sequential(nn.Linear(d_model, d_ff), nn.GELU(), nn.Dropout(dropout),
+                                   nn.Linear(d_ff, d_model), nn.Dropout(dropout))
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+    def forward(self, x, pad_mask=None):
+        x = self.norm1(x + self.mha(x, x, x, mask=pad_mask))
+        return self.norm2(x + self.ff(x))
+
+
+class DecoderBlock(nn.Module):
+    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
+        super().__init__()
+        self.self_attn  = MultiHeadAttention(d_model, n_heads, dropout)
+        self.cross_attn = MultiHeadAttention(d_model, n_heads, dropout)
+        self.ff         = nn.Sequential(nn.Linear(d_model, d_ff), nn.GELU(), nn.Dropout(dropout),
+                                        nn.Linear(d_ff, d_model), nn.Dropout(dropout))
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.norm3 = nn.LayerNorm(d_model)
+    def forward(self, x, memory, tgt_pad_mask=None, src_pad_mask=None):
+        x = self.norm1(x + self.self_attn(x, x, x, mask=tgt_pad_mask))
+        x = self.norm2(x + self.cross_attn(x, memory, memory, mask=src_pad_mask))
+        return self.norm3(x + self.ff(x))
+
+
+class D3PMCrossAttention(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg           = cfg
+        self.mask_token_id = cfg['diffusion']['mask_token_id']
+        d      = cfg['model']['d_model']
+        nhead  = cfg['model']['n_heads']
+        d_ff   = cfg['model']['d_ff']
+        drop   = cfg['model']['dropout']
+        seqlen = cfg['model']['max_seq_len']
+        nlayer = cfg['model']['n_layers']
+        src_vocab = cfg['model'].get('src_vocab_size', cfg['model']['vocab_size'])
+        tgt_vocab = cfg['model'].get('tgt_vocab_size', cfg['model']['vocab_size'])
+
+        # Separate embeddings: Roman src, Devanagari tgt
+        self.src_embed = SanskritEmbeddings(src_vocab, d, seqlen)
+        self.tgt_embed = SanskritEmbeddings(tgt_vocab, d, seqlen)
+
+        self.scheduler       = OptimizedCosineScheduler(cfg)
+        self.forward_process = AbsorbingForwardProcess(self.scheduler)
+
+        self.encoder_blocks = nn.ModuleList([EncoderBlock(d, nhead, d_ff, drop) for _ in range(nlayer)])
+        self.decoder_blocks = nn.ModuleList([DecoderBlock(d, nhead, d_ff, drop) for _ in range(nlayer)])
+
+        self.time_mlp  = nn.Sequential(nn.Linear(1, d//4), nn.SiLU(), nn.Linear(d//4, d))
+        self.hint_gate = nn.Sequential(nn.Linear(d, d), nn.Sigmoid())
+
+        # Output head: predict Devanagari tokens, tied to tgt_embed
+        self.head = nn.Linear(d, tgt_vocab, bias=False)
+        self.head.weight = self.tgt_embed.token_embedding.weight
+
+    def forward(self, src, tgt, t, x0_hint=None, inference_mode=False):
+        PAD = 1
+        src_pad_mask = (src == PAD)
+        # BUG 1 FIX: no tgt mask during inference
+        tgt_pad_mask = None if inference_mode else (tgt == PAD)
+
+        # Encode Roman source
+        memory = self.src_embed(src)
+        for block in self.encoder_blocks:
+            memory = block(memory, pad_mask=src_pad_mask)
+
+        # BUG 2 FIX: skip q_sample at final step t=0
+        if inference_mode and (t == 0).all():
+            x_t_ids = tgt
+        else:
+            _, x_t_ids = self.forward_process.q_sample(tgt, t)
+
+        x = self.tgt_embed(x_t_ids)
+
+        # BUG 3 FIX: time embedding BEFORE hint gate
+        t_norm = t.float() / self.scheduler.num_timesteps
+        t_emb  = self.time_mlp(t_norm.unsqueeze(-1))
+        x      = x + t_emb.unsqueeze(1)
+
+        if x0_hint is not None:
+            hint_emb = self.tgt_embed(x0_hint)
+            gate     = self.hint_gate(x)   # time-aware gate
+            x        = x + gate * hint_emb
+
+        for block in self.decoder_blocks:
+            x = block(x, memory, tgt_pad_mask=tgt_pad_mask, src_pad_mask=src_pad_mask)
+
+        return self.head(x), None
+
+    @torch.no_grad()
+    def generate(self, src, num_steps=None, temperature=0.8, top_k=50,
+                 repetition_penalty=1.2, diversity_penalty=0.0):
+        if src.dim() == 1:
+            src = src.unsqueeze(0)
+        device = src.device
+        B, L   = src.shape
+        T      = self.scheduler.num_timesteps
+        steps  = num_steps or T
+        step_size = max(1, T // steps)
+        timesteps = list(range(T - 1, -1, -step_size))
+        if timesteps[-1] != 0:
+            timesteps.append(0)
+
+        mask_id = self.mask_token_id
+        x0_est  = torch.full((B, L), mask_id, dtype=torch.long, device=device)
+        hint    = None
+
+        self.eval()
+        with torch.no_grad():
+            for step_idx, t_val in enumerate(timesteps):
+                t       = torch.full((B,), t_val, dtype=torch.long, device=device)
+                is_last = (step_idx == len(timesteps) - 1)
+                logits, _ = self.forward(src, x0_est, t, x0_hint=hint, inference_mode=True)
+                if repetition_penalty != 1.0:
+                    logits = _apply_repetition_penalty(logits, x0_est, repetition_penalty)
+                if diversity_penalty > 0.0:
+                    logits = _apply_diversity_penalty_fixed(logits, diversity_penalty)  # BUG 4 FIX
+                logits = logits / max(temperature, 1e-5)
+                if top_k > 0:
+                    logits = _top_k_filter(logits, top_k)
+                probs = F.softmax(logits, dim=-1)
+                x0_est = torch.argmax(probs, dim=-1) if is_last else _batch_multinomial(probs)
+                hint = x0_est
+        return x0_est
+
+
+class BaselineCrossAttention(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        d = cfg['model']['d_model']; nhead = cfg['model']['n_heads']
+        d_ff = cfg['model']['d_ff']; drop = cfg['model']['dropout']
+        seqlen = cfg['model']['max_seq_len']; nlayer = cfg['model']['n_layers']
+        src_vocab = cfg['model'].get('src_vocab_size', cfg['model']['vocab_size'])
+        tgt_vocab = cfg['model'].get('tgt_vocab_size', cfg['model']['vocab_size'])
+        self.src_embed = SanskritEmbeddings(src_vocab, d, seqlen)
+        self.tgt_embed = SanskritEmbeddings(tgt_vocab, d, seqlen)
+        self.encoder_blocks = nn.ModuleList([EncoderBlock(d, nhead, d_ff, drop) for _ in range(nlayer)])
+        self.decoder_blocks = nn.ModuleList([DecoderBlock(d, nhead, d_ff, drop) for _ in range(nlayer)])
+        self.head = nn.Linear(d, tgt_vocab, bias=False)
+        self.head.weight = self.tgt_embed.token_embedding.weight
+
+    def forward(self, src, tgt, t=None, x0_hint=None):
+        PAD = 1
+        memory = self.src_embed(src)
+        for b in self.encoder_blocks: memory = b(memory, pad_mask=(src==PAD))
+        x = self.tgt_embed(tgt)
+        for b in self.decoder_blocks: x = b(x, memory, tgt_pad_mask=(tgt==PAD), src_pad_mask=(src==PAD))
+        return (self.head(x),)
+
+    @torch.no_grad()
+    def generate(self, src, max_len=None, start_token_id=2, **kwargs):
+        if max_len is None: max_len = src.size(1)
+        B, device = src.size(0), src.device
+        memory = self.src_embed(src)
+        for b in self.encoder_blocks: memory = b(memory, pad_mask=(src==1))
+        ys = torch.full((B, 1), start_token_id, dtype=torch.long, device=device)
+        for _ in range(max_len):
+            x = self.tgt_embed(ys)
+            for b in self.decoder_blocks: x = b(x, memory, tgt_pad_mask=None, src_pad_mask=(src==1))
+            ys = torch.cat([ys, torch.argmax(self.head(x)[:,-1,:], dim=-1, keepdim=True)], dim=1)
+        return ys[:, 1:max_len+1]
+
+
+# helpers
+def _top_k_filter(logits, k):
+    B, L, V = logits.shape
+    if k >= V: return logits
+    topk_vals, _ = torch.topk(logits, k, dim=-1)
+    return logits.masked_fill(logits < topk_vals[..., -1].unsqueeze(-1), float('-inf'))
+
+def _batch_multinomial(probs):
+    B, L, V = probs.shape
+    flat = probs.view(B*L, V) + 1e-9
+    return torch.multinomial(flat/flat.sum(-1,keepdim=True), 1).squeeze(-1).view(B, L)
+
+def _apply_repetition_penalty(logits, prev_tokens, penalty):
+    for b in range(logits.shape[0]):
+        for tid in set(prev_tokens[b].tolist()):
+            if tid > 4: logits[b, :, tid] = logits[b, :, tid] / penalty
+    return logits
+
+def _apply_diversity_penalty(logits, penalty):          # legacy wrong version
+    return logits + penalty * logits.var(dim=-1, keepdim=True)
+
+def _apply_diversity_penalty_fixed(logits, penalty):    # correct version
+    return logits - penalty * logits.mean(dim=1, keepdim=True)

d3pm_model_encoder_decoder.py

@@ -0,0 +1,227 @@
+import torch
+import torch.nn as nn
+from diffusion.scheduler import OptimizedCosineScheduler
+from diffusion.forward_process import AbsorbingForwardProcess
+# Import shared classes to guarantee identical architectures
+from model.d3pm_model_cross_attention import SanskritEmbeddings, EncoderBlock, MultiHeadAttention
+class DecoderBlock(nn.Module):
+    def __init__(self, d_model, n_heads, d_ff, dropout=0.15):
+        super().__init__()
+        self.self_attn  = MultiHeadAttention(d_model, n_heads, dropout)
+        self.cross_attn = MultiHeadAttention(d_model, n_heads, dropout)  # ← restored
+        self.ff = nn.Sequential(
+            nn.Linear(d_model, d_ff),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(d_ff, d_model),
+            nn.Dropout(dropout),
+        )
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.norm3 = nn.LayerNorm(d_model)  # ← restored (for cross-attn residual)
+
+    def forward(self, x, memory, tgt_pad_mask=None):
+        # 1. Masked self-attention on target
+        x = self.norm1(x + self.self_attn(x, x, x, mask=tgt_pad_mask))
+        # 2. Cross-attention: queries from decoder, keys/values from encoder memory
+        x = self.norm2(x + self.cross_attn(x, memory, memory))
+        # 3. Feed-forward
+        return self.norm3(x + self.ff(x))
+
+
+class DecoderBlockNoCrossAttn(nn.Module):
+    """Kept for reference — NOT used by D3PMEncoderDecoder."""
+    def __init__(self, d_model, n_heads, d_ff, dropout=0.15):
+        super().__init__()
+        self.self_attn = MultiHeadAttention(d_model, n_heads, dropout)
+        self.ff = nn.Sequential(
+            nn.Linear(d_model, d_ff), nn.ReLU(), nn.Dropout(dropout),
+            nn.Linear(d_ff, d_model), nn.Dropout(dropout),
+        )
+        self.norm1, self.norm2 = nn.LayerNorm(d_model), nn.LayerNorm(d_model)
+
+    def forward(self, x, tgt_pad_mask=None, causal_mask=None):
+        combined_mask = None
+        if tgt_pad_mask is not None and causal_mask is not None:
+            combined_mask = tgt_pad_mask | causal_mask
+        elif causal_mask is not None:
+            combined_mask = causal_mask
+        elif tgt_pad_mask is not None:
+            combined_mask = tgt_pad_mask
+        x = self.norm1(x + self.self_attn(x, x, x, mask=combined_mask))
+        return self.norm2(x + self.ff(x))
+
+
+# ============================================================
+# 1. D3PM Encoder-Decoder Model
+# ============================================================
+class D3PMEncoderDecoder(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg           = cfg
+        self.mask_token_id = cfg['diffusion']['mask_token_id']
+
+        src_vocab = cfg['model'].get('src_vocab_size', cfg['model']['vocab_size'])
+        tgt_vocab = cfg['model'].get('tgt_vocab_size', cfg['model']['vocab_size'])
+        d_model   = cfg['model']['d_model']
+        n_heads   = cfg['model']['n_heads']
+        d_ff      = cfg['model']['d_ff']
+        dropout   = cfg['model']['dropout']
+        n_layers  = cfg['model']['n_layers']
+        max_len   = cfg['model']['max_seq_len']
+
+        self.src_embed = SanskritEmbeddings(src_vocab, d_model, max_len)
+        self.tgt_embed = SanskritEmbeddings(tgt_vocab, d_model, max_len)
+
+        self.scheduler       = OptimizedCosineScheduler(cfg)
+        self.forward_process = AbsorbingForwardProcess(self.scheduler)
+
+        self.encoder_blocks = nn.ModuleList([
+            EncoderBlock(d_model, n_heads, d_ff, dropout) for _ in range(n_layers)
+        ])
+        # DecoderBlock now has cross-attention — matches saved checkpoint
+        self.decoder_blocks = nn.ModuleList([
+            DecoderBlock(d_model, n_heads, d_ff, dropout) for _ in range(n_layers)
+        ])
+
+        self.time_mlp = nn.Sequential(
+            nn.Linear(1, d_model // 4), nn.SiLU(),
+            nn.Linear(d_model // 4, d_model),
+        )
+        self.head        = nn.Linear(d_model, tgt_vocab)
+        self.head.weight = self.tgt_embed.token_embedding.weight
+
+    def forward(self, src, tgt, t, x0_hint=None):
+        src_pad_mask = (src == 1)
+        tgt_pad_mask = (tgt == 1)
+
+        # Encode source (Roman IAST)
+        memory = self.src_embed(src)
+        for block in self.encoder_blocks:
+            memory = block(memory, pad_mask=src_pad_mask)
+
+        # Corrupt target with forward diffusion
+        _, x_t_ids = self.forward_process.q_sample(tgt, t)
+
+        # Optionally blend in x0_hint (self-conditioning)
+        if x0_hint is not None:
+            hint_prob  = 0.5
+            blend_mask = (torch.rand(x_t_ids.shape, device=x_t_ids.device) < hint_prob)
+            still_mask = (x_t_ids == self.mask_token_id)
+            x_t_ids    = torch.where(blend_mask & still_mask, x0_hint, x_t_ids)
+
+        x     = self.tgt_embed(x_t_ids)
+        t_emb = self.time_mlp(t.float().unsqueeze(-1)).unsqueeze(1)
+        x     = x + t_emb.expand(-1, tgt.shape[1], -1)
+
+        # Decode with cross-attention over encoder memory
+        for block in self.decoder_blocks:
+            x = block(x, memory, tgt_pad_mask=tgt_pad_mask)
+
+        return self.head(x), None
+
+    @torch.no_grad()
+    def generate(
+        self,
+        src,
+        num_steps          = None,
+        temperature        = 0.75,
+        top_k              = 50,
+        repetition_penalty = 1.15,
+        diversity_penalty  = 0.0,
+    ):
+        """
+        Iterative D3PM reverse diffusion — same signature as
+        D3PMCrossAttention.generate() so SanskritModel.generate() works
+        identically for both model types.
+        """
+        device   = src.device
+        B, L     = src.shape[0], self.cfg['model']['max_seq_len']
+        T        = num_steps or self.scheduler.num_timesteps
+        mask_id  = self.mask_token_id
+        pad_id   = 1
+
+        x0_est = torch.full((B, L), mask_id, dtype=torch.long, device=device)
+
+        for step in range(T - 1, -1, -1):
+            t_tensor = torch.full((B,), step, dtype=torch.long, device=device)
+            hint     = x0_est.clone()
+
+            logits, _ = self.forward(src, x0_est, t_tensor, x0_hint=hint)
+
+            # Repetition penalty
+            if repetition_penalty != 1.0:
+                for b in range(B):
+                    for tok in set(x0_est[b].tolist()):
+                        if tok > pad_id:
+                            logits[b, :, tok] /= repetition_penalty
+
+            # Diversity penalty (suppress common tokens)
+            if diversity_penalty > 0.0:
+                logits = logits - diversity_penalty * logits.mean(dim=1, keepdim=True)
+
+            # Temperature + top-k sampling
+            logits = logits / max(temperature, 1e-8)
+            if top_k > 0:
+                vals, _ = torch.topk(logits, top_k, dim=-1)
+                logits  = logits.masked_fill(logits < vals[..., -1:], float('-inf'))
+
+            probs  = torch.softmax(logits, dim=-1)
+            # Only update positions that are still masked
+            still  = (x0_est == mask_id)
+            sample = torch.multinomial(probs.view(-1, probs.size(-1)), 1).view(B, L)
+            x0_est = torch.where(still, sample, x0_est)
+
+        return x0_est
+
+
+# ============================================================
+# 2. Baseline Encoder-Decoder Model (unchanged)
+# ============================================================
+class BaselineEncoderDecoder(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg       = cfg
+        self.src_embed = SanskritEmbeddings(cfg['model']['vocab_size'], cfg['model']['d_model'],
+                                            cfg['model']['max_seq_len'])
+        self.tgt_embed = SanskritEmbeddings(cfg['model']['vocab_size'], cfg['model']['d_model'],
+                                            cfg['model']['max_seq_len'])
+        self.encoder_blocks = nn.ModuleList([
+            EncoderBlock(cfg['model']['d_model'], cfg['model']['n_heads'],
+                         cfg['model']['d_ff'], cfg['model']['dropout'])
+            for _ in range(cfg['model']['n_layers'])
+        ])
+        self.decoder_blocks = nn.ModuleList([
+            DecoderBlock(cfg['model']['d_model'], cfg['model']['n_heads'],
+                         cfg['model']['d_ff'], cfg['model']['dropout'])
+            for _ in range(cfg['model']['n_layers'])
+        ])
+        self.head        = nn.Linear(cfg['model']['d_model'], cfg['model']['vocab_size'])
+        self.head.weight = self.tgt_embed.token_embedding.weight
+
+    def forward(self, src, tgt):
+        src_pad_mask, tgt_pad_mask = (src == 1), (tgt == 1)
+        memory = self.src_embed(src)
+        for block in self.encoder_blocks:
+            memory = block(memory, pad_mask=src_pad_mask)
+        x = self.tgt_embed(tgt)
+        for block in self.decoder_blocks:
+            x = block(x, memory, tgt_pad_mask=tgt_pad_mask)
+        return self.head(x)
+
+    @torch.no_grad()
+    def generate(self, src, max_len=80, start_token_id=2):
+        batch_size, device = src.size(0), src.device
+        src_pad_mask = (src == 1)
+        memory = self.src_embed(src)
+        for block in self.encoder_blocks:
+            memory = block(memory, pad_mask=src_pad_mask)
+        ys = torch.ones(batch_size, 1, dtype=torch.long, device=device) * start_token_id
+        for _ in range(max_len):
+            x = self.tgt_embed(ys)
+            for block in self.decoder_blocks:
+                x = block(x, memory, tgt_pad_mask=None)
+            logits     = self.head(x)
+            next_token = torch.argmax(logits[:, -1, :], dim=-1, keepdim=True)
+            ys         = torch.cat([ys, next_token], dim=1)
+        return ys[:, 1:]

dataset.py

@@ -0,0 +1,152 @@
+"""
+dataset.py  — Cross-Script Translation Fix
+==========================================
+INPUT  : quote_text       (Roman/IAST transliteration of Sanskrit)
+TARGET : quote_devanagari (Devanagari script)
+
+This is the CORRECT task: the model learns to transliterate / translate
+Roman Sanskrit → Devanagari, which is a meaningful, learnable mapping
+(far better than devanagari→devanagari reconstruction which teaches nothing).
+
+KEY CHANGES from original:
+  1. _input_field  = 'quote_text'        (was 'quote_devanagari')
+  2. _target_field = 'quote_devanagari'  (unchanged)
+  3. Separate source/target tokenizers — Roman and Devanagari have
+     completely different character sets; a shared BPE vocab forces the
+     model to learn both scripts in one embedding table, which wastes
+     capacity and confuses the attention mechanism.
+  4. Negative example generation fixed — reversal now operates on
+     DEVANAGARI target only (not accidentally on Roman source).
+  5. curriculum_sort uses target length (Devanagari) for difficulty proxy.
+"""
+
+from datasets import load_dataset
+from torch.utils.data import Dataset
+import torch
+import torch.nn.functional as F
+import random
+
+
+class OptimizedSanskritDataset(Dataset):
+    def __init__(self, split='train', tokenizer=None, max_len=80, cfg=None,
+                 src_tokenizer=None, tgt_tokenizer=None):
+        """
+        Args:
+            tokenizer     : shared tokenizer (legacy — used if src/tgt not provided)
+            src_tokenizer : tokenizer for quote_text  (Roman script)
+            tgt_tokenizer : tokenizer for quote_devanagari (Devanagari script)
+                            If None, falls back to shared `tokenizer`.
+        """
+        from config import CONFIG
+        self.cfg = cfg or CONFIG
+        self.max_len = max_len
+        self.pad_id  = 1
+        self.mask_id = self.cfg['diffusion']['mask_token_id']
+        self.include_negatives = self.cfg['data']['include_negative_examples']
+
+        # ── Tokenizer setup ───────────────────────────────────────────
+        # Support both legacy (shared) and new (separate src/tgt) tokenizers
+        self.src_tokenizer = src_tokenizer or tokenizer
+        self.tgt_tokenizer = tgt_tokenizer or tokenizer
+
+        if self.src_tokenizer is None:
+            raise ValueError("Provide at least one tokenizer.")
+
+        print(f"📥 Loading '{split}' split …")
+        raw = load_dataset("paws/sanskrit-verses-gretil", split=split)
+        cols = raw.column_names
+
+        # ── Field selection ───────────────────────────────────────────
+        if 'quote_text' in cols and 'quote_devanagari' in cols:
+            # CORRECT setup: Roman input → Devanagari output
+            self._input_field  = 'quote_text'
+            self._target_field = 'quote_devanagari'
+            print("   Format: quote_text (Roman) → quote_devanagari (Devanagari) ✓")
+        elif 'sentence1' in cols and 'sentence2' in cols:
+            # PAWS paraphrase pairs fallback
+            self._input_field  = 'sentence1'
+            self._target_field = 'sentence2'
+            print("   Format: PAWS sentence pairs ✓")
+        else:
+            # Last resort: same field both sides
+            self._input_field  = 'quote_devanagari'
+            self._target_field = 'quote_devanagari'
+            print("   ⚠️  Format: Devanagari→Devanagari (suboptimal — no quote_text found)")
+
+        # ── Filter empty rows ─────────────────────────────────────────
+        # Some rows have empty quote_text — skip them
+        raw = raw.filter(
+            lambda ex: (
+                bool(ex[self._input_field].strip()) and
+                bool(ex[self._target_field].strip())
+            )
+        )
+        print(f"   After empty-filter: {len(raw)} samples")
+
+        self.dataset = raw
+
+        if split == 'train':
+            self.dataset = self._curriculum_sort()
+
+        print(f"✅ {len(self.dataset)} samples loaded.")
+
+    # ── Encoding ──────────────────────────────────────────────────────
+
+    def _encode_src(self, text):
+        """Encode source (Roman) text."""
+        ids = self.src_tokenizer.encode(text)[:self.max_len]
+        t   = torch.tensor(ids, dtype=torch.long)
+        t   = F.pad(t, (0, max(0, self.max_len - len(t))), value=self.pad_id)
+        return t
+
+    def _encode_tgt(self, text):
+        """Encode target (Devanagari) text."""
+        ids = self.tgt_tokenizer.encode(text)[:self.max_len]
+        t   = torch.tensor(ids, dtype=torch.long)
+        t   = F.pad(t, (0, max(0, self.max_len - len(t))), value=self.pad_id)
+        return t
+
+    # ── Curriculum ───��────────────────────────────────────────────────
+
+    def _curriculum_sort(self):
+        """Short, common Devanagari targets first → long, rare targets last."""
+        scores = []
+        for s in self.dataset:
+            text         = s[self._target_field]
+            length       = len(text.split())
+            rarity_score = len(set(text)) / max(1, len(text))
+            scores.append(length * (1 - rarity_score))
+        order = sorted(range(len(self.dataset)), key=lambda i: scores[i])
+        return self.dataset.select(order)
+
+    # ── Item ──────────────────────────────────────────────────────────
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        sample = self.dataset[idx]
+
+        src_text = sample[self._input_field].strip()
+        tgt_text = sample[self._target_field].strip()
+
+        input_ids  = self._encode_src(src_text)   # Roman encoded with src_tokenizer
+        target_ids = self._encode_tgt(tgt_text)   # Devanagari encoded with tgt_tokenizer
+
+        out = {
+            'input_ids':   input_ids,
+            'target_ids':  target_ids,
+            'input_text':  src_text,
+            'target_text': tgt_text,
+        }
+
+        if self.include_negatives:
+            neg_ids = target_ids.clone()
+            # Reverse a random chunk of the DEVANAGARI target
+            non_pad = (neg_ids != self.pad_id).sum().item()
+            if non_pad > 4:
+                i1, i2 = sorted(random.sample(range(non_pad), 2))
+                neg_ids[i1:i2] = torch.flip(neg_ids[i1:i2], dims=[0])
+            out['negative_target_ids'] = neg_ids
+
+        return out

forward_process.py

@@ -0,0 +1,21 @@
+"""
+forward_process.py  — Verified Correct (no changes needed)
+===========================================================
+Absorbing (mask) diffusion. PAD never masked. At t=0 alpha=1.0 exactly
+so x_t == x_0 (nothing masked). Works correctly with the fixed scheduler.
+"""
+import torch
+
+class AbsorbingForwardProcess:
+    def __init__(self, scheduler, mask_id=0, pad_id=1):
+        self.scheduler = scheduler
+        self.mask_id   = mask_id
+        self.pad_id    = pad_id
+
+    def q_sample(self, x_0, t):
+        alpha_t = self.scheduler.get_alpha(t).to(x_0.device).view(-1, 1)
+        r   = torch.rand(x_0.shape, device=x_0.device)
+        x_t = x_0.clone()
+        x_t[r > alpha_t]          = self.mask_id
+        x_t[x_0 == self.pad_id]   = self.pad_id   # PAD stays PAD always
+        return x_0, x_t

inference.py

@@ -0,0 +1,300 @@
+"""
+inference.py
+============
+Correct D3PM inference for Sanskrit paraphrase generation.
+
+The model's forward() takes CLEAN tgt and noises it internally.
+So inference passes x0_estimate (starting all-[MASK]) as tgt each step,
+letting the model noise it and then predict a cleaner version.
+
+Also includes: robust checkpoint loading (auto-detects architecture
+from saved weights — no CONFIG mismatch crashes).
+"""
+
+import torch
+import torch.nn.functional as F
+import os, sys
+from tqdm import tqdm
+from torch.utils.data import DataLoader, Subset
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)))
+from config import CONFIG
+
+
+# ── Checkpoint loader ─────────────────────────────────────────────────
+
+def load_model(ckpt_path: str, base_cfg: dict, device: torch.device):
+    """
+    Auto-detect architecture from checkpoint weight shapes,
+    then load. Never fails due to CONFIG vs checkpoint mismatch.
+    """
+    import copy
+    from model.sanskrit_model import SanskritModel
+
+    cfg   = copy.deepcopy(base_cfg)
+    state = torch.load(ckpt_path, map_location='cpu')
+
+    # d_model + vocab_size
+    ek = 'model.src_embed.token_emb.weight'
+    if ek in state:
+        vocab, d          = state[ek].shape
+        cfg['model']['vocab_size'] = vocab
+        cfg['model']['d_model']    = d
+        cfg['model']['d_ff']       = d * 4
+
+    # n_layers
+    ids = {int(k.split('.')[2]) for k in state if k.startswith('model.encoder_blocks.')}
+    if ids:
+        cfg['model']['n_layers'] = max(ids) + 1
+
+    # max_seq_len
+    pk = 'model.src_embed.pos_enc.pe'
+    if pk in state:
+        cfg['model']['max_seq_len'] = state[pk].shape[1]
+
+    # n_heads
+    d = cfg['model']['d_model']
+    h = cfg['model'].get('n_heads', 6)
+    if d % h != 0:
+        h = next(x for x in [8, 6, 4, 2, 1] if d % x == 0)
+    cfg['model']['n_heads'] = h
+
+    print(f"🔍 Detected: d_model={cfg['model']['d_model']}, "
+          f"n_layers={cfg['model']['n_layers']}, "
+          f"max_seq_len={cfg['model']['max_seq_len']}, "
+          f"n_heads={cfg['model']['n_heads']}")
+
+    model = SanskritModel(cfg).to(device)
+    missing, unexpected = model.load_state_dict(
+        torch.load(ckpt_path, map_location=device), strict=False
+    )
+
+    # hint_gate may be absent in older checkpoints — initialise safely
+    allowed = {'model.hint_gate.0.weight', 'model.hint_gate.0.bias'}
+    real_missing = [k for k in missing if k not in allowed]
+    if real_missing:
+        print(f"⚠️  Missing keys: {real_missing[:3]} …")
+    if unexpected:
+        print(f"⚠️  Unexpected keys: {unexpected[:3]} …")
+    if hasattr(model.model, 'hint_gate') and 'model.hint_gate.0.weight' in missing:
+        with torch.no_grad():
+            w = model.model.hint_gate[0].weight
+            torch.nn.init.zeros_(model.model.hint_gate[0].bias)
+            torch.nn.init.eye_(w) if w.shape[0] == w.shape[1] \
+                else torch.nn.init.xavier_uniform_(w)
+        print("ℹ️  hint_gate initialised to identity (not in checkpoint).")
+
+    print("✅ Model loaded.")
+    return model, cfg
+
+
+# ── Core inference function ───────────────────────────────────────────
+
+def run_inference(model, input_ids, cfg):
+    """
+    Correct D3PM iterative refinement.
+
+    x0_est starts as all [MASK].
+    Each step: model(src, x0_est, t) noises x0_est internally,
+    then predicts a cleaner version.  x0_est is updated each step.
+    """
+    inf    = cfg['inference']
+    device = input_ids.device
+    B, L   = input_ids.shape
+
+    inner   = model.model
+    T       = inner.scheduler.num_timesteps
+    steps   = inf['num_steps']           # must equal T (set in config)
+    step_size = max(1, T // steps)
+    timesteps = list(range(T - 1, -1, -step_size))
+    if timesteps[-1] != 0:
+        timesteps.append(0)
+
+    mask_id = inner.mask_token_id
+    x0_est  = torch.full((B, L), mask_id, dtype=torch.long, device=device)
+    hint    = None
+
+    model.eval()
+    with torch.no_grad():
+        for step_idx, t_val in enumerate(timesteps):
+            t            = torch.full((B,), t_val, dtype=torch.long, device=device)
+            is_last      = (step_idx == len(timesteps) - 1)
+
+            logits, _    = model(input_ids, x0_est, t, x0_hint=hint)
+
+            # Penalties
+            if inf['repetition_penalty'] != 1.0:
+                from model.d3pm_model_cross_attention import _apply_repetition_penalty
+                logits = _apply_repetition_penalty(
+                    logits, x0_est, inf['repetition_penalty']
+                )
+            if inf['diversity_penalty'] > 0.0:
+                from model.d3pm_model_cross_attention import _apply_diversity_penalty
+                logits = _apply_diversity_penalty(logits, inf['diversity_penalty'])
+
+            logits = logits / max(inf['temperature'], 1e-5)
+            if inf['top_k'] > 0:
+                from model.d3pm_model_cross_attention import _top_k_filter
+                logits = _top_k_filter(logits, inf['top_k'])
+
+            probs  = F.softmax(logits, dim=-1)
+
+            if is_last:
+                x0_est = torch.argmax(probs, dim=-1)
+            else:
+                from model.d3pm_model_cross_attention import _batch_multinomial
+                x0_est = _batch_multinomial(probs)
+
+            hint = x0_est
+
+    return x0_est
+
+
+# ── Interactive demo ──────────────────────────────────────────────────
+
+def interactive_demo():
+    from model.tokenizer import SanskritTokenizer
+
+    cfg    = CONFIG
+    device = torch.device(cfg['training']['device'])
+
+    model_name = cfg['model_type']
+    has_neg    = cfg['data']['include_negative_examples']
+    ckpt       = f"results/{model_name}_neg_{has_neg}/best_model.pt"
+
+    if not os.path.exists(ckpt):
+        raise FileNotFoundError(f"No checkpoint at {ckpt} — train first.")
+
+    model, cfg = load_model(ckpt, cfg, device)
+    model.eval()
+
+    tokenizer = SanskritTokenizer(cfg['model']['vocab_size'])
+    PAD_ID    = tokenizer.tokenizer.token_to_id('[PAD]') or 1
+    MASK_ID   = cfg['diffusion']['mask_token_id']
+
+    print("\n" + "="*55)
+    print("Sanskrit D3PM Paraphrase — type verse, get paraphrase")
+    print("="*55 + "\n")
+
+    while True:
+        try:
+            text = input("INPUT > ").strip()
+        except (EOFError, KeyboardInterrupt):
+            break
+        if not text or text.lower() in ('quit', 'exit', 'q'):
+            break
+
+        ids = torch.tensor(
+            [tokenizer.encode(text)[:cfg['model']['max_seq_len']]],
+            dtype=torch.long, device=device
+        )
+        out   = run_inference(model, ids, cfg)
+        clean = [i for i in out[0].tolist() if i not in (MASK_ID, PAD_ID)]
+        print(f"PARAPHRASE → {tokenizer.decode(clean).strip()}\n")
+
+
+# ── Batch evaluation ──────────────────────────────────────────────────
+
+def batch_evaluate(sample_size=500):
+    from data.dataset import OptimizedSanskritDataset
+    from model.tokenizer import SanskritTokenizer
+
+    cfg    = CONFIG
+    device = torch.device(cfg['training']['device'])
+
+    model_name = cfg['model_type']
+    has_neg    = cfg['data']['include_negative_examples']
+    exp_dir    = f"results/{model_name}_neg_{has_neg}"
+    ckpt       = f"{exp_dir}/best_model.pt"
+
+    if not os.path.exists(ckpt):
+        raise FileNotFoundError(f"No checkpoint at {ckpt}")
+
+    model, cfg = load_model(ckpt, cfg, device)
+    model.eval()
+
+    tokenizer = SanskritTokenizer(cfg['model']['vocab_size'])
+    PAD_ID    = tokenizer.tokenizer.token_to_id('[PAD]') or 1
+    MASK_ID   = cfg['diffusion']['mask_token_id']
+
+    def collate(batch):
+        return {
+            'input_ids':   torch.stack([b['input_ids'].long()  for b in batch]),
+            'target_text': [b['target_text'] for b in batch],
+            'input_text':  [b['input_text']  for b in batch],
+        }
+
+    dataset = OptimizedSanskritDataset('test', tokenizer, cfg['model']['max_seq_len'], cfg)
+    indices = list(range(min(sample_size, len(dataset))))
+    loader  = DataLoader(
+        Subset(dataset, indices),
+        batch_size=cfg['training']['batch_size'],
+        shuffle=False, collate_fn=collate
+    )
+
+    all_preds, all_refs, all_inputs = [], [], []
+    print(f"⏳ Generating {len(indices)} paraphrases …")
+
+    for batch in tqdm(loader):
+        ids = batch['input_ids'].to(device)
+        out = run_inference(model, ids, cfg)
+        for i in range(out.size(0)):
+            clean = [x for x in out[i].tolist() if x not in (MASK_ID, PAD_ID)]
+            all_preds.append(tokenizer.decode(clean).strip())
+            all_refs.append(batch['target_text'][i].strip())
+            all_inputs.append(batch['input_text'][i].strip())
+
+    # Metrics
+    bleu_score, bert_f1 = 0.0, 0.0
+    try:
+        from nltk.translate.bleu_score import corpus_bleu
+        bleu_score = corpus_bleu(
+            [[r.split()] for r in all_refs],
+            [p.split() for p in all_preds]
+        )
+    except Exception:
+        pass
+
+    try:
+        import evaluate as hf_eval
+        res    = hf_eval.load('bertscore').compute(
+            predictions=all_preds, references=all_refs, lang='hi'
+        )
+        bert_f1 = sum(res['f1']) / len(res['f1'])
+    except Exception:
+        pass
+
+    # Save
+    out_path = f"{exp_dir}/evaluation_results.txt"
+    with open(out_path, 'w', encoding='utf-8') as f:
+        f.write(f"Model    : {model_name}\n")
+        f.write(f"Negatives: {has_neg}\n")
+        f.write(f"Steps    : {cfg['inference']['num_steps']}\n")
+        f.write(f"Temp     : {cfg['inference']['temperature']}\n")
+        f.write(f"RepPen   : {cfg['inference']['repetition_penalty']}\n")
+        f.write(f"DivPen   : {cfg['inference']['diversity_penalty']}\n")
+        f.write(f"BLEU     : {bleu_score:.4f}\n")
+        f.write(f"BERTScore: {bert_f1:.4f}\n\n")
+        f.write("=== SAMPLES ===\n")
+        for i in range(min(20, len(all_preds))):
+            f.write(f"IN  : {all_inputs[i]}\n")
+            f.write(f"REF : {all_refs[i]}\n")
+            f.write(f"PRED: {all_preds[i]}\n")
+            f.write("-" * 60 + "\n")
+
+    print(f"\n✅ Results → {out_path}")
+    print(f"📊 BLEU: {bleu_score:.4f}  |  BERTScore: {bert_f1:.4f}")
+    return all_preds, all_refs
+
+
+if __name__ == '__main__':
+    import argparse
+    p = argparse.ArgumentParser()
+    p.add_argument('--mode',    choices=['demo', 'eval'], default='demo')
+    p.add_argument('--samples', type=int, default=500)
+    args = p.parse_args()
+
+    if args.mode == 'demo':
+        interactive_demo()
+    else:
+        batch_evaluate(args.samples)

kv_cache_benchmark.py

@@ -0,0 +1,233 @@
+"""
+analysis/kv_cache_benchmark.py
+================================
+Task 1: Benchmark KV cache vs standard generate().
+
+Measures:
+  - Wall-clock time for generate() vs generate_cached()
+  - Encoder time as % of total generation time (before/after)
+  - Speedup ratio at src_len = 16, 32, 64 tokens
+
+How it works:
+  Standard generate():
+    For each of T=128 steps:
+      src → encoder → memory → decoder → logits    (encoder runs 128 times)
+
+  generate_cached():
+    src → encoder → memory (once)
+    For each of T=128 steps:
+      cached_memory → decoder → logits              (encoder runs 1 time)
+
+  Expected speedup:
+    If encoder = 30% of per-step time:
+      Saved = 127/128 * 30% ≈ 29.7% of total time
+    If encoder = 50% of per-step time:
+      Saved ≈ 49.6% of total time
+
+Usage:
+    python -m analysis.kv_cache_benchmark
+    or:
+    from analysis.kv_cache_benchmark import run_benchmark
+    results = run_benchmark(model, src_tokenizer, device)
+"""
+
+import torch
+import time
+import numpy as np
+from typing import Dict, List
+
+
+def _make_src(src_len: int, src_vocab: int, device: torch.device, batch_size: int = 1):
+    """Create a random source tensor of given length."""
+    # Random real tokens (ids 5..src_vocab-1), padded to src_len
+    ids = torch.randint(5, src_vocab, (batch_size, src_len), device=device)
+    return ids
+
+
+def _time_fn(fn, n_warmup: int = 2, n_runs: int = 5) -> float:
+    """
+    Time a zero-argument callable.
+    Returns mean wall-clock seconds over n_runs after n_warmup warmup calls.
+    """
+    # Warmup
+    for _ in range(n_warmup):
+        fn()
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+        elif torch.backends.mps.is_available():
+            torch.mps.synchronize()
+
+    times = []
+    for _ in range(n_runs):
+        start = time.perf_counter()
+        fn()
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+        elif torch.backends.mps.is_available():
+            torch.mps.synchronize()
+        times.append(time.perf_counter() - start)
+
+    return float(np.mean(times))
+
+
+def benchmark_encoder_cost(
+    model,
+    src:    torch.Tensor,
+) -> Dict[str, float]:
+    """
+    Measure encoder time as a fraction of one full forward pass.
+
+    Returns:
+        encoder_s   : seconds for one encoder call
+        full_step_s : seconds for one full forward_cached decoder step
+        encoder_pct : encoder_s / (encoder_s + full_step_s) * 100
+    """
+    inner = model.model
+    if not hasattr(inner, 'encode_source'):
+        raise ValueError("Model does not support KV cache (not D3PMCrossAttention).")
+
+    device = src.device
+    B      = src.shape[0]
+    T      = inner.scheduler.num_timesteps
+    tgt_len = inner.max_seq_len
+    mask_id = inner.mask_token_id
+
+    x0_est = torch.full((B, tgt_len), mask_id, dtype=torch.long, device=device)
+    t      = torch.zeros(B, dtype=torch.long, device=device)
+
+    # Time encoder alone
+    encoder_s = _time_fn(lambda: inner.encode_source(src))
+
+    # Pre-compute memory for decoder timing
+    memory, src_pad_mask = inner.encode_source(src)
+
+    # Time one decoder step (cached)
+    decoder_s = _time_fn(
+        lambda: inner.forward_cached(memory, src_pad_mask, x0_est, t,
+                                     inference_mode=True)
+    )
+
+    # Time one full step (non-cached = encoder + decoder)
+    full_s = _time_fn(
+        lambda: inner.forward(src, x0_est, t, inference_mode=True)
+    )
+
+    encoder_pct = 100.0 * encoder_s / max(full_s, 1e-9)
+
+    return {
+        "encoder_s":   encoder_s,
+        "decoder_s":   decoder_s,
+        "full_step_s": full_s,
+        "encoder_pct": encoder_pct,
+    }
+
+
+def run_benchmark(
+    model,
+    src_tokenizer,
+    device:        torch.device,
+    src_lens:      List[int] = [16, 32, 64],
+    n_runs:        int       = 5,
+) -> Dict:
+    """
+    Full benchmark: compare generate() vs generate_cached() at multiple src lengths.
+
+    Args:
+        model         : SanskritModel (D3PMCrossAttention)
+        src_tokenizer : SanskritSourceTokenizer
+        device        : torch.device
+        src_lens      : list of source lengths to benchmark
+        n_runs        : number of timing runs per condition
+
+    Returns:
+        results dict with timing and speedup for each src_len
+    """
+    inner = model.model
+    if not hasattr(inner, 'generate_cached'):
+        raise ValueError("Model does not support KV cache (not D3PMCrossAttention).")
+
+    src_vocab = inner.src_embed.token_emb.weight.shape[0]
+    results   = {}
+
+    print("\n" + "=" * 65)
+    print("  KV CACHE BENCHMARK")
+    print("=" * 65)
+    print(f"  {'src_len':>8}  {'standard(s)':>12}  {'cached(s)':>10}  "
+          f"{'speedup':>8}  {'encoder%':>9}")
+    print("-" * 65)
+
+    for src_len in src_lens:
+        src = _make_src(src_len, src_vocab, device)
+
+        # Encoder cost breakdown
+        enc_cost = benchmark_encoder_cost(model, src)
+
+        # Time standard generate() — encoder runs T times
+        def run_standard():
+            return inner.generate(src, temperature=0.8, top_k=40)
+
+        # Time generate_cached() — encoder runs once
+        def run_cached():
+            return inner.generate_cached(src, temperature=0.8, top_k=40)
+
+        t_standard = _time_fn(run_standard, n_warmup=1, n_runs=n_runs)
+        t_cached   = _time_fn(run_cached,   n_warmup=1, n_runs=n_runs)
+        speedup    = t_standard / max(t_cached, 1e-9)
+
+        results[src_len] = {
+            "standard_s":  t_standard,
+            "cached_s":    t_cached,
+            "speedup":     speedup,
+            "encoder_pct": enc_cost["encoder_pct"],
+        }
+
+        print(f"  {src_len:>8}  {t_standard:>12.3f}  {t_cached:>10.3f}  "
+              f"{speedup:>7.2f}x  {enc_cost['encoder_pct']:>8.1f}%")
+
+    print("=" * 65)
+    print(f"\n  Encoder cost = % of one full forward pass")
+    print(f"  Speedup = standard_time / cached_time")
+    print(f"  Expected: speedup ≈ 1 / (1 - encoder_pct/100 * (T-1)/T)")
+
+    return results
+
+
+def print_summary(results: Dict):
+    """Print a human-readable summary of benchmark results."""
+    print("\n  SUMMARY")
+    print("  -------")
+    for src_len, r in results.items():
+        saved_pct = (1.0 - 1.0 / r["speedup"]) * 100
+        print(f"  src_len={src_len}: {r['speedup']:.2f}x speedup "
+              f"({saved_pct:.1f}% time saved, "
+              f"encoder was {r['encoder_pct']:.1f}% of total)")
+
+
+if __name__ == "__main__":
+    import sys, os
+    sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+    from config import CONFIG
+    from inference import load_model
+    from models.tokenizer import SanskritSourceTokenizer
+
+    cfg    = CONFIG
+    device = torch.device(cfg['training']['device'])
+
+    model_name = cfg['model_type']
+    has_neg    = cfg['data']['include_negative_examples']
+    ckpt       = f"results7/{model_name}_neg_{has_neg}/best_model.pt"
+
+    if not os.path.exists(ckpt):
+        print(f"No checkpoint at {ckpt}. Train first.")
+        sys.exit(1)
+
+    model, cfg = load_model(ckpt, cfg, device)
+    model.eval()
+
+    src_tokenizer = SanskritSourceTokenizer(
+        vocab_size = cfg['model'].get('src_vocab_size', 500),
+        max_len    = cfg['model']['max_seq_len'],
+    )
+
+    results = run_benchmark(model, src_tokenizer, device)
+    print_summary(results)

quality_classifier.py

@@ -0,0 +1,514 @@
+"""
+analysis/quality_classifier.py
+================================
+Task 5: Classifier-Free Guidance for Paraphrase Quality Control
+
+Two steps — only Step 2 requires training a SMALL model (not the main D3PM):
+
+STEP 1 — Collect training data (no training):
+  Run existing model on val set, record (hidden_state, CER) pairs.
+  Hidden states come from model.model._last_hidden after forward_cached().
+  CER score = quality label (lower CER = higher quality).
+
+STEP 2 — Train quality classifier:
+  Small 2-layer MLP: d_model → 64 → 1
+  Input: pooled decoder hidden state [B, d_model]
+  Output: predicted quality score in [0, 1]  (1 = high quality)
+  Loss: MSE against normalized CER labels
+  Training time: ~5-10 minutes on CPU for 10k examples
+
+STEP 3 — Guided inference (no retraining):
+  At each diffusion step, use classifier gradient to shift logits:
+    guided_logits = logits + λ * ∂(quality_score)/∂(logits)
+  Higher λ → model biased toward high-quality outputs
+  λ=0 → standard generation (no guidance)
+
+Key: main D3PM model is FROZEN throughout. Only the 10k-param classifier trains.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import os
+import json
+from typing import List, Dict, Optional, Tuple
+
+
+# ── Quality classifier architecture ──────────────────────────────────
+
+class QualityClassifier(nn.Module):
+    """
+    Lightweight MLP that predicts transliteration quality from decoder
+    hidden states.
+
+    Architecture:
+      d_model → 128 → 64 → 1 → Sigmoid
+
+    Input:  mean-pooled decoder hidden state [B, d_model]
+    Output: quality score [B, 1] ∈ [0, 1]  (1 = high quality)
+
+    ~10k parameters. Trains in minutes on CPU.
+    """
+    def __init__(self, d_model: int):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(d_model, 128),
+            nn.ReLU(),
+            nn.Dropout(0.1),
+            nn.Linear(128, 64),
+            nn.ReLU(),
+            nn.Linear(64, 1),
+            nn.Sigmoid(),
+        )
+        self.d_model = d_model
+
+    def forward(self, hidden: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            hidden : [B, tgt_len, d_model] OR [B, d_model] (already pooled)
+
+        Returns:
+            score : [B, 1] quality score in [0, 1]
+        """
+        if hidden.dim() == 3:
+            # Pool over sequence length
+            hidden = hidden.mean(dim=1)   # [B, d_model]
+        return self.net(hidden)           # [B, 1]
+
+
+# ── Training data collection ──────────────────────────────────────────
+
+@torch.no_grad()
+def collect_quality_data(
+    model,
+    src_list:      List[torch.Tensor],
+    ref_list:      List[str],
+    tgt_tokenizer,
+    t_capture:     int   = 0,
+    temperature:   float = 0.8,
+    top_k:         int   = 40,
+    max_samples:   int   = 5000,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """
+    Collect (hidden_state, quality_score) pairs for classifier training.
+
+    For each sample:
+      1. Run generate_cached() on src
+      2. Capture decoder hidden state at t=t_capture
+      3. Compute CER between output and reference
+      4. Quality = 1 - CER  (normalize to [0,1])
+
+    Args:
+        model         : SanskritModel
+        src_list      : list of [1, src_len] tensors
+        ref_list      : list of reference Devanagari strings
+        tgt_tokenizer : SanskritTargetTokenizer
+        t_capture     : which step to capture hidden states (0 = final)
+        max_samples   : cap number of training examples
+
+    Returns:
+        hidden_matrix : np.ndarray [N, d_model]
+        quality_scores: np.ndarray [N]  values in [0, 1]
+    """
+    inner  = model.model
+    T      = inner.scheduler.num_timesteps
+    device = next(inner.parameters()).device
+
+    hidden_list  = []
+    quality_list = []
+    n            = min(len(src_list), max_samples)
+
+    def cer(pred, ref):
+        if not ref:
+            return 1.0
+        def ed(s1, s2):
+            m, n = len(s1), len(s2)
+            dp = list(range(n + 1))
+            for i in range(1, m + 1):
+                prev, dp[0] = dp[0], i
+                for j in range(1, n + 1):
+                    temp = dp[j]
+                    dp[j] = prev if s1[i-1] == s2[j-1] else 1 + min(prev, dp[j], dp[j-1])
+                    prev = temp
+            return dp[n]
+        return ed(pred, ref) / max(len(ref), 1)
+
+    print(f"Collecting quality data from {n} examples...")
+    for i, (src, ref) in enumerate(zip(src_list[:n], ref_list[:n])):
+        if i % 200 == 0:
+            print(f"  {i}/{n}")
+
+        if src.dim() == 1:
+            src = src.unsqueeze(0)
+        src = src.to(device)
+
+        B       = src.shape[0]
+        tgt_len = inner.max_seq_len
+        mask_id = inner.mask_token_id
+
+        memory, src_pad_mask = inner.encode_source(src)
+        x0_est  = torch.full((B, tgt_len), mask_id, dtype=torch.long, device=device)
+        hint    = None
+        h_cap   = None
+
+        for t_val in range(T - 1, -1, -1):
+            t       = torch.full((B,), t_val, dtype=torch.long, device=device)
+            is_last = (t_val == 0)
+
+            logits, _ = inner.forward_cached(
+                memory, src_pad_mask, x0_est, t,
+                x0_hint=hint, inference_mode=True,
+            )
+
+            if t_val == t_capture and hasattr(inner, '_last_hidden'):
+                h_cap = inner._last_hidden[0].mean(dim=0).detach().cpu()  # [d_model]
+
+            logits = logits / max(temperature, 1e-8)
+            if top_k > 0:
+                V = logits.shape[-1]
+                if top_k < V:
+                    vals, _ = torch.topk(logits, top_k, dim=-1)
+                    logits  = logits.masked_fill(logits < vals[..., -1:], float('-inf'))
+
+            probs  = F.softmax(logits, dim=-1)
+            x0_est = torch.argmax(probs, dim=-1) if is_last else _sample(probs)
+            hint   = x0_est
+
+        if h_cap is None:
+            continue
+
+        ids  = [x for x in x0_est[0].tolist() if x > 4]
+        pred = tgt_tokenizer.decode(ids).strip()
+        q    = max(0.0, 1.0 - cer(pred, ref))   # quality = 1 - CER
+
+        hidden_list.append(h_cap.numpy())
+        quality_list.append(q)
+
+    print(f"Collected {len(hidden_list)} quality examples.")
+    print(f"Quality stats: mean={np.mean(quality_list):.3f}  "
+          f"min={np.min(quality_list):.3f}  max={np.max(quality_list):.3f}")
+
+    return np.stack(hidden_list), np.array(quality_list, dtype=np.float32)
+
+
+def _sample(probs):
+    B, L, V = probs.shape
+    flat    = probs.view(B * L, V).clamp(min=1e-9)
+    flat    = flat / flat.sum(dim=-1, keepdim=True)
+    return torch.multinomial(flat, 1).squeeze(-1).view(B, L)
+
+
+# ── Training ──────────────────────────────────────────────────────────
+
+def train_quality_classifier(
+    hidden_matrix:  np.ndarray,
+    quality_scores: np.ndarray,
+    d_model:        int,
+    epochs:         int   = 30,
+    batch_size:     int   = 64,
+    lr:             float = 1e-3,
+    val_frac:       float = 0.1,
+    save_path:      Optional[str] = None,
+) -> QualityClassifier:
+    """
+    Train QualityClassifier on collected (hidden, quality) pairs.
+
+    Args:
+        hidden_matrix  : [N, d_model] from collect_quality_data()
+        quality_scores : [N] quality labels in [0, 1]
+        d_model        : hidden dimension
+        epochs         : training epochs
+        save_path      : if given, save trained classifier weights here
+
+    Returns:
+        trained QualityClassifier
+    """
+    device = torch.device("cpu")   # classifier is tiny, CPU is fine
+
+    X = torch.tensor(hidden_matrix, dtype=torch.float32)
+    y = torch.tensor(quality_scores, dtype=torch.float32).unsqueeze(-1)
+
+    N     = len(X)
+    n_val = max(1, int(N * val_frac))
+    idx   = torch.randperm(N)
+    val_idx   = idx[:n_val]
+    train_idx = idx[n_val:]
+
+    X_train, y_train = X[train_idx], y[train_idx]
+    X_val,   y_val   = X[val_idx],   y[val_idx]
+
+    clf       = QualityClassifier(d_model).to(device)
+    optimizer = torch.optim.Adam(clf.parameters(), lr=lr)
+
+    print(f"\nTraining QualityClassifier: {sum(p.numel() for p in clf.parameters())} params")
+    print(f"Train: {len(X_train)}  Val: {len(X_val)}")
+
+    best_val_loss = float('inf')
+    best_state    = None
+
+    for epoch in range(epochs):
+        clf.train()
+        perm       = torch.randperm(len(X_train))
+        train_loss = 0.0
+        n_batches  = 0
+
+        for start in range(0, len(X_train), batch_size):
+            batch_idx = perm[start:start + batch_size]
+            xb, yb    = X_train[batch_idx], y_train[batch_idx]
+            pred      = clf(xb)
+            loss      = F.mse_loss(pred, yb)
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            train_loss += loss.item()
+            n_batches  += 1
+
+        clf.eval()
+        with torch.no_grad():
+            val_pred = clf(X_val)
+            val_loss = F.mse_loss(val_pred, y_val).item()
+
+        if epoch % 5 == 0 or epoch == epochs - 1:
+            print(f"  Ep {epoch+1:3d}  train={train_loss/n_batches:.4f}  val={val_loss:.4f}")
+
+        if val_loss < best_val_loss:
+            best_val_loss = val_loss
+            best_state    = {k: v.clone() for k, v in clf.state_dict().items()}
+
+    if best_state:
+        clf.load_state_dict(best_state)
+        print(f"  Best val loss: {best_val_loss:.4f}")
+
+    if save_path:
+        os.makedirs(os.path.dirname(save_path) or ".", exist_ok=True)
+        torch.save(clf.state_dict(), save_path)
+        print(f"  Classifier saved: {save_path}")
+
+    return clf
+
+
+# ── Guided inference ──────────────────────────────────────────────────
+
+def generate_guided(
+    model,
+    src:        torch.Tensor,
+    classifier: QualityClassifier,
+    guidance_scale: float = 1.0,
+    temperature:    float = 0.8,
+    top_k:          int   = 40,
+) -> torch.Tensor:
+    """
+    Classifier-guided generation.
+
+    At each diffusion step:
+      1. Run forward_cached() → logits, hidden states
+      2. Compute classifier gradient: ∂(quality_score) / ∂(hidden)
+      3. Project gradient back to logit space (approximate)
+      4. guided_logits = logits + λ * gradient_signal
+      5. Sample from guided_logits
+
+    guidance_scale λ:
+      0.0 → no guidance (standard generation)
+      0.5 → weak guidance
+      1.0 → moderate guidance (recommended starting point)
+      2.0 → strong guidance (may reduce diversity)
+      3.0 → very strong (may collapse to repetitive output)
+
+    Args:
+        model           : SanskritModel (frozen)
+        src             : [1, src_len] IAST token ids
+        classifier      : trained QualityClassifier
+        guidance_scale  : λ — guidance strength
+
+    Returns:
+        x0_est : [1, tgt_len] generated token ids
+    """
+    inner  = model.model
+    T      = inner.scheduler.num_timesteps
+    device = next(inner.parameters()).device
+    clf_device = next(classifier.parameters()).device
+
+    if src.dim() == 1:
+        src = src.unsqueeze(0)
+    src = src.to(device)
+
+    B       = src.shape[0]
+    tgt_len = inner.max_seq_len
+    mask_id = inner.mask_token_id
+
+    memory, src_pad_mask = inner.encode_source(src)
+    x0_est  = torch.full((B, tgt_len), mask_id, dtype=torch.long, device=device)
+    hint    = None
+
+    inner.eval()
+    classifier.eval()
+
+    for t_val in range(T - 1, -1, -1):
+        t       = torch.full((B,), t_val, dtype=torch.long, device=device)
+        is_last = (t_val == 0)
+
+        if guidance_scale > 0.0:
+            # Need gradients for classifier guidance
+            with torch.enable_grad():
+                # Run forward_cached and get hidden states
+                PAD = 1
+                if t_val > 0:
+                    _, x_t_ids = inner.forward_process.q_sample(x0_est, t)
+                else:
+                    x_t_ids = x0_est
+
+                x      = inner.tgt_embed(x_t_ids)
+                t_norm = t.float() / T
+                t_emb  = inner.time_mlp(t_norm.unsqueeze(-1))
+                x      = x + t_emb.unsqueeze(1)
+
+                if hint is not None:
+                    hint_emb = inner.tgt_embed(hint)
+                    gate     = inner.hint_gate(x)
+                    x        = x + gate * hint_emb
+
+                for block in inner.decoder_blocks:
+                    x = block(x, memory, tgt_pad_mask=None, src_pad_mask=src_pad_mask)
+
+                # hidden: [B, tgt_len, d_model] — detach from graph for clf
+                hidden = x.detach().requires_grad_(True).to(clf_device)
+
+                # Classifier quality score
+                quality = classifier(hidden)   # [B, 1]
+                quality.sum().backward()
+
+                # Gradient of quality w.r.t. hidden: [B, tgt_len, d_model]
+                grad = hidden.grad.to(device)   # [B, tgt_len, d_model]
+
+                # Project gradient to logit space via output head weight
+                # logit_grad ≈ grad @ head.weight   [B, tgt_len, tgt_vocab]
+                logit_grad = grad @ inner.head.weight.T
+
+                # Compute standard logits (no gradient needed)
+                with torch.no_grad():
+                    logits = inner.head(x)
+
+                # Apply guidance
+                logits = logits + guidance_scale * logit_grad
+
+        else:
+            with torch.no_grad():
+                logits, _ = inner.forward_cached(
+                    memory, src_pad_mask, x0_est, t,
+                    x0_hint=hint, inference_mode=True,
+                )
+
+        with torch.no_grad():
+            logits = logits / max(temperature, 1e-8)
+            if top_k > 0:
+                V = logits.shape[-1]
+                if top_k < V:
+                    vals, _ = torch.topk(logits, top_k, dim=-1)
+                    logits  = logits.masked_fill(logits < vals[..., -1:], float('-inf'))
+
+            probs  = F.softmax(logits, dim=-1)
+            x0_est = torch.argmax(probs, dim=-1) if is_last else _sample_no_grad(probs)
+            hint   = x0_est
+
+    return x0_est
+
+
+def _sample_no_grad(probs):
+    B, L, V = probs.shape
+    flat    = probs.view(B * L, V).clamp(min=1e-9)
+    flat    = flat / flat.sum(dim=-1, keepdim=True)
+    return torch.multinomial(flat, 1).squeeze(-1).view(B, L)
+
+
+# ── Guidance scale sweep ──────────────────────────────────────────────
+
+def sweep_guidance_scales(
+    model,
+    classifier: QualityClassifier,
+    src_list:   List[torch.Tensor],
+    ref_list:   List[str],
+    tgt_tokenizer,
+    scales:     List[float] = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0],
+    n_samples:  int         = 50,
+    device:     torch.device = None,
+    output_dir: str          = "analysis/outputs",
+) -> Dict:
+    """
+    Evaluate CER at each guidance scale.
+    Produces quality-diversity tradeoff plot.
+    """
+    def cer(pred, ref):
+        if not ref:
+            return 1.0
+        def ed(s1, s2):
+            m, n = len(s1), len(s2)
+            dp = list(range(n + 1))
+            for i in range(1, m + 1):
+                prev, dp[0] = dp[0], i
+                for j in range(1, n + 1):
+                    temp = dp[j]
+                    dp[j] = prev if s1[i-1] == s2[j-1] else 1 + min(prev, dp[j], dp[j-1])
+                    prev = temp
+            return dp[n]
+        return ed(pred, ref) / max(len(ref), 1)
+
+    device  = device or next(model.parameters()).device
+    results = {}
+    n       = min(n_samples, len(src_list))
+
+    print("\nGuidance scale sweep...")
+    for scale in scales:
+        cer_list   = []
+        output_set = []
+        for src, ref in zip(src_list[:n], ref_list[:n]):
+            if src.dim() == 1:
+                src = src.unsqueeze(0)
+            out      = generate_guided(model, src.to(device), classifier,
+                                        guidance_scale=scale)
+            ids      = [x for x in out[0].tolist() if x > 4]
+            pred     = tgt_tokenizer.decode(ids).strip()
+            cer_list.append(cer(pred, ref))
+            output_set.append(pred)
+
+        mean_cer = float(np.mean(cer_list))
+
+        # Self-diversity: unique outputs / total (proxy for diversity)
+        unique_frac = len(set(output_set)) / max(len(output_set), 1)
+
+        results[scale] = {"mean_cer": mean_cer, "diversity": unique_frac}
+        print(f"  λ={scale:.1f}  CER={mean_cer:.4f}  diversity={unique_frac:.3f}")
+
+    # Plot
+    os.makedirs(output_dir, exist_ok=True)
+    try:
+        import matplotlib.pyplot as plt
+        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
+
+        sc_list  = sorted(results.keys())
+        cers     = [results[s]["mean_cer"]   for s in sc_list]
+        diversities = [results[s]["diversity"] for s in sc_list]
+
+        ax1.plot(sc_list, cers, 'o-', color='coral', linewidth=1.8, markersize=7)
+        ax1.set_xlabel("Guidance scale λ", fontsize=10)
+        ax1.set_ylabel("CER (↓ better)", fontsize=10)
+        ax1.set_title("Quality vs guidance scale", fontsize=10)
+
+        ax2.plot(sc_list, diversities, 'o-', color='steelblue', linewidth=1.8, markersize=7)
+        ax2.set_xlabel("Guidance scale λ", fontsize=10)
+        ax2.set_ylabel("Output diversity (unique fraction)", fontsize=10)
+        ax2.set_title("Diversity vs guidance scale", fontsize=10)
+
+        plt.suptitle("Quality-Diversity Tradeoff (Guidance Scale Sweep)", fontsize=11)
+        plt.tight_layout()
+        path = os.path.join(output_dir, "guidance_scale_sweep.png")
+        plt.savefig(path, dpi=150, bbox_inches='tight')
+        plt.close()
+        print(f"  Saved: {path}")
+    except ImportError:
+        pass
+
+    with open(os.path.join(output_dir, "guidance_results.json"), "w") as f:
+        json.dump({str(k): v for k, v in results.items()}, f, indent=2)
+
+    return results

reverse_process.py

@@ -0,0 +1,302 @@
+"""
+reverse_process.py — Fixed
+===========================
+Two bugs fixed from the original:
+
+BUG 1 (critical): generate_beam() passed x_t (noisy) as `tgt` to model.
+  The model does q_sample(tgt, t) internally — so x_t got double-noised.
+  Fix: pass x0_estimate (current clean guess) as tgt. Model noises it correctly.
+
+BUG 2: apply_diversity_penalty used logits.var(dim=-1) — this adds the
+  variance of each position's own distribution back to itself, which is
+  mathematically meaningless and just injects noise.
+  Fix: penalize tokens that are uniformly high-probability across ALL positions
+  (global common tokens). This genuinely promotes diversity.
+"""
+
+import torch
+import torch.nn.functional as F
+
+
+class ReverseDiffusion:
+    def __init__(self, scheduler):
+        self.scheduler = scheduler
+
+    def p_sample_step(
+        self,
+        model,
+        x_t,
+        t,
+        condition,
+        beam_width=3,
+        temperature=1.0,
+        repetition_penalty=1.2,
+        diversity_penalty=0.3
+    ):
+        """
+        Single reverse step with temperature + penalties.
+        """
+
+        with torch.no_grad():
+
+            # ---- Shape safety ----
+            if x_t.dim() == 1:
+                x_t = x_t.unsqueeze(0)
+
+            if condition.dim() == 1:
+                condition = condition.unsqueeze(0)
+
+            if t.dim() == 0:
+                t = t.unsqueeze(0)
+
+            if t.shape[0] != x_t.shape[0]:
+                t = t.expand(x_t.shape[0])
+
+            # ---- Model forward ----
+            logits, _ = model(condition, x_t, t)
+
+            # ---- Temperature scaling ----
+            logits = logits / temperature
+
+            # ---- Repetition penalty (FIXED VERSION) ----
+            if repetition_penalty != 1.0:
+                logits = apply_repetition_penalty(
+                    logits, x_t, repetition_penalty
+                )
+
+            # ---- Diversity penalty ----
+            if diversity_penalty > 0:
+                logits = apply_diversity_penalty(
+                    logits, diversity_penalty
+                )
+
+            probs = F.softmax(logits, dim=-1)
+
+            B, L, V = probs.shape
+
+            # ---- Top-k beam expansion ----
+            topk_probs, topk_ids = torch.topk(
+                probs, beam_width, dim=-1
+            )
+
+            candidates = []
+
+            for k in range(beam_width):
+                next_tokens = topk_ids[:, :, k]
+                score = torch.log(
+                    topk_probs[:, :, k] + 1e-9
+                ).sum()
+                candidates.append((next_tokens, score))
+
+            return candidates
+
+    def generate_beam(
+        self,
+        model,
+        condition,
+        beam_width=3,
+        num_steps=None,
+        temperature=1.0,
+        repetition_penalty=1.2,
+        diversity_penalty=0.3
+    ):
+        """
+        Beam-search reverse diffusion with temperature.
+        """
+
+        if num_steps is None:
+            num_steps = self.scheduler.num_timesteps
+
+        device = condition.device
+
+        if condition.dim() == 1:
+            condition = condition.unsqueeze(0)
+
+        B, L = condition.shape
+
+        # 🔥 Better initialization: start from MASK
+        x_init = torch.full(
+            (B, L),
+            fill_value=model.mask_token_id,
+            dtype=torch.long,
+            device=device
+        )
+
+        beams = [(x_init, 0.0)]
+
+        for step in reversed(range(num_steps)):
+
+            new_beams = []
+
+            for x_t, score in beams:
+
+                t_tensor = torch.full(
+                    (B,),
+                    step,
+                    dtype=torch.long,
+                    device=device
+                )
+
+                candidates = self.p_sample_step(
+                    model,
+                    x_t,
+                    t_tensor,
+                    condition,
+                    beam_width,
+                    temperature,
+                    repetition_penalty,
+                    diversity_penalty
+                )
+
+                for tokens, new_score in candidates:
+                    new_beams.append(
+                        (tokens, score + new_score)
+                    )
+
+            # ---- Keep top beams ----
+            new_beams = sorted(
+                new_beams,
+                key=lambda x: x[1],
+                reverse=True
+            )
+
+            beams = new_beams[:beam_width]
+
+        best_tokens, best_score = beams[0]
+        return best_tokens
+
+
+
+    def generate(
+        self,
+        model,
+        condition,
+        num_steps=None,
+        temperature=0.8,
+        top_k=50,
+        repetition_penalty=1.2,
+        diversity_penalty=0.0,
+    ):
+        """
+        Correct D3PM iterative refinement.
+
+        x0_est starts as all [MASK].
+        Each step: forward(src=condition, tgt=x0_est, t)
+          → model applies q_sample(x0_est, t) internally
+          → predicts cleaner x0
+          → x0_est updated
+
+        diversity_penalty: reduces probability of tokens that are
+        globally dominant across all sequence positions (not logits.var()).
+        """
+        if num_steps is None:
+            num_steps = self.scheduler.num_timesteps
+
+        device = condition.device
+        if condition.dim() == 1:
+            condition = condition.unsqueeze(0)
+        B, L = condition.shape
+
+        T         = self.scheduler.num_timesteps
+        step_size = max(1, T // num_steps)
+        timesteps = list(range(T - 1, -1, -step_size))
+        if timesteps[-1] != 0:
+            timesteps.append(0)
+
+        mask_id = model.mask_token_id
+        # Start: know nothing → all MASK is our initial clean estimate
+        x0_est = torch.full((B, L), mask_id, dtype=torch.long, device=device)
+        hint   = None
+
+        model.eval()
+        with torch.no_grad():
+            for step_idx, t_val in enumerate(timesteps):
+                t       = torch.full((B,), t_val, dtype=torch.long, device=device)
+                is_last = (step_idx == len(timesteps) - 1)
+
+                # KEY: pass x0_est as tgt — model noises it internally
+                import inspect
+                sig = inspect.signature(model.forward).parameters
+                if 'x0_hint' in sig:
+                    outputs = model(condition, x0_est, t, x0_hint=hint)
+                else:
+                    outputs = model(condition, x0_est, t)
+
+                logits = outputs[0] if isinstance(outputs, tuple) else outputs
+
+                # Repetition penalty: down-weight tokens already in sequence
+                if repetition_penalty != 1.0:
+                    logits = apply_repetition_penalty(logits, x0_est, repetition_penalty)
+
+                # Diversity penalty: reduce globally dominant tokens
+                if diversity_penalty > 0.0:
+                    logits = apply_diversity_penalty(logits, diversity_penalty)
+
+                # Temperature + top-k
+                logits = logits / max(temperature, 1e-5)
+                if top_k > 0:
+                    logits = top_k_filter(logits, top_k)
+
+                probs = F.softmax(logits, dim=-1)
+
+                if is_last:
+                    x0_est = torch.argmax(probs, dim=-1)
+                else:
+                    x0_est = batch_multinomial(probs)
+
+                hint = x0_est
+
+        return x0_est
+
+
+# ── Penalty functions ─────────────────────────────────────────────────
+
+def apply_repetition_penalty(logits, prev_tokens, penalty=1.2):
+    """
+    Down-weight tokens that already appear in the current sequence.
+    Prevents मनो मनो मनो repetition loops.
+    penalty=1.0 → no effect
+    penalty=1.2 → mild suppression of repeated tokens
+    penalty=2.0 → strong suppression
+    """
+    B, L, V = logits.shape
+    for b in range(B):
+        for token_id in set(prev_tokens[b].tolist()):
+            if token_id > 4:   # don't penalize special tokens
+                logits[b, :, token_id] = logits[b, :, token_id] / penalty
+    return logits
+
+
+def apply_diversity_penalty(logits, penalty=0.5):
+    """
+    Correct diversity penalty: penalize tokens that are globally dominant
+    across ALL sequence positions. This forces the model to use less
+    common tokens, increasing output diversity.
+
+    Method: compute mean probability across positions, subtract penalty
+    times that mean. Tokens uniformly high everywhere get suppressed.
+
+    penalty=0.0 → no diversity enforcement
+    penalty=0.5 → moderate diversity
+    penalty=1.0 → strong diversity (may hurt coherence)
+    """
+    # Mean logit across all positions: [B, V]
+    global_mean = logits.mean(dim=1, keepdim=True)   # [B, 1, V]
+    # Subtract scaled global mean — suppresses globally common tokens
+    return logits - penalty * global_mean
+
+
+def top_k_filter(logits, k):
+    B, L, V = logits.shape
+    if k >= V:
+        return logits
+    topk_vals, _ = torch.topk(logits, k, dim=-1)
+    threshold = topk_vals[..., -1].unsqueeze(-1)
+    return logits.masked_fill(logits < threshold, float('-inf'))
+
+
+def batch_multinomial(probs):
+    B, L, V = probs.shape
+    flat = probs.view(B * L, V) + 1e-9
+    flat = flat / flat.sum(dim=-1, keepdim=True)
+    return torch.multinomial(flat, 1).squeeze(-1).view(B, L)

reverse_process1.py

@@ -0,0 +1,154 @@
+import torch
+import torch.nn.functional as F
+
+
+class ReverseDiffusion:
+    """
+    Stable reverse diffusion with:
+    - Beam search
+    - Self conditioning
+    - Temperature sampling
+    - Repetition penalty
+    - Diversity penalty
+    """
+
+    def __init__(self, scheduler):
+
+        self.scheduler = scheduler
+
+        self.temperature = 0.75
+        self.repetition_penalty = 1.15
+        self.diversity_penalty = 0.0
+        self.length_penalty = 1.0
+
+    # ------------------------------------------------
+    # penalties
+    # ------------------------------------------------
+
+    def apply_repetition_penalty(self, logits, tokens):
+
+        B, L, V = logits.shape
+
+        for b in range(B):
+
+            used = set(tokens[b].tolist())
+
+            for token_id in used:
+                logits[b, :, token_id] /= self.repetition_penalty
+
+        return logits
+
+    def apply_diversity_penalty(self, logits):
+
+        if self.diversity_penalty == 0:
+            return logits
+
+        logits_var = logits.var(dim=-1, keepdim=True)
+        return logits + self.diversity_penalty * logits_var
+
+    # ------------------------------------------------
+    # single reverse step
+    # ------------------------------------------------
+
+    def p_sample_step(self, model, x_t, t, condition, self_cond=None, beam_width=3):
+
+        with torch.no_grad():
+
+            logits, hidden = model(condition, x_t, t, self_cond)
+
+            logits = logits / self.temperature
+
+            logits = self.apply_repetition_penalty(logits, x_t)
+            logits = self.apply_diversity_penalty(logits)
+
+            probs = F.softmax(logits, dim=-1)
+
+            B, L, V = probs.shape
+
+            topk_probs, topk_ids = torch.topk(probs, beam_width, dim=-1)
+
+            candidates = []
+
+            for k in range(beam_width):
+
+                tokens = topk_ids[:, :, k]
+
+                score = torch.log(topk_probs[:, :, k] + 1e-9).sum()
+
+                candidates.append((tokens, score))
+
+            return candidates
+
+    # ------------------------------------------------
+    # beam reverse diffusion
+    # ------------------------------------------------
+
+    def generate_beam(self, model, condition, beam_width=3, num_steps=None):
+
+        if num_steps is None:
+            num_steps = self.scheduler.num_timesteps
+
+        device = condition.device
+
+        if condition.dim() == 1:
+            condition = condition.unsqueeze(0)
+
+        B, L = condition.shape
+
+        # ------------------------------------------------
+        # BETTER LATENT INITIALIZATION
+        # ------------------------------------------------
+
+        x_init = condition.clone()
+
+        mask = torch.rand_like(x_init.float()) < 0.5
+        x_init[mask] = model.mask_token_id
+
+        beams = [(x_init, 0.0)]
+
+        self_cond = None
+
+        for step in reversed(range(num_steps)):
+
+            new_beams = []
+
+            for x_t, score in beams:
+
+                t_tensor = torch.full(
+                    (B,),
+                    step,
+                    dtype=torch.long,
+                    device=device
+                )
+
+                candidates = self.p_sample_step(
+                    model,
+                    x_t,
+                    t_tensor,
+                    condition,
+                    self_cond,
+                    beam_width
+                )
+
+                for tokens, new_score in candidates:
+
+                    length_norm = tokens.shape[1] ** self.length_penalty
+
+                    final_score = (score + new_score) / length_norm
+
+                    new_beams.append((tokens, final_score))
+
+            new_beams = sorted(
+                new_beams,
+                key=lambda x: x[1],
+                reverse=True
+            )
+
+            beams = new_beams[:beam_width]
+
+            # self conditioning
+            self_cond = beams[0][0]
+
+        best_tokens, best_score = beams[0]
+
+        return best_tokens

reverse_process2.py

@@ -0,0 +1,275 @@
+"""
+reverse_process.py  — Final Correct Version
+=============================================
+
+KEY PRINCIPLE: generate() must be byte-for-byte identical to run_inference()
+in inference.py, which is what produced BERTScore 0.75 at validation.
+
+CRITICAL BUG IN PREVIOUS VERSION:
+  We passed inference_mode=True to the model, but the model was NEVER
+  called with inference_mode=True during training or validation.
+  run_inference() (the validated path) does:
+      model(input_ids, x0_est, t, x0_hint=hint)
+  → inference_mode defaults to False.
+
+  With inference_mode=True the model does two things differently:
+    1. tgt_pad_mask = None  (training used tgt_pad_mask = tgt==PAD)
+    2. Skips q_sample at t=0 (training always called q_sample)
+  The model was never trained to handle these conditions → garbage output.
+
+  Fix: do NOT pass inference_mode. Let it default to False, exactly
+  as run_inference() did.
+
+BUGS FIXED (vs original reverse_process.py)
+--------------------------------------------
+BUG 1  generate_beam() used for D3PM → all-Ṛ repetition.
+       Use generate() (iterative refinement) from app1.py instead.
+BUG 2  apply_diversity_penalty used logits.var() → noise injection.
+       Fixed to logits - penalty * logits.mean(dim=1) — global suppression.
+BUG 3  x0_hint (self-conditioning) never passed to model.
+       Fixed: generate() passes x0_hint=hint every step.
+BUG 4  params not forwarded from generate_beam() to p_sample_step().
+       Fixed in generate_beam() (kept for reference, not for production use).
+"""
+
+import torch
+import torch.nn.functional as F
+
+
+class ReverseDiffusion:
+
+    def __init__(self, scheduler):
+        self.scheduler = scheduler
+
+        # Attribute-style defaults for backward compat with any code
+        # that sets  reverse_diffusion.temperature = 0.9 etc.
+        # generate() prefers explicit kwargs and falls back to these.
+        self.temperature        = 0.75
+        self.repetition_penalty = 1.15
+        self.diversity_penalty  = 0.0
+        self.top_k              = 50
+
+    # ------------------------------------------------------------------ #
+    #  generate  — CORRECT D3PM iterative refinement                      #
+    #  Exact equivalent of run_inference() in inference.py                #
+    # ------------------------------------------------------------------ #
+    def generate(
+        self,
+        model,
+        condition,
+        num_steps          = None,
+        temperature        = None,
+        top_k              = None,
+        repetition_penalty = None,
+        diversity_penalty  = None,
+    ):
+        """
+        D3PM iterative refinement — identical to run_inference() in inference.py,
+        which is the validated path (BERTScore 0.75).
+
+        Algorithm:
+          x0_est = all [MASK]
+          for t = T-1 down to 0:
+            logits = model(src, x0_est, t, x0_hint=hint)
+                     ↑ inference_mode NOT passed (defaults to False)
+                     ↑ this exactly matches training/validation
+            apply penalties, temperature, top_k
+            if t > 0: x0_est = multinomial(softmax(logits))   ← stochastic
+            if t = 0: x0_est = argmax(softmax(logits))         ← deterministic
+            hint = x0_est
+        """
+        # Resolve: explicit kwarg > object attribute
+        temperature        = temperature        if temperature        is not None else self.temperature
+        top_k              = top_k              if top_k              is not None else self.top_k
+        repetition_penalty = repetition_penalty if repetition_penalty is not None else self.repetition_penalty
+        diversity_penalty  = diversity_penalty  if diversity_penalty  is not None else self.diversity_penalty
+
+        if num_steps is None:
+            num_steps = self.scheduler.num_timesteps
+
+        device = condition.device
+        if condition.dim() == 1:
+            condition = condition.unsqueeze(0)
+        B, L = condition.shape
+
+        T         = self.scheduler.num_timesteps
+        step_size = max(1, T // num_steps)
+        timesteps = list(range(T - 1, -1, -step_size))
+        if timesteps[-1] != 0:
+            timesteps.append(0)
+
+        mask_id = model.mask_token_id
+        x0_est  = torch.full((B, L), mask_id, dtype=torch.long, device=device)
+        hint    = None
+
+        model.eval()
+        with torch.no_grad():
+            for step_idx, t_val in enumerate(timesteps):
+                t       = torch.full((B,), t_val, dtype=torch.long, device=device)
+                is_last = (step_idx == len(timesteps) - 1)
+
+                # ── CRITICAL: do NOT pass inference_mode ──────────────────
+                # inference_mode defaults to False inside SanskritModel /
+                # D3PMCrossAttention. This matches run_inference() exactly.
+                # Passing inference_mode=True changes tgt_pad_mask and
+                # q_sample behaviour — the model was never trained for that.
+                logits, _ = model(condition, x0_est, t, x0_hint=hint)
+
+                # Repetition penalty
+                if repetition_penalty != 1.0:
+                    logits = apply_repetition_penalty(
+                        logits, x0_est, repetition_penalty
+                    )
+
+                # Diversity penalty (correct: global mean suppression)
+                if diversity_penalty > 0.0:
+                    logits = apply_diversity_penalty(logits, diversity_penalty)
+
+                logits = logits / max(temperature, 1e-5)
+
+                if top_k > 0:
+                    logits = top_k_filter(logits, top_k)
+
+                probs = F.softmax(logits, dim=-1)
+
+                # Stochastic at every step except the last (argmax at t=0)
+                if is_last:
+                    x0_est = torch.argmax(probs, dim=-1)
+                else:
+                    x0_est = batch_multinomial(probs)
+
+                hint = x0_est
+
+        return x0_est   # (B, L)
+
+    # ------------------------------------------------------------------ #
+    #  p_sample_step  — used by generate_beam (not for production)        #
+    # ------------------------------------------------------------------ #
+    def p_sample_step(
+        self,
+        model,
+        x_t,
+        t,
+        condition,
+        beam_width         = 3,
+        temperature        = 1.0,
+        repetition_penalty = 1.2,
+        diversity_penalty  = 0.3,
+        x0_hint            = None,
+    ):
+        with torch.no_grad():
+            if x_t.dim() == 1:       x_t       = x_t.unsqueeze(0)
+            if condition.dim() == 1: condition  = condition.unsqueeze(0)
+            if t.dim() == 0:         t          = t.unsqueeze(0)
+            if t.shape[0] != x_t.shape[0]:
+                t = t.expand(x_t.shape[0])
+
+            # No inference_mode — matches training convention
+            logits, _ = model(condition, x_t, t, x0_hint=x0_hint)
+
+            logits = logits / max(temperature, 1e-5)
+
+            if repetition_penalty != 1.0:
+                logits = apply_repetition_penalty(logits, x_t, repetition_penalty)
+            if diversity_penalty > 0.0:
+                logits = apply_diversity_penalty(logits, diversity_penalty)
+
+            probs = F.softmax(logits, dim=-1)
+            B, L, V = probs.shape
+
+            topk_probs, topk_ids = torch.topk(probs, beam_width, dim=-1)
+            candidates = []
+            for k in range(beam_width):
+                next_tokens = topk_ids[:, :, k]
+                score       = torch.log(topk_probs[:, :, k] + 1e-9).sum()
+                candidates.append((next_tokens, score))
+            return candidates
+
+    # ------------------------------------------------------------------ #
+    #  generate_beam  — kept for reference; NOT the correct D3PM method   #
+    # ------------------------------------------------------------------ #
+    def generate_beam(
+        self,
+        model,
+        condition,
+        beam_width         = 3,
+        num_steps          = None,
+        temperature        = None,
+        repetition_penalty = None,
+        diversity_penalty  = None,
+    ):
+        """
+        WARNING: do NOT call this from app1.py for D3PM generation.
+        generate_beam() forces every position to the same top-k token
+        → all-Ṛ / all-rud repetition. Use generate() instead.
+        Kept only for experimental reference.
+        """
+        temperature        = temperature        if temperature        is not None else self.temperature
+        repetition_penalty = repetition_penalty if repetition_penalty is not None else self.repetition_penalty
+        diversity_penalty  = diversity_penalty  if diversity_penalty  is not None else self.diversity_penalty
+        if num_steps is None:
+            num_steps = self.scheduler.num_timesteps
+
+        device = condition.device
+        if condition.dim() == 1: condition = condition.unsqueeze(0)
+        B, L = condition.shape
+
+        x_init = torch.full((B, L), fill_value=model.mask_token_id,
+                            dtype=torch.long, device=device)
+        beams     = [(x_init, 0.0)]
+        best_hint = None
+
+        for step in reversed(range(num_steps)):
+            t_tensor  = torch.full((B,), step, dtype=torch.long, device=device)
+            new_beams = []
+            for x_t, score in beams:
+                candidates = self.p_sample_step(
+                    model, x_t, t_tensor, condition,
+                    beam_width         = beam_width,
+                    temperature        = temperature,
+                    repetition_penalty = repetition_penalty,
+                    diversity_penalty  = diversity_penalty,
+                    x0_hint            = best_hint,
+                )
+                for tokens, new_score in candidates:
+                    new_beams.append((tokens, score + new_score.item()))
+
+            new_beams = sorted(new_beams, key=lambda x: x[1], reverse=True)
+            beams     = new_beams[:beam_width]
+            best_hint = beams[0][0]
+
+        return beams[0][0]   # (B, L)
+
+
+# ── Penalty helpers ────────────────────────────────────────────────────────
+
+def apply_repetition_penalty(logits, prev_tokens, penalty=1.2):
+    """Down-weight tokens already present in the sequence."""
+    for b in range(logits.shape[0]):
+        for token_id in set(prev_tokens[b].tolist()):
+            if token_id > 4:
+                logits[b, :, token_id] = logits[b, :, token_id] / penalty
+    return logits
+
+
+def apply_diversity_penalty(logits, penalty=0.3):
+    """
+    Correct diversity penalty: suppress globally dominant tokens.
+    logits -= penalty * mean(logits, dim=1)  [sequence dimension]
+    """
+    global_mean = logits.mean(dim=1, keepdim=True)   # [B, 1, V]
+    return logits - penalty * global_mean
+
+
+def top_k_filter(logits, k):
+    B, L, V = logits.shape
+    if k >= V: return logits
+    topk_vals, _ = torch.topk(logits, k, dim=-1)
+    return logits.masked_fill(logits < topk_vals[..., -1].unsqueeze(-1), float('-inf'))
+
+
+def batch_multinomial(probs):
+    B, L, V = probs.shape
+    flat = probs.view(B * L, V) + 1e-9
+    flat = flat / flat.sum(dim=-1, keepdim=True)
+    return torch.multinomial(flat, 1).squeeze(-1).view(B, L)

run_analysis.py

@@ -0,0 +1,407 @@
+"""
+analysis/run_analysis.py
+=========================
+Entry point for all 5 tasks.
+
+Tasks:
+  Task 1 — KV Cache benchmark          (no retraining)
+  Task 2 — Attention viz + drift        (no retraining)
+  Task 3 — Concept vectors + PCA steer  (no retraining)
+  Task 4 — Step ablation                (REQUIRES retraining for each T)
+  Task 5 — Classifier-free guidance     (trains small 10k-param classifier)
+
+Usage:
+  python analysis/run_analysis.py --task 1
+  python analysis/run_analysis.py --task 2 --input "dharmo rakṣati rakṣitaḥ"
+  python analysis/run_analysis.py --task 3
+  python analysis/run_analysis.py --task 4 --phase generate_configs
+  python analysis/run_analysis.py --task 4 --phase analyze
+  python analysis/run_analysis.py --task 5
+  python analysis/run_analysis.py --task all --input "satyameva jayate"
+
+Output files: analysis/outputs/
+"""
+
+import torch
+import os, sys, argparse, json
+import numpy as np
+
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from config import CONFIG
+from inference import load_model
+from model.tokenizer import SanskritSourceTokenizer, SanskritTargetTokenizer
+
+OUTPUT_DIR = "analysis/outputs"
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+
+
+# ── Shared loader ─────────────────────────────────────────────────────
+
+def load_everything(cfg, device):
+    model_name = cfg['model_type']
+    has_neg    = cfg['data']['include_negative_examples']
+    ckpt       = f"results7/{model_name}_neg_{has_neg}/best_model.pt"
+    if not os.path.exists(ckpt):
+        raise FileNotFoundError(f"No checkpoint at {ckpt}. Train first.")
+    model, cfg = load_model(ckpt, cfg, device)
+    model.eval()
+    src_tok = SanskritSourceTokenizer(
+        vocab_size=cfg['model'].get('src_vocab_size', 500),
+        max_len=cfg['model']['max_seq_len'])
+    tgt_tok = SanskritTargetTokenizer(
+        vocab_size=cfg['model'].get('tgt_vocab_size', 500),
+        max_len=cfg['model']['max_seq_len'])
+    return model, src_tok, tgt_tok, cfg
+
+
+def load_val_data(cfg, src_tok, tgt_tok, n=500):
+    """Load validation set as (src_tensors, ref_strings, input_strings)."""
+    from Data.data import OptimizedSanskritDataset
+    from torch.utils.data import Subset
+    from sklearn.model_selection import train_test_split
+
+    dataset = OptimizedSanskritDataset(
+        'train', max_len=cfg['model']['max_seq_len'],
+        cfg=cfg, src_tokenizer=src_tok, tgt_tokenizer=tgt_tok)
+    total = min(cfg['data']['dataset_size'], len(dataset))
+    _, val_idx = train_test_split(list(range(total)), train_size=0.8, random_state=42)
+    val_idx = val_idx[:n]
+
+    src_list, ref_list, inp_list = [], [], []
+    for i in val_idx:
+        item = dataset[i]
+        src_list.append(item['input_ids'].unsqueeze(0))
+        ref_list.append(item['target_text'])
+        inp_list.append(item['input_text'])
+    return src_list, ref_list, inp_list
+
+
+# ── Task 1 ────────────────────────────────────────────────────────────
+
+def run_task1(model, src_tok, device):
+    print("\n" + "="*65)
+    print("  TASK 1 — KV Cache Benchmark")
+    print("="*65)
+    if not hasattr(model.model, 'generate_cached'):
+        print("  SKIP: not D3PMCrossAttention.")
+        return
+    from analysis.kv_cache_benchmark import run_benchmark, print_summary
+    results = run_benchmark(model, src_tok, device, src_lens=[16, 32, 64])
+    print_summary(results)
+    path = os.path.join(OUTPUT_DIR, "task1_kv_cache.txt")
+    with open(path, "w") as f:
+        f.write("TASK 1 — KV CACHE BENCHMARK\n" + "="*40 + "\n\n")
+        f.write(f"{'src_len':>8}  {'standard(s)':>12}  {'cached(s)':>10}  "
+                f"{'speedup':>8}  {'encoder%':>9}\n")
+        for src_len, r in results.items():
+            f.write(f"{src_len:>8}  {r['standard_s']:>12.3f}  {r['cached_s']:>10.3f}  "
+                    f"{r['speedup']:>7.2f}x  {r['encoder_pct']:>8.1f}%\n")
+    print(f"  Saved: {path}")
+
+
+# ── Task 2 ────────────────────────────────────────────────────────────
+
+def run_task2(model, src_tok, tgt_tok, device, input_text):
+    print("\n" + "="*65)
+    print("  TASK 2 — Attention Visualization + Semantic Drift")
+    print("="*65)
+    print(f"  Input: {input_text}")
+    if not hasattr(model.model, 'encode_source'):
+        print("  SKIP: not D3PMCrossAttention.")
+        return
+
+    src_ids    = src_tok.encode(input_text)
+    src_tensor = torch.tensor([src_ids], dtype=torch.long, device=device)
+    src_chars  = list(input_text.strip())
+
+    from analysis.attention_viz import (AttentionCapture, plot_attn_heatmap,
+                                         plot_attn_evolution, plot_all_layers)
+    from analysis.semantic_drift import (capture_intermediate_outputs,
+                                          compute_drift, compute_token_stability,
+                                          plot_drift_curve)
+
+    # Attention capture
+    print("  Capturing attention weights...")
+    capturer     = AttentionCapture(model)
+    step_weights = capturer.capture(src_tensor, capture_every=10)
+
+    with torch.no_grad():
+        out_ids  = model.generate_cached(src_tensor)
+    tgt_ids   = [x for x in out_ids[0].tolist() if x > 4]
+    tgt_text  = tgt_tok.decode(tgt_ids).strip()
+    tgt_chars = list(tgt_text)
+    print(f"  Output: {tgt_text}")
+
+    first_t = max(step_weights.keys())
+    plot_attn_heatmap(step_weights, t_val=first_t, layer=0,
+        src_tokens=src_chars[:20], tgt_tokens=tgt_chars[:20],
+        save_path=os.path.join(OUTPUT_DIR, f"task2_attn_t{first_t}.png"),
+        title=f"Attention t={first_t} (noisy)  Layer 0")
+    plot_attn_heatmap(step_weights, t_val=0, layer=0,
+        src_tokens=src_chars[:20], tgt_tokens=tgt_chars[:20],
+        save_path=os.path.join(OUTPUT_DIR, "task2_attn_t0.png"),
+        title="Attention t=0 (final)  Layer 0")
+    plot_all_layers(step_weights, t_val=0,
+        src_tokens=src_chars[:20], tgt_tokens=tgt_chars[:20],
+        save_path=os.path.join(OUTPUT_DIR, "task2_all_layers_t0.png"))
+    if len(src_chars) > 0 and len(tgt_chars) > 0:
+        plot_attn_evolution(step_weights, src_token_idx=0, tgt_token_idx=0,
+            layer=0, src_token_str=src_chars[0], tgt_token_str=tgt_chars[0],
+            save_path=os.path.join(OUTPUT_DIR, "task2_attn_evolution.png"))
+
+    # Semantic drift
+    print("  Computing semantic drift...")
+    step_outputs, final_out = capture_intermediate_outputs(
+        model, src_tensor, tgt_tok, capture_every=5)
+    drift   = compute_drift(step_outputs, final_out)
+    stab    = compute_token_stability(step_outputs, final_out, tgt_tok)
+    plot_drift_curve(drift, src_text=input_text,
+        save_path=os.path.join(OUTPUT_DIR, "task2_semantic_drift.png"))
+
+    print(f"  Lock-in timestep: t={drift['lock_in_t']}")
+    print(f"  Mean position lock-in: t={stab['mean_lock_t']:.1f} ± {stab['std_lock_t']:.1f}")
+
+    report = os.path.join(OUTPUT_DIR, "task2_report.txt")
+    with open(report, "w", encoding="utf-8") as f:
+        f.write("TASK 2 — ATTENTION + DRIFT REPORT\n" + "="*50 + "\n\n")
+        f.write(f"Input  : {input_text}\nOutput : {final_out}\n\n")
+        f.write(f"Lock-in t : {drift['lock_in_t']}\n")
+        f.write(f"Mean pos lock-in : {stab['mean_lock_t']:.1f} ± {stab['std_lock_t']:.1f}\n\n")
+        f.write("Step → Output → CER-to-final\n" + "-"*60 + "\n")
+        for tv, cer in zip(drift["t_vals"], drift["cer_to_final"]):
+            f.write(f"  t={tv:4d}  |  {step_outputs.get(tv,'')[:40]:40s}  |  {cer:.4f}\n")
+    print(f"  Report: {report}")
+
+
+# ── Task 3 ────────────────────────────────────────────────────────────
+
+def run_task3(model, src_tok, tgt_tok, device, src_list, ref_list):
+    print("\n" + "="*65)
+    print("  TASK 3 — Concept Vectors + PCA Steering")
+    print("="*65)
+    if not hasattr(model.model, 'encode_source'):
+        print("  SKIP: not D3PMCrossAttention.")
+        return
+
+    from analysis.concept_vectors import (collect_hidden_states, fit_pca,
+        find_diversity_direction, generate_diversity_spectrum, plot_pca_space)
+
+    # Collect hidden states from val set
+    n = min(500, len(src_list))
+    print(f"  Collecting hidden states from {n} examples...")
+    hidden, _ = collect_hidden_states(
+        model, src_list[:n], t_capture=0, max_samples=n)
+
+    # Compute output lengths for diversity direction
+    lengths = []
+    for src in src_list[:n]:
+        with torch.no_grad():
+            out = model.generate_cached(src.to(device))
+        ids = [x for x in out[0].tolist() if x > 4]
+        lengths.append(len(tgt_tok.decode(ids)))
+
+    # Fit PCA + find diversity direction
+    pca = fit_pca(hidden, n_components=min(50, n-1))
+    direction, best_pc, corr = find_diversity_direction(hidden, lengths, pca)
+
+    # Plot concept space
+    plot_pca_space(hidden, lengths, pca, best_pc,
+        save_path=os.path.join(OUTPUT_DIR, "task3_concept_space.png"))
+
+    # Generate diversity spectrum for first example
+    print("\n  Diversity spectrum for first example:")
+    src0  = src_list[0]
+    inp0  = src_tok.decode([x for x in src0[0].tolist() if x > 4])
+    print(f"  Input: {inp0}")
+    spectrum = generate_diversity_spectrum(
+        model, src0.to(device), direction, tgt_tok,
+        alphas=[-2.0, -1.0, 0.0, 1.0, 2.0])
+
+    # Save diversity direction + results
+    np.save(os.path.join(OUTPUT_DIR, "task3_diversity_direction.npy"), direction)
+
+    report = os.path.join(OUTPUT_DIR, "task3_report.txt")
+    with open(report, "w", encoding="utf-8") as f:
+        f.write("TASK 3 — CONCEPT VECTORS + PCA STEERING\n" + "="*50 + "\n\n")
+        f.write(f"PCA: {pca.n_components_} components, "
+                f"{pca.explained_variance_ratio_.sum()*100:.1f}% variance\n")
+        f.write(f"Diversity PC: {best_pc}  (|r|={corr:.3f} with output length)\n\n")
+        f.write("Diversity spectrum:\n")
+        for alpha, text in sorted(spectrum.items()):
+            f.write(f"  alpha={alpha:+.1f}  →  {text}\n")
+    print(f"  Report: {report}")
+
+
+# ── Task 4 ────────────────────────────────────────────────────────────
+
+def run_task4(phase, model, src_tok, tgt_tok, device, cfg,
+              src_list, ref_list):
+    print("\n" + "="*65)
+    print(f"  TASK 4 — Step Ablation  (phase={phase})")
+    print("="*65)
+
+    from analysis.step_ablation import (generate_ablation_configs,
+        run_ablation_analysis, plot_ablation_3d, run_adversarial_test)
+
+    if phase == "generate_configs":
+        print("  Generating ablation configs...")
+        generate_ablation_configs(output_dir="ablation_configs")
+        print("\n  NEXT STEPS:")
+        print("  1. bash ablation_configs/train_all.sh")
+        print("  2. python analysis/run_analysis.py --task 4 --phase analyze")
+
+    elif phase == "analyze":
+        # Check which models exist
+        existing = [T for T in [4, 8, 16, 32, 64]
+                    if os.path.exists(f"ablation_results/T{T}/best_model.pt")]
+        if not existing:
+            print("  No ablation models found at ablation_results/T*/best_model.pt")
+            print("  Run: python analysis/run_analysis.py --task 4 --phase generate_configs")
+            print("  Then: bash ablation_configs/train_all.sh")
+            return
+
+        print(f"  Found models for T={existing}")
+        results = run_ablation_analysis(
+            ablation_dir="ablation_results", base_cfg=cfg,
+            src_list=src_list[:200], ref_list=ref_list[:200],
+            tgt_tokenizer=tgt_tok, device=device,
+            output_dir=OUTPUT_DIR)
+        plot_ablation_3d(results,
+            save_path=os.path.join(OUTPUT_DIR, "task4_ablation_3d.png"))
+
+    # Adversarial robustness always runs on existing model (no retraining)
+    print("\n  Running adversarial robustness test...")
+    inp_texts = [src_tok.decode([x for x in s[0].tolist() if x > 4])
+                 for s in src_list[:50]]
+    run_adversarial_test(
+        model, src_tok, tgt_tok,
+        test_inputs=inp_texts, test_refs=ref_list[:50],
+        device=device, output_dir=OUTPUT_DIR)
+
+
+# ── Task 5 ────────────────────────────────────────────────────────────
+
+def run_task5(model, src_tok, tgt_tok, device, cfg, src_list, ref_list):
+    print("\n" + "="*65)
+    print("  TASK 5 — Classifier-Free Guidance")
+    print("="*65)
+    if not hasattr(model.model, 'encode_source'):
+        print("  SKIP: not D3PMCrossAttention.")
+        return
+
+    from analysis.quality_classifier import (
+        QualityClassifier, collect_quality_data,
+        train_quality_classifier, sweep_guidance_scales)
+
+    clf_path = os.path.join(OUTPUT_DIR, "task5_quality_classifier.pt")
+    d_model  = cfg['model']['d_model']
+
+    # Step 1: collect or load training data
+    data_path = os.path.join(OUTPUT_DIR, "task5_quality_data.npz")
+    if os.path.exists(data_path):
+        print("  Loading cached quality data...")
+        data    = np.load(data_path)
+        hidden  = data["hidden"]
+        quality = data["quality"]
+    else:
+        print("  Collecting quality data (this takes a few minutes)...")
+        n       = min(2000, len(src_list))
+        hidden, quality = collect_quality_data(
+            model, src_list[:n], ref_list[:n], tgt_tok,
+            t_capture=0, max_samples=n)
+        np.savez(data_path, hidden=hidden, quality=quality)
+        print(f"  Saved quality data: {data_path}")
+
+    # Step 2: train or load classifier
+    if os.path.exists(clf_path):
+        print(f"  Loading cached classifier: {clf_path}")
+        clf = QualityClassifier(d_model)
+        clf.load_state_dict(torch.load(clf_path, map_location='cpu'))
+        clf.eval()
+    else:
+        print("  Training quality classifier...")
+        clf = train_quality_classifier(
+            hidden, quality, d_model=d_model,
+            epochs=30, batch_size=64, lr=1e-3,
+            save_path=clf_path)
+        clf.eval()
+
+    # Step 3: guidance scale sweep
+    print("\n  Guidance scale sweep (λ ∈ {0.0, 0.5, 1.0, 1.5, 2.0, 3.0})...")
+    n_sweep = min(50, len(src_list))
+    results = sweep_guidance_scales(
+        model, clf, src_list[:n_sweep], ref_list[:n_sweep],
+        tgt_tok, scales=[0.0, 0.5, 1.0, 1.5, 2.0, 3.0],
+        n_samples=n_sweep, device=device, output_dir=OUTPUT_DIR)
+
+    # Find optimal scale
+    best_scale = min(results, key=lambda s: results[s]["mean_cer"])
+    print(f"\n  Optimal guidance scale: λ={best_scale:.1f}  "
+          f"CER={results[best_scale]['mean_cer']:.4f}")
+
+    report = os.path.join(OUTPUT_DIR, "task5_report.txt")
+    with open(report, "w") as f:
+        f.write("TASK 5 — CLASSIFIER-FREE GUIDANCE\n" + "="*50 + "\n\n")
+        f.write(f"Classifier params: {sum(p.numel() for p in clf.parameters())}\n")
+        f.write(f"Training samples : {len(hidden)}\n\n")
+        f.write("Guidance scale sweep:\n")
+        f.write(f"  {'λ':>6}  {'CER':>8}  {'diversity':>10}\n")
+        f.write("  " + "-"*28 + "\n")
+        for s in sorted(results.keys()):
+            r = results[s]
+            marker = " ← optimal" if s == best_scale else ""
+            f.write(f"  {s:>6.1f}  {r['mean_cer']:>8.4f}  {r['diversity']:>10.3f}{marker}\n")
+    print(f"  Report: {report}")
+
+
+# ── Main ──────────────────────────────────────────────────────────────
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--task",
+        choices=["1","2","3","4","5","all"], default="all")
+    parser.add_argument("--input",
+        default="dharmo rakṣati rakṣitaḥ",
+        help="IAST input text for Task 2")
+    parser.add_argument("--phase",
+        choices=["generate_configs", "analyze"], default="analyze",
+        help="Task 4 phase: generate_configs (before training) or analyze (after)")
+    args = parser.parse_args()
+
+    cfg    = CONFIG
+    device = torch.device(cfg['training']['device'])
+
+    print("Loading model and tokenizers...")
+    model, src_tok, tgt_tok, cfg = load_everything(cfg, device)
+
+    # Load val data for tasks that need it (Tasks 3, 4, 5)
+    needs_data = args.task in ("3", "4", "5", "all")
+    if needs_data:
+        print("Loading validation data...")
+        src_list, ref_list, inp_list = load_val_data(cfg, src_tok, tgt_tok, n=500)
+    else:
+        src_list, ref_list, inp_list = [], [], []
+
+    tasks = (["1","2","3","4","5"] if args.task == "all"
+             else [args.task])
+
+    for task in tasks:
+        if task == "1":
+            run_task1(model, src_tok, device)
+        elif task == "2":
+            run_task2(model, src_tok, tgt_tok, device, args.input)
+        elif task == "3":
+            run_task3(model, src_tok, tgt_tok, device, src_list, ref_list)
+        elif task == "4":
+            run_task4(args.phase, model, src_tok, tgt_tok, device, cfg,
+                      src_list, ref_list)
+        elif task == "5":
+            run_task5(model, src_tok, tgt_tok, device, cfg, src_list, ref_list)
+
+    print(f"\n{'='*65}")
+    print(f"  All outputs saved to: {OUTPUT_DIR}/")
+    print("="*65)
+
+
+if __name__ == "__main__":
+    main()

sanskrit_model.py

@@ -0,0 +1,61 @@
+"""
+sanskrit_model.py  — Fixed
+===========================
+Added inference_mode parameter to forward() so reverse_process.py can
+pass inference_mode=True without a TypeError.
+
+The wrapper introspects each inner model's signature and only passes
+kwargs that model actually accepts — safe across all four architectures.
+"""
+
+import torch
+import torch.nn as nn
+import inspect
+
+
+class SanskritModel(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        model_type = cfg['model_type']
+
+        if model_type == 'd3pm_cross_attention':
+            from model.d3pm_model_cross_attention import D3PMCrossAttention
+            self.model = D3PMCrossAttention(cfg)
+
+        elif model_type == 'd3pm_encoder_decoder':
+            from model.d3pm_model_encoder_decoder import D3PMEncoderDecoder
+            self.model = D3PMEncoderDecoder(cfg)
+
+        elif model_type == 'baseline_cross_attention':
+            from model.d3pm_model_cross_attention import BaselineCrossAttention
+            self.model = BaselineCrossAttention(cfg)
+
+        elif model_type == 'baseline_encoder_decoder':
+            from model.d3pm_model_encoder_decoder import BaselineEncoderDecoder
+            self.model = BaselineEncoderDecoder(cfg)
+
+        else:
+            raise ValueError(f"Unknown model_type: {model_type}")
+
+    def forward(self, input_ids, target_ids, t, x0_hint=None, inference_mode=False):
+        """
+        Forward pass.  Introspects the inner model's signature so only
+        supported kwargs are passed — works with all four architectures.
+        """
+        sig    = inspect.signature(self.model.forward).parameters
+        kwargs = {}
+        if 'x0_hint'        in sig:
+            kwargs['x0_hint']        = x0_hint
+        if 'inference_mode' in sig:
+            kwargs['inference_mode'] = inference_mode
+
+        if 't' in sig:
+            return self.model(input_ids, target_ids, t, **kwargs)
+        else:
+            return self.model(input_ids, target_ids, **kwargs)
+
+    @torch.no_grad()
+    def generate(self, src, **kwargs):
+        sig      = inspect.signature(self.model.generate).parameters
+        filtered = {k: v for k, v in kwargs.items() if k in sig}
+        return self.model.generate(src, **filtered)

scheduler.py

@@ -0,0 +1,34 @@
+"""
+scheduler.py  — Fixed & Upgraded
+==================================
+Changes:
+  1. T=64 (was 16). More timesteps = richer denoising curriculum per epoch.
+  2. alpha at t=0 is EXACTLY 1.0 — fixes Bug 2 (final-step re-noise).
+  3. sample_timestep samples [0, T-1] including t=0, so model trains on
+     fully-clean inputs (learns the identity at t=0 explicitly).
+"""
+import torch, math
+
+class OptimizedCosineScheduler:
+    def __init__(self, cfg, device=None):
+        self.num_timesteps  = cfg['model']['diffusion_steps']   # 64
+        self.mask_token_id  = cfg['diffusion']['mask_token_id']
+        self.device         = device or torch.device('cpu')
+        self.alphas_cumprod = self._build_schedule().to(self.device)
+
+    def _build_schedule(self):
+        T   = self.num_timesteps
+        t   = torch.arange(T + 1, dtype=torch.float32)
+        f_t = torch.cos((t / T + 0.008) / 1.008 * math.pi / 2) ** 2
+        alphas_bar = f_t / f_t[0]
+        alphas_bar = alphas_bar[1:]       # shape [T]
+        alphas_bar[0]  = 1.0              # FIX: exact 1.0 at t=0
+        alphas_bar[-1] = alphas_bar[-1].clamp(max=0.001)
+        return alphas_bar
+
+    def sample_timestep(self, batch_size):
+        """Uniform [0, T-1] — includes t=0 so model sees clean inputs."""
+        return torch.randint(0, self.num_timesteps, (batch_size,))
+
+    def get_alpha(self, t):
+        return self.alphas_cumprod[t.to(self.alphas_cumprod.device).long()]

semantic_drift.py

@@ -0,0 +1,279 @@
+"""
+analysis/semantic_drift.py
+===========================
+Task 2: Semantic drift metric — how much does the intermediate generation
+diverge from the final output as we walk through diffusion steps T → 0?
+
+Metric: CER between x0_estimate at each step vs the final x0 at t=0.
+
+A well-trained model should show:
+  - High drift at t=T-1 (near-random initial estimate)
+  - Rapid decrease in drift around t=T//2 (model finds the right structure)
+  - Near-zero drift at t=10 (output is stable, only fine corrections remain)
+
+If drift stays high until t=5 then suddenly collapses → model is doing all
+its work in the last few steps → consider reducing T.
+
+Also measures:
+  - Token stability: fraction of positions that don't change between steps
+  - Lock-in time: first step where each position "commits" to its final token
+
+No retraining required. Uses generate_cached() with intermediate snapshots.
+"""
+
+import torch
+import torch.nn.functional as F
+import numpy as np
+from typing import List, Dict, Optional, Tuple
+
+
+def compute_cer_between(pred: str, ref: str) -> float:
+    """CER between two strings."""
+    if not ref:
+        return 1.0 if pred else 0.0
+
+    def edit_distance(s1, s2):
+        m, n = len(s1), len(s2)
+        dp = list(range(n + 1))
+        for i in range(1, m + 1):
+            prev, dp[0] = dp[0], i
+            for j in range(1, n + 1):
+                temp = dp[j]
+                dp[j] = prev if s1[i-1] == s2[j-1] else 1 + min(prev, dp[j], dp[j-1])
+                prev = temp
+        return dp[n]
+
+    return edit_distance(pred, ref) / len(ref)
+
+
+@torch.no_grad()
+def capture_intermediate_outputs(
+    model,
+    src:          torch.Tensor,
+    tgt_tokenizer,
+    capture_every: int = 5,
+    temperature:   float = 0.8,
+    top_k:         int   = 40,
+) -> Tuple[Dict[int, str], str]:
+    """
+    Run generation while recording the decoded x0_estimate at every
+    `capture_every` diffusion steps.
+
+    Args:
+        model         : SanskritModel (D3PMCrossAttention)
+        src           : [1, src_len] IAST token ids (single sample)
+        tgt_tokenizer : SanskritTargetTokenizer for decoding intermediate outputs
+        capture_every : record every N steps
+        temperature   : sampling temperature
+        top_k         : top-k filter
+
+    Returns:
+        step_outputs : dict mapping t_val → decoded Devanagari string at that step
+        final_output : decoded string at t=0 (final result)
+    """
+    if src.dim() == 1:
+        src = src.unsqueeze(0)
+
+    inner  = model.model
+    T      = inner.scheduler.num_timesteps
+    device = src.device
+
+    # Encode source once (KV cache)
+    memory, src_pad_mask = inner.encode_source(src)
+
+    B       = src.shape[0]
+    tgt_len = inner.max_seq_len
+    mask_id = inner.mask_token_id
+
+    x0_est = torch.full((B, tgt_len), mask_id, dtype=torch.long, device=device)
+    hint   = None
+
+    step_outputs: Dict[int, str] = {}
+    inner.eval()
+
+    for t_val in range(T - 1, -1, -1):
+        t       = torch.full((B,), t_val, dtype=torch.long, device=device)
+        is_last = (t_val == 0)
+
+        logits, _ = inner.forward_cached(
+            memory, src_pad_mask, x0_est, t,
+            x0_hint=hint, inference_mode=True,
+        )
+
+        logits = logits / max(temperature, 1e-8)
+        if top_k > 0:
+            V = logits.shape[-1]
+            if top_k < V:
+                topk_vals, _ = torch.topk(logits, top_k, dim=-1)
+                threshold    = topk_vals[..., -1].unsqueeze(-1)
+                logits       = logits.masked_fill(logits < threshold, float('-inf'))
+
+        probs  = F.softmax(logits, dim=-1)
+        x0_est = torch.argmax(probs, dim=-1) if is_last else _sample(probs)
+        hint   = x0_est
+
+        # Capture at this step
+        if (T - 1 - t_val) % capture_every == 0 or is_last:
+            ids  = [x for x in x0_est[0].tolist() if x > 4]
+            text = tgt_tokenizer.decode(ids).strip()
+            step_outputs[t_val] = text
+
+    final_output = step_outputs.get(0, "")
+    return step_outputs, final_output
+
+
+def _sample(probs):
+    B, L, V = probs.shape
+    flat    = probs.view(B * L, V).clamp(min=1e-9)
+    flat    = flat / flat.sum(dim=-1, keepdim=True)
+    return torch.multinomial(flat, 1).squeeze(-1).view(B, L)
+
+
+def compute_drift(
+    step_outputs:  Dict[int, str],
+    final_output:  str,
+) -> Dict[str, object]:
+    """
+    Compute drift metrics comparing each intermediate output to the final.
+
+    Returns dict with:
+      t_vals      : list of captured timesteps (T-1 → 0)
+      cer_to_final: CER between each step's output and the final output
+                    0.0 = identical to final, 1.0 = completely different
+      lock_in_t   : first t_val where CER drops and stays below 0.1
+                    (step at which output "commits" to final form)
+    """
+    t_vals       = sorted(step_outputs.keys(), reverse=True)   # T-1 → 0
+    cer_to_final = []
+
+    for t_val in t_vals:
+        cer = compute_cer_between(step_outputs[t_val], final_output)
+        cer_to_final.append(cer)
+
+    # Find lock-in: first step where CER stays below threshold for rest of run
+    threshold = 0.1
+    lock_in_t = 0   # default: never locked in early
+    for i, (t_val, cer) in enumerate(zip(t_vals, cer_to_final)):
+        if all(c <= threshold for c in cer_to_final[i:]):
+            lock_in_t = t_val
+            break
+
+    return {
+        "t_vals":       t_vals,
+        "cer_to_final": cer_to_final,
+        "lock_in_t":    lock_in_t,
+        "final_output": final_output,
+    }
+
+
+def compute_token_stability(
+    step_outputs:  Dict[int, str],
+    final_output:  str,
+    tgt_tokenizer,
+) -> Dict[str, object]:
+    """
+    Token-level stability: for each position, at which diffusion step
+    does it first match its final token and stay matched?
+
+    Returns:
+      position_lock_times: list of t_val at which each position locks in
+      mean_lock_t        : average lock-in timestep across positions
+    """
+    T      = max(step_outputs.keys())
+    t_vals = sorted(step_outputs.keys(), reverse=True)   # T-1 → 0
+
+    # Encode all intermediate outputs and the final
+    def encode(text):
+        return tgt_tokenizer.encode(text)
+
+    final_ids = encode(final_output)
+    L         = len(final_ids)
+
+    # Build matrix: [n_steps, L]
+    step_ids = []
+    for t_val in t_vals:
+        step_ids.append(encode(step_outputs.get(t_val, "")))
+
+    # Pad all to same length
+    max_len = max(len(s) for s in step_ids)
+    step_ids = [s + [1] * (max_len - len(s)) for s in step_ids]   # 1=PAD
+    final_ids_padded = final_ids + [1] * (max_len - len(final_ids))
+
+    step_arr  = np.array(step_ids)                # [n_steps, L]
+    final_arr = np.array(final_ids_padded)         # [L]
+
+    # For each position: find first step index where it matches final
+    # and stays matched for all subsequent steps
+    position_lock_steps = []
+    for pos in range(min(L, max_len)):
+        col = step_arr[:, pos]   # [n_steps]
+        fin = final_arr[pos]
+        locked_at = len(t_vals) - 1   # default: never locks early
+        for i in range(len(t_vals)):
+            if all(col[i:] == fin):
+                locked_at = i
+                break
+        position_lock_steps.append(t_vals[locked_at] if locked_at < len(t_vals) else 0)
+
+    return {
+        "position_lock_times": position_lock_steps,
+        "mean_lock_t":         float(np.mean(position_lock_steps)),
+        "std_lock_t":          float(np.std(position_lock_steps)),
+    }
+
+
+def plot_drift_curve(
+    drift_result: Dict,
+    src_text:     str = "",
+    save_path:    Optional[str] = None,
+):
+    """
+    Plot CER-to-final vs diffusion step.
+    Shows where the model "commits" to the final output.
+    """
+    try:
+        import matplotlib.pyplot as plt
+    except ImportError:
+        print("pip install matplotlib.")
+        return
+
+    t_vals  = drift_result["t_vals"]
+    cers    = drift_result["cer_to_final"]
+    lock_t  = drift_result["lock_in_t"]
+
+    fig, ax = plt.subplots(figsize=(12, 4))
+    ax.plot(range(len(t_vals)), cers, linewidth=1.8, color='coral', label='CER to final')
+    ax.fill_between(range(len(t_vals)), cers, alpha=0.15, color='coral')
+
+    # Mark lock-in point
+    if lock_t in t_vals:
+        lock_idx = t_vals.index(lock_t)
+        ax.axvline(lock_idx, color='steelblue', linestyle='--', linewidth=1.2,
+                   label=f"Lock-in at t={lock_t}")
+
+    ax.axhline(0.1, color='gray', linestyle=':', linewidth=1, alpha=0.7)
+
+    n = len(t_vals)
+    tick_positions = list(range(0, n, max(1, n // 10)))
+    ax.set_xticks(tick_positions)
+    ax.set_xticklabels([str(t_vals[i]) for i in tick_positions], fontsize=8)
+    ax.set_xlabel("Diffusion step t  (T-1 → 0)", fontsize=11)
+    ax.set_ylabel("CER vs final output", fontsize=11)
+    ax.set_ylim(0, 1.05)
+    ax.set_xlim(0, n - 1)
+    ax.legend(fontsize=10)
+
+    title = f"Semantic drift"
+    if src_text:
+        title += f"  |  src: {src_text[:50]}"
+    ax.set_title(title, fontsize=11)
+    plt.tight_layout()
+
+    if save_path:
+        import os
+        os.makedirs(os.path.dirname(save_path) or ".", exist_ok=True)
+        plt.savefig(save_path, dpi=150, bbox_inches='tight')
+        print(f"Saved: {save_path}")
+    else:
+        plt.show()
+    plt.close()

step_ablation.py

@@ -0,0 +1,389 @@
+"""
+analysis/step_ablation.py
+==========================
+Task 4: Semantic Robustness — Ablation of Diffusion Steps vs Meaning Preservation
+
+Two-phase workflow (retraining IS required for different T values):
+
+  PHASE 1 — Generate configs + train (run once per T value):
+    python analysis/step_ablation.py --phase generate_configs
+    # Creates configs: ablation_configs/T4.py, T8.py, T16.py, T32.py, T64.py
+    # Then train each: MODEL_TYPE=d3pm_cross_attention python train.py  (for each config)
+
+  PHASE 2 — Analyze trained models (no retraining needed):
+    python analysis/step_ablation.py --phase analyze
+    # Loads each trained model, generates 200 paraphrases, computes CER
+    # Produces 3D plot: X=steps, Y=generation_speed, Z=CER
+
+Why retraining is needed:
+  A model trained with T=128 learns to denoise from x_t~Uniform[0,128].
+  Running it with T=4 means the model only sees t∈{0,1,2,3} — which it
+  was never trained on at those scales. Outputs are meaningless.
+  You must train a separate model for each T value.
+
+Also implements adversarial robustness test (no retraining):
+  Takes your existing T=128 model and tests whether corrupted IAST
+  inputs (typos, character swaps) cause proportional output degradation.
+"""
+
+import torch
+import torch.nn.functional as F
+import numpy as np
+import os
+import sys
+import time
+import json
+import copy
+from typing import List, Dict, Optional
+
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+
+# ── Phase 1: Config generation ────────────────────────────────────────
+
+T_VALUES = [4, 8, 16, 32, 64]
+
+def generate_ablation_configs(base_config_path: str = "config.py",
+                               output_dir: str = "ablation_configs"):
+    """
+    Generate one config file per T value.
+    Each config is a copy of the base config with diffusion_steps changed.
+
+    After running this, train each model:
+        for T in 4 8 16 32 64; do
+            cp ablation_configs/config_T${T}.py config.py
+            python train.py
+            mv results7/d3pm_cross_attention_neg_False \
+               ablation_results/T${T}
+        done
+    """
+    os.makedirs(output_dir, exist_ok=True)
+
+    # Read base config
+    with open(base_config_path, "r") as f:
+        base_src = f.read()
+
+    for T in T_VALUES:
+        # Replace diffusion_steps and num_steps
+        cfg_src = base_src
+        cfg_src = cfg_src.replace(
+            '"diffusion_steps": 128',
+            f'"diffusion_steps": {T}'
+        )
+        cfg_src = cfg_src.replace(
+            "'diffusion_steps': 128",
+            f"'diffusion_steps': {T}"
+        )
+        cfg_src = cfg_src.replace(
+            '"num_steps": 128',
+            f'"num_steps": {T}'
+        )
+        cfg_src = cfg_src.replace(
+            "'num_steps': 128",
+            f"'num_steps': {T}"
+        )
+        out_path = os.path.join(output_dir, f"config_T{T}.py")
+        with open(out_path, "w") as f:
+            f.write(f"# Ablation config: T={T} diffusion steps\n")
+            f.write(cfg_src)
+        print(f"  Wrote: {out_path}")
+
+    # Write a shell script to train all
+    shell_script = os.path.join(output_dir, "train_all.sh")
+    with open(shell_script, "w") as f:
+        f.write("#!/bin/bash\n")
+        f.write("# Run this script to train all ablation models\n\n")
+        for T in T_VALUES:
+            f.write(f"echo '=== Training T={T} ==='\n")
+            f.write(f"cp {output_dir}/config_T{T}.py config.py\n")
+            f.write(f"python train.py\n")
+            f.write(f"mkdir -p ablation_results/T{T}\n")
+            f.write(f"cp -r results7/d3pm_cross_attention_neg_False/best_model.pt "
+                    f"ablation_results/T{T}/best_model.pt\n")
+            f.write(f"cp -r results7/d3pm_cross_attention_neg_False/train.log "
+                    f"ablation_results/T{T}/train.log\n\n")
+    os.chmod(shell_script, 0o755)
+    print(f"\nTraining script: {shell_script}")
+    print(f"Run: bash {shell_script}")
+
+
+# ── Phase 2: Analysis (after models are trained) ──────────────────────
+
+def compute_cer(pred: str, ref: str) -> float:
+    if not ref:
+        return 1.0
+
+    def edit_distance(s1, s2):
+        m, n = len(s1), len(s2)
+        dp = list(range(n + 1))
+        for i in range(1, m + 1):
+            prev, dp[0] = dp[0], i
+            for j in range(1, n + 1):
+                temp = dp[j]
+                dp[j] = prev if s1[i-1] == s2[j-1] else 1 + min(prev, dp[j], dp[j-1])
+                prev = temp
+        return dp[n]
+
+    return edit_distance(pred, ref) / max(len(ref), 1)
+
+
+def evaluate_model(
+    model,
+    src_list:      List[torch.Tensor],
+    ref_list:      List[str],
+    tgt_tokenizer,
+    n_samples:     int   = 200,
+    temperature:   float = 0.8,
+    top_k:         int   = 40,
+) -> Dict:
+    """
+    Generate n_samples outputs and compute CER + generation speed.
+
+    Returns dict with:
+        mean_cer      : average CER over samples
+        generation_s  : total wall-clock seconds for all generations
+        speed_per_sample: seconds per sample
+        cer_list      : per-sample CER values
+    """
+    device   = next(model.parameters()).device
+    n        = min(n_samples, len(src_list))
+    cer_list = []
+
+    start = time.perf_counter()
+    for i, (src, ref) in enumerate(zip(src_list[:n], ref_list[:n])):
+        if src.dim() == 1:
+            src = src.unsqueeze(0)
+
+        with torch.no_grad():
+            if hasattr(model.model, 'generate_cached'):
+                out = model.model.generate_cached(
+                    src.to(device), temperature=temperature, top_k=top_k
+                )
+            else:
+                out = model.generate(
+                    src.to(device), temperature=temperature, top_k=top_k
+                )
+
+        ids  = [x for x in out[0].tolist() if x > 4]
+        pred = tgt_tokenizer.decode(ids).strip()
+        cer  = compute_cer(pred, ref)
+        cer_list.append(cer)
+
+    elapsed = time.perf_counter() - start
+
+    return {
+        "mean_cer":          float(np.mean(cer_list)),
+        "std_cer":           float(np.std(cer_list)),
+        "generation_s":      elapsed,
+        "speed_per_sample":  elapsed / max(n, 1),
+        "cer_list":          cer_list,
+        "n_samples":         n,
+    }
+
+
+def run_ablation_analysis(
+    ablation_dir:  str = "ablation_results",
+    base_cfg:      dict = None,
+    src_list:      List[torch.Tensor] = None,
+    ref_list:      List[str] = None,
+    tgt_tokenizer  = None,
+    device:        torch.device = None,
+    output_dir:    str = "analysis/outputs",
+) -> Dict:
+    """
+    Load each trained model and evaluate.
+    Produces results dict and 3D plot.
+
+    Expects ablation_results/T{N}/best_model.pt for each T in T_VALUES.
+    """
+    from inference import load_model
+
+    results = {}
+    for T in T_VALUES:
+        ckpt = os.path.join(ablation_dir, f"T{T}", "best_model.pt")
+        if not os.path.exists(ckpt):
+            print(f"  SKIP T={T}: no checkpoint at {ckpt}")
+            continue
+
+        print(f"\nEvaluating T={T}...")
+        cfg_T = copy.deepcopy(base_cfg)
+        cfg_T['model']['diffusion_steps'] = T
+        cfg_T['inference']['num_steps']   = T
+
+        model, cfg_T = load_model(ckpt, cfg_T, device)
+        model.eval()
+
+        metrics = evaluate_model(
+            model, src_list, ref_list, tgt_tokenizer, n_samples=200
+        )
+        results[T] = metrics
+        print(f"  T={T}  CER={metrics['mean_cer']:.4f}  "
+              f"speed={metrics['speed_per_sample']:.3f}s/sample")
+
+        del model
+
+    # Save results
+    os.makedirs(output_dir, exist_ok=True)
+    results_path = os.path.join(output_dir, "ablation_results.json")
+    with open(results_path, "w") as f:
+        json.dump({str(k): {kk: vv for kk, vv in v.items() if kk != 'cer_list'}
+                   for k, v in results.items()}, f, indent=2)
+    print(f"\nResults saved: {results_path}")
+
+    return results
+
+
+def plot_ablation_3d(
+    results:   Dict,
+    save_path: Optional[str] = None,
+):
+    """
+    3D plot: X=diffusion_steps, Y=generation_speed(s/sample), Z=CER.
+    Also produces a 2D summary plot.
+    """
+    try:
+        import matplotlib.pyplot as plt
+        from mpl_toolkits.mplot3d import Axes3D
+    except ImportError:
+        print("pip install matplotlib.")
+        return
+
+    T_list    = sorted(results.keys())
+    cers      = [results[T]["mean_cer"] for T in T_list]
+    speeds    = [results[T]["speed_per_sample"] for T in T_list]
+
+    # ── 3D plot ───────────────────────────────────────────────────────
+    fig = plt.figure(figsize=(14, 5))
+
+    ax3d = fig.add_subplot(121, projection='3d')
+    ax3d.scatter(T_list, speeds, cers, c=cers, cmap='RdYlGn_r', s=80)
+    for T, s, c in zip(T_list, speeds, cers):
+        ax3d.text(T, s, c, f"T={T}", fontsize=8)
+    ax3d.set_xlabel("Diffusion steps T", fontsize=9)
+    ax3d.set_ylabel("Speed (s/sample)", fontsize=9)
+    ax3d.set_zlabel("CER (↓ better)", fontsize=9)
+    ax3d.set_title("T vs speed vs CER", fontsize=10)
+
+    # ── 2D CER vs T (find the knee) ──────────────────────────────────
+    ax2d = fig.add_subplot(122)
+    ax2d.plot(T_list, cers, 'o-', linewidth=1.8, color='coral', markersize=7)
+    for T, c in zip(T_list, cers):
+        ax2d.annotate(f"{c:.3f}", (T, c), textcoords="offset points",
+                      xytext=(0, 8), fontsize=8, ha='center')
+
+    # Find knee: largest CER drop per unit T (elbow method)
+    if len(T_list) >= 3:
+        drops  = [cers[i] - cers[i+1] for i in range(len(cers)-1)]
+        knee_i = int(np.argmax(drops))
+        knee_T = T_list[knee_i + 1]
+        ax2d.axvline(knee_T, color='steelblue', linestyle='--', linewidth=1.2,
+                     label=f"Knee at T={knee_T}")
+        ax2d.legend(fontsize=9)
+
+    ax2d.set_xlabel("Diffusion steps T", fontsize=10)
+    ax2d.set_ylabel("CER (lower = better)", fontsize=10)
+    ax2d.set_title("CER vs diffusion steps", fontsize=10)
+    ax2d.set_ylim(0, max(cers) * 1.1)
+
+    plt.tight_layout()
+    if save_path:
+        os.makedirs(os.path.dirname(save_path) or ".", exist_ok=True)
+        plt.savefig(save_path, dpi=150, bbox_inches='tight')
+        print(f"Saved: {save_path}")
+    else:
+        plt.show()
+    plt.close()
+
+
+# ── Adversarial robustness test (no retraining needed) ───────────────
+
+def corrupt_iast(text: str, corruption_rate: float = 0.05) -> str:
+    """
+    Introduce random corruption into IAST text:
+      - Character swap (adjacent chars swapped)
+      - Character deletion
+      - Random character insertion
+
+    Models rate as 5% to 20% corruption to test robustness.
+    """
+    import random
+    chars = list(text)
+    n_corrupt = max(1, int(len(chars) * corruption_rate))
+
+    for _ in range(n_corrupt):
+        op  = random.choice(['swap', 'delete', 'insert'])
+        pos = random.randint(0, len(chars) - 1)
+
+        if op == 'swap' and pos < len(chars) - 1:
+            chars[pos], chars[pos+1] = chars[pos+1], chars[pos]
+        elif op == 'delete' and len(chars) > 1:
+            chars.pop(pos)
+        elif op == 'insert':
+            chars.insert(pos, random.choice('abcdeimnostu'))
+
+    return "".join(chars)
+
+
+@torch.no_grad()
+def run_adversarial_test(
+    model,
+    src_tokenizer,
+    tgt_tokenizer,
+    test_inputs:    List[str],
+    test_refs:      List[str],
+    corruption_rates: List[float] = [0.0, 0.05, 0.10, 0.15, 0.20],
+    device:         torch.device  = None,
+    output_dir:     str           = "analysis/outputs",
+) -> Dict:
+    """
+    Test if CER degrades proportionally with IAST corruption.
+    Uses existing trained model — no retraining.
+    """
+    device = device or next(model.parameters()).device
+    results = {}
+
+    print("\nAdversarial robustness test...")
+    for rate in corruption_rates:
+        cer_list = []
+        for text, ref in zip(test_inputs, test_refs):
+            corrupted = corrupt_iast(text, rate)
+            ids       = src_tokenizer.encode(corrupted)
+            src       = torch.tensor([ids], dtype=torch.long, device=device)
+
+            if hasattr(model.model, 'generate_cached'):
+                out = model.model.generate_cached(src)
+            else:
+                out = model.generate(src)
+
+            pred_ids = [x for x in out[0].tolist() if x > 4]
+            pred     = tgt_tokenizer.decode(pred_ids).strip()
+            cer_list.append(compute_cer(pred, ref))
+
+        mean_cer = float(np.mean(cer_list))
+        results[rate] = mean_cer
+        print(f"  corruption={rate*100:.0f}%  →  CER={mean_cer:.4f}")
+
+    # Save + plot
+    os.makedirs(output_dir, exist_ok=True)
+    try:
+        import matplotlib.pyplot as plt
+        fig, ax = plt.subplots(figsize=(8, 4))
+        rates   = [r * 100 for r in corruption_rates]
+        cers    = [results[r] for r in corruption_rates]
+        ax.plot(rates, cers, 'o-', linewidth=1.8, color='steelblue', markersize=7)
+        ax.set_xlabel("IAST corruption rate (%)", fontsize=11)
+        ax.set_ylabel("CER", fontsize=11)
+        ax.set_title("Model robustness to IAST input corruption", fontsize=11)
+        ax.set_ylim(0, max(cers) * 1.2)
+        plt.tight_layout()
+        plt.savefig(os.path.join(output_dir, "adversarial_robustness.png"),
+                    dpi=150, bbox_inches='tight')
+        plt.close()
+        print(f"  Saved: {output_dir}/adversarial_robustness.png")
+    except ImportError:
+        pass
+
+    with open(os.path.join(output_dir, "adversarial_results.json"), "w") as f:
+        json.dump({str(k): v for k, v in results.items()}, f, indent=2)
+
+    return results

tokenizer.py

@@ -0,0 +1,222 @@
+"""
+tokenizer.py  — Dual Tokenizer Fix
+====================================
+Two separate BPE tokenizers:
+
+  SanskritSourceTokenizer  — trained on quote_text (Roman/IAST script)
+  SanskritTargetTokenizer  — trained on quote_devanagari (Devanagari script)
+
+WHY SEPARATE?
+  Roman Sanskrit and Devanagari are fundamentally different character sets.
+  Roman uses a-z + diacritics (~60 unique chars), Devanagari uses ā-ह + matras
+  (~100+ unique chars). A shared BPE tokenizer wastes half its vocab on
+  character combos that never cross scripts, and forces the embedding table
+  to encode both scripts in one space — confusing the model's cross-attention.
+
+  With separate tokenizers:
+  - src vocab captures Roman subwords cleanly (ā, ś, ṭ, ṃ etc.)
+  - tgt vocab captures Devanagari akshara clusters cleanly (क्ष, त्र, etc.)
+  - The model learns a true cross-script mapping in its cross-attention
+
+SPECIAL TOKENS (same IDs in both):
+  [MASK] = 0   ← required by absorbing diffusion
+  [PAD]  = 1
+  [UNK]  = 2
+  [CLS]  = 3
+  [SEP]  = 4
+"""
+
+from tokenizers import Tokenizer
+from tokenizers.models import BPE
+from tokenizers.trainers import BpeTrainer
+from tokenizers.pre_tokenizers import Whitespace
+from datasets import load_dataset
+from pathlib import Path
+
+
+SPECIAL_TOKENS = ["[MASK]", "[PAD]", "[UNK]", "[CLS]", "[SEP]"]
+
+
+def _build_bpe(texts, vocab_size):
+    """Build a BPE tokenizer from an iterator of strings."""
+    tok = Tokenizer(BPE(unk_token="[UNK]"))
+    tok.pre_tokenizer = Whitespace()
+    trainer = BpeTrainer(
+        vocab_size=vocab_size,
+        special_tokens=SPECIAL_TOKENS,   # [MASK] MUST be first → id=0
+        min_frequency=2,
+    )
+    tok.train_from_iterator(texts, trainer)
+    return tok
+
+
+def _validate(tok, name):
+    mask_id = tok.token_to_id("[MASK]")
+    pad_id  = tok.token_to_id("[PAD]")
+    assert mask_id == 0, f"{name}: [MASK] must be id=0, got {mask_id}"
+    assert pad_id  == 1, f"{name}: [PAD] must be id=1, got {pad_id}"
+    print(f"✅ {name}: [MASK]=0, [PAD]=1 confirmed. Vocab size={tok.get_vocab_size()}")
+
+
+# ── Source tokenizer (Roman/IAST Sanskrit) ────────────────────────────
+
+class SanskritSourceTokenizer:
+    """
+    Tokenizer for quote_text — Roman transliteration of Sanskrit.
+    Examples: "dharmo rakṣati rakṣitaḥ", "yatra nāryastu pūjyante"
+    """
+    MODEL_PATH = "sanskrit_src_tokenizer.json"
+
+    def __init__(self, vocab_size=8000, max_len=80, n_train_samples=50000):
+        self.vocab_size = vocab_size
+        self.max_len    = max_len
+        self.mask_token_id = 0
+
+        if Path(self.MODEL_PATH).exists():
+            print(f"📖 Loading source tokenizer from {self.MODEL_PATH} …")
+            self.tokenizer = Tokenizer.from_file(self.MODEL_PATH)
+        else:
+            print("🎓 Training source tokenizer on quote_text …")
+            self._train(vocab_size, n_train_samples)
+
+        _validate(self.tokenizer, "SrcTokenizer")
+
+    def _train(self, vocab_size, n_samples):
+        dataset = load_dataset("paws/sanskrit-verses-gretil", split="train")
+        n = min(n_samples, len(dataset))
+        texts = [s["quote_text"] for s in dataset.select(range(n))
+                 if s["quote_text"].strip()]
+        self.tokenizer = _build_bpe(texts, vocab_size)
+        self.tokenizer.save(self.MODEL_PATH)
+        print(f"✅ Source tokenizer trained on {len(texts)} Roman texts.")
+
+    def encode(self, text):
+        ids = self.tokenizer.encode(text).ids[:self.max_len]
+        pad = self.tokenizer.token_to_id("[PAD]")
+        ids += [pad] * max(0, self.max_len - len(ids))
+        return ids[:self.max_len]
+
+    def decode(self, ids):
+        clean = [i for i in ids if i > 4]   # skip special tokens
+        return self.tokenizer.decode(clean)
+
+    def __len__(self):
+        return self.vocab_size
+
+
+# ── Target tokenizer (Devanagari Sanskrit) ───────────────────────────
+
+class SanskritTargetTokenizer:
+    """
+    Tokenizer for quote_devanagari — Devanagari script.
+    Examples: "धर्मो रक्षति रक्षितः", "यत्र नार्यस्तु पूज्यन्ते"
+    """
+    MODEL_PATH = "sanskrit_tgt_tokenizer.json"
+
+    def __init__(self, vocab_size=8000, max_len=80, n_train_samples=50000):
+        self.vocab_size = vocab_size
+        self.max_len    = max_len
+        self.mask_token_id = 0
+
+        if Path(self.MODEL_PATH).exists():
+            print(f"📖 Loading target tokenizer from {self.MODEL_PATH} …")
+            self.tokenizer = Tokenizer.from_file(self.MODEL_PATH)
+        else:
+            print("🎓 Training target tokenizer on quote_devanagari …")
+            self._train(vocab_size, n_train_samples)
+
+        _validate(self.tokenizer, "TgtTokenizer")
+
+    def _train(self, vocab_size, n_samples):
+        dataset = load_dataset("paws/sanskrit-verses-gretil", split="train")
+        n = min(n_samples, len(dataset))
+        texts = [s["quote_devanagari"] for s in dataset.select(range(n))
+                 if s["quote_devanagari"].strip()]
+        self.tokenizer = _build_bpe(texts, vocab_size)
+        self.tokenizer.save(self.MODEL_PATH)
+        print(f"✅ Target tokenizer trained on {len(texts)} Devanagari texts.")
+
+    def encode(self, text):
+        ids = self.tokenizer.encode(text).ids[:self.max_len]
+        pad = self.tokenizer.token_to_id("[PAD]")
+        ids += [pad] * max(0, self.max_len - len(ids))
+        return ids[:self.max_len]
+
+    def decode(self, ids):
+        clean = [i for i in ids if i > 4]
+        return self.tokenizer.decode(clean)
+
+    # Methods required by BERTScore
+    def build_inputs_with_special_tokens(self, token_ids):
+        return list(token_ids)
+
+    def get_vocab(self):
+        return {str(i): i for i in range(self.vocab_size)}
+
+    def convert_ids_to_tokens(self, ids):
+        return [str(i) for i in ids]
+
+    def __len__(self):
+        return self.vocab_size
+
+
+# ── Legacy shared tokenizer (kept for backward compat) ───────────────
+
+class SanskritTokenizer:
+    """
+    LEGACY: single shared tokenizer trained on BOTH scripts.
+    Still works but suboptimal — use SanskritSourceTokenizer +
+    SanskritTargetTokenizer for the quote_text → quote_devanagari task.
+    """
+    MODEL_PATH = "sanskrit_tokenizer_m4pro.json"
+
+    def __init__(self, vocab_size=16000, max_len=80):
+        self.vocab_size    = vocab_size
+        self.max_len       = max_len
+        self.mask_token_id = 0
+
+        if Path(self.MODEL_PATH).exists():
+            print("📖 Loading shared tokenizer …")
+            self.tokenizer = Tokenizer.from_file(self.MODEL_PATH)
+        else:
+            print("🎓 Training shared tokenizer on both scripts …")
+            self._train(vocab_size)
+
+        _validate(self.tokenizer, "SharedTokenizer")
+
+    def _train(self, vocab_size):
+        dataset = load_dataset("paws/sanskrit-verses-gretil", split="train")
+        n = min(50000, len(dataset))
+        texts = []
+        for s in dataset.select(range(n)):
+            if s["quote_text"].strip():
+                texts.append(s["quote_text"])
+            if s["quote_devanagari"].strip():
+                texts.append(s["quote_devanagari"])
+        self.tokenizer = _build_bpe(texts, vocab_size)
+        self.tokenizer.save(self.MODEL_PATH)
+        print(f"✅ Shared tokenizer trained ({len(texts)} texts).")
+
+    def encode(self, text):
+        ids = self.tokenizer.encode(text).ids[:self.max_len]
+        pad = self.tokenizer.token_to_id("[PAD]")
+        ids += [pad] * max(0, self.max_len - len(ids))
+        return ids[:self.max_len]
+
+    def decode(self, ids):
+        if ids and isinstance(ids[0], list):
+            raise TypeError("decode() got 2D list — pass a 1D list.")
+        clean = [i for i in ids if i > 4]
+        return self.tokenizer.decode(clean)
+
+    def build_inputs_with_special_tokens(self, token_ids):
+        return list(token_ids)
+
+    def get_vocab(self):
+        return {str(i): i for i in range(self.vocab_size)}
+
+    def convert_ids_to_tokens(self, ids):
+        return [str(i) for i in ids]
+
+    def __len__(self):
+        return self.vocab_size

train_all.sh

@@ -0,0 +1,28 @@
+#!/bin/bash
+set -euo pipefail
+# Run this script to train all ablation models
+
+MODEL_TYPE=${MODEL_TYPE:-d3pm_cross_attention}
+INCLUDE_NEG=${INCLUDE_NEG:-False}
+TRAIN_DEVICE=${TRAIN_DEVICE:-mps}
+
+echo '=== Training T=4 ==='
+mkdir -p ablation_results/T4
+MODEL_TYPE="$MODEL_TYPE" INCLUDE_NEG="$INCLUDE_NEG" TRAIN_DEVICE="$TRAIN_DEVICE" DIFFUSION_STEPS=4 INFERENCE_NUM_STEPS=4 TRAIN_OUTPUT_DIR="ablation_results/T4" python train.py
+
+echo '=== Training T=8 ==='
+mkdir -p ablation_results/T8
+MODEL_TYPE="$MODEL_TYPE" INCLUDE_NEG="$INCLUDE_NEG" TRAIN_DEVICE="$TRAIN_DEVICE" DIFFUSION_STEPS=8 INFERENCE_NUM_STEPS=8 TRAIN_OUTPUT_DIR="ablation_results/T8" python train.py
+
+echo '=== Training T=16 ==='
+mkdir -p ablation_results/T16
+MODEL_TYPE="$MODEL_TYPE" INCLUDE_NEG="$INCLUDE_NEG" TRAIN_DEVICE="$TRAIN_DEVICE" DIFFUSION_STEPS=16 INFERENCE_NUM_STEPS=16 TRAIN_OUTPUT_DIR="ablation_results/T16" python train.py
+
+echo '=== Training T=32 ==='
+mkdir -p ablation_results/T32
+MODEL_TYPE="$MODEL_TYPE" INCLUDE_NEG="$INCLUDE_NEG" TRAIN_DEVICE="$TRAIN_DEVICE" DIFFUSION_STEPS=32 INFERENCE_NUM_STEPS=32 TRAIN_OUTPUT_DIR="ablation_results/T32" python train.py
+
+echo '=== Training T=64 ==='
+mkdir -p ablation_results/T64
+MODEL_TYPE="$MODEL_TYPE" INCLUDE_NEG="$INCLUDE_NEG" TRAIN_DEVICE="$TRAIN_DEVICE" DIFFUSION_STEPS=64 INFERENCE_NUM_STEPS=64 TRAIN_OUTPUT_DIR="ablation_results/T64" python train.py
+