Update inference cleanup + model card + runtime

README.md

@@ -9,16 +9,13 @@ tags:
 - diffusion
 - d3pm
 - pytorch
-pipeline_tag: text-generation
+pipeline_tag: text2text-generation
 ---
 
 # Sanskrit D3PM Paraphrase Model
 
 Roman/IAST Sanskrit input to Devanagari output using a D3PM cross-attention model.
 
-This is a **custom PyTorch architecture** (not a native `transformers.AutoModel` checkpoint).
-You can still use it in a transformer-like workflow (load once, pass text, get generated text) via `inference_api.py`.
-
 ## Files Included
 
 - `best_model.pt` — trained checkpoint
@@ -28,7 +25,6 @@ You can still use it in a transformer-like workflow (load once, pass text, get g
 - `handler.py` — Hugging Face Endpoint handler
 - `model/`, `diffusion/` — architecture modules
 - `sanskrit_src_tokenizer.json`, `sanskrit_tgt_tokenizer.json` — tokenizers
-- `LOCAL_SETUP_GUIDE.md` — full laptop setup and execution guide
 
 ## Quick Local Test
 
@@ -37,52 +33,30 @@ from inference_api import predict
 print(predict("dharmo rakṣati rakṣitaḥ")["output"])
 ```
 
-## Transformer-Style Usage (Recommended)
+## Transformer-Style Usage (Custom Runtime)
 
-Use this model like a transformer pipeline pattern: load once, call `generate(text)` many times.
+This checkpoint is a custom D3PM architecture (`.pt`), not a native `transformers` `AutoModel` format.
+Use it in a transformer-like way via the provided runtime:
 
 ```python
 import torch
 from config import CONFIG
-from inference import load_model, _build_tokenizers
+from inference import load_model, run_inference, _decode_clean
+from model.tokenizer import SanskritSourceTokenizer, SanskritTargetTokenizer
 
-cfg = CONFIG
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model, cfg = load_model("best_model.pt", CONFIG, device)
 
-model, cfg = load_model("best_model.pt", cfg, device)
-src_tok, tgt_tok = _build_tokenizers(cfg)
-
-def generate(text: str):
-    input_ids = torch.tensor([src_tok.encode(text)], dtype=torch.long, device=device)
-    output_ids = model.generate(
-        input_ids,
-        num_steps=cfg["inference"]["num_steps"],
-        temperature=cfg["inference"]["temperature"],
-        top_k=cfg["inference"]["top_k"],
-        repetition_penalty=cfg["inference"]["repetition_penalty"],
-        diversity_penalty=cfg["inference"]["diversity_penalty"],
-    )
-    ids = [x for x in output_ids[0].tolist() if x > 4]
-    return tgt_tok.decode(ids).strip()
-
-print(generate("yadā mano nivarteta viṣayebhyaḥ svabhāvataḥ"))
-```
-
-### Minimal 3-Step Pattern
-
-1. `load_model(...)` once at app startup  
-2. `encode -> model.generate(...) -> decode` for each request  
-3. Reuse loaded model/tokenizers for all requests
+src_tok = SanskritSourceTokenizer(vocab_size=16000, max_len=cfg["model"]["max_seq_len"])
+tgt_tok = SanskritTargetTokenizer(vocab_size=16000, max_len=cfg["model"]["max_seq_len"])
 
-## About `transformers` Compatibility
+text = "dharmo rakṣati rakṣitaḥ"
+ids = torch.tensor([src_tok.encode(text)], dtype=torch.long, device=device)
+out = run_inference(model, ids, cfg)
+print(_decode_clean(tgt_tok, out[0].tolist()))
+```
 
-- This repo does not expose `config.json` + `model.safetensors` in `transformers` format.
-- This is not a PEFT/LoRA adapter repository.
-- If you want full `AutoModel`/`pipeline` compatibility, you must create a wrapper architecture and export weights into HF Transformers conventions.
-- For production today, use:
-  - `inference_api.py` for Python apps
-  - `handler.py` for HF Inference Endpoints
-  - `space_repo/app.py` for Gradio UI
+If you need full `transformers` compatibility (`AutoModel.from_pretrained`), export weights to a Hugging Face Transformers model format first.
 
 ## Endpoint Payload
 
@@ -113,9 +87,3 @@ git add .
 git commit -m "Initial model release"
 git push -u origin main
 ```
-
-## Full Local Laptop Guide
-
-For complete setup (training, inference, UI, tasks 1-5, ablation, and deployment), see:
-
-- `LOCAL_SETUP_GUIDE.md`

inference.py

@@ -1,122 +1,554 @@
-import copy
+"""
+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 json
 import torch
-import torch.nn.functional as F
+import os, sys
+import re
+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
 
 
-def _resolve_device(cfg: dict) -> torch.device:
-    requested = cfg["training"]["device"]
-    if requested == "cuda" and not torch.cuda.is_available():
-        requested = "cpu"
-    if requested == "mps" and not torch.backends.mps.is_available():
-        requested = "cpu"
-    cfg["training"]["device"] = requested
-    return torch.device(requested)
+# ── Checkpoint loader ─────────────────────────────────────────────────
+
+def _resolve_device(cfg_device: str) -> torch.device:
+    cfg_device = (cfg_device or "").lower()
+    if cfg_device == "cuda" and torch.cuda.is_available():
+        return torch.device("cuda")
+    if cfg_device == "mps" and torch.backends.mps.is_available():
+        return torch.device("mps")
+    if cfg_device in {"cpu", "cuda", "mps"}:
+        return torch.device("cpu")
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    if torch.backends.mps.is_available():
+        return torch.device("mps")
+    return torch.device("cpu")
+
+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)
+    raw_state = torch.load(ckpt_path, map_location=device)
+    model_state = model.state_dict()
+    filtered_state = {}
+    skipped_mismatch = []
+    for k, v in raw_state.items():
+        if k in model_state and hasattr(v, "shape") and hasattr(model_state[k], "shape"):
+            if tuple(v.shape) != tuple(model_state[k].shape):
+                skipped_mismatch.append((k, tuple(v.shape), tuple(model_state[k].shape)))
+                continue
+        filtered_state[k] = v
+
+    missing, unexpected = model.load_state_dict(filtered_state, 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 skipped_mismatch:
+        print(f"⚠️  Shape-mismatched keys skipped: {len(skipped_mismatch)}")
+
+    # Enable compact-attention branch only when checkpoint actually provides it.
+    has_compact = any(".compact_out_proj.weight" in k for k in filtered_state.keys())
+    if has_compact and hasattr(model, "model") and hasattr(model.model, "decoder_blocks"):
+        for block in model.model.decoder_blocks:
+            if hasattr(block, "cross_attn") and hasattr(block.cross_attn, "use_compact"):
+                block.cross_attn.use_compact = True
+        print("ℹ️  Compact cross-attention branch enabled from checkpoint.")
+    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 (same path as validation) ────────────────
+
+@torch.no_grad()
+def run_inference(model, input_ids, cfg):
+    """
+    Reverse diffusion sampling (clean path).
+    Uses cached reverse diffusion when available, otherwise model.generate().
+    """
+    inf = cfg['inference']
+    model.eval()
+    kwargs = dict(
+        num_steps=inf['num_steps'],
+        temperature=inf['temperature'],
+        top_k=inf['top_k'],
+        repetition_penalty=inf.get('repetition_penalty', 1.2),
+        diversity_penalty=inf.get('diversity_penalty', 0.0),
+    )
+    if hasattr(model, "generate_cached"):
+        out = model.generate_cached(input_ids, **kwargs)
+    else:
+        out = model.generate(input_ids, **kwargs)
+
+    # Optional retry with stronger anti-repetition settings.
+    if inf.get("auto_retry_on_repetition", True):
+        repeat_threshold = float(inf.get("repeat_ratio_threshold", 0.40))
+        max_repeat_run = int(inf.get("max_repeat_run", 4))
+        if _mean_repeat_ratio(out) >= repeat_threshold:
+            retry_kwargs = dict(kwargs)
+            retry_kwargs["temperature"] = max(0.6, float(kwargs["temperature"]) - 0.1)
+            retry_kwargs["top_k"] = max(20, int(kwargs["top_k"]) - 10)
+            retry_kwargs["repetition_penalty"] = max(float(kwargs["repetition_penalty"]), 1.6)
+            retry_kwargs["diversity_penalty"] = max(float(kwargs["diversity_penalty"]), 0.3)
+            if hasattr(model, "generate_cached"):
+                retry = model.generate_cached(input_ids, **retry_kwargs)
+            else:
+                retry = model.generate(input_ids, **retry_kwargs)
+            if _mean_repeat_ratio(retry) < _mean_repeat_ratio(out):
+                out = retry
+        out = _dedup_repeated_ids(out, max_repeat_run=max_repeat_run)
+
+    return out
+
+
+def _mean_repeat_ratio(ids_tensor: torch.Tensor) -> float:
+    if ids_tensor is None or ids_tensor.numel() == 0:
+        return 0.0
+    ratios = []
+    for row in ids_tensor:
+        ids = [int(x) for x in row.tolist() if int(x) > 4]
+        if len(ids) < 2:
+            ratios.append(0.0)
+            continue
+        repeats = sum(1 for i in range(1, len(ids)) if ids[i] == ids[i - 1])
+        ratios.append(repeats / max(1, len(ids) - 1))
+    return float(sum(ratios) / max(1, len(ratios)))
+
+
+def _dedup_repeated_ids(ids_tensor: torch.Tensor, max_repeat_run: int = 4) -> torch.Tensor:
+    """
+    Keep generation path unchanged, but clean extreme run-on token loops in final output ids.
+    """
+    if ids_tensor is None or ids_tensor.numel() == 0:
+        return ids_tensor
+    cleaned_rows = []
+    for row in ids_tensor.tolist():
+        out = []
+        prev = None
+        run = 0
+        for tok in row:
+            if tok <= 4:
+                out.append(tok)
+                prev = tok
+                run = 1
+                continue
+            if tok == prev:
+                run += 1
+                if run > max_repeat_run:
+                    continue
+            else:
+                run = 1
+            out.append(tok)
+            prev = tok
+        # Preserve original length for downstream decode assumptions.
+        if len(out) < len(row):
+            out.extend([1] * (len(row) - len(out)))
+        else:
+            out = out[:len(row)]
+        cleaned_rows.append(out)
+    return torch.tensor(cleaned_rows, dtype=ids_tensor.dtype, device=ids_tensor.device)
+
+
+def _decode_clean(tgt_tok, ids):
+    out = []
+    for x in ids:
+        if x in (1, 4) and out:
+            break
+        if x > 4:
+            out.append(x)
+    text = tgt_tok.decode(out).strip()
+    return _clean_repetition_text(text)
+
+
+def _clean_repetition_text(text: str, max_repeat_run: int = 3) -> str:
+    words = [w for w in text.split() if w.strip()]
+    if not words:
+        return text.strip()
+    cleaned = []
+    prev = None
+    run = 0
+    for w in words:
+        if w == prev:
+            run += 1
+            if run > max_repeat_run:
+                continue
+        else:
+            run = 1
+        cleaned.append(w)
+        prev = w
+    return " ".join(cleaned).strip()
+
+
+# ── Cleanup heuristics from UI inference pipeline ─────────────────────
+
+_IAST_VOWELS = [
+    ("ai", "ऐ"), ("au", "औ"),
+    ("ā", "आ"), ("ī", "ई"), ("ū", "ऊ"),
+    ("ṛ", "ऋ"), ("ṝ", "ॠ"), ("ḷ", "ऌ"), ("ḹ", "ॡ"),
+    ("a", "अ"), ("i", "इ"), ("u", "उ"),
+    ("e", "ए"), ("o", "ओ"),
+]
+_IAST_MATRAS = [
+    ("ai", "ै"), ("au", "ौ"),
+    ("ā", "ा"), ("ī", "ी"), ("ū", "ू"),
+    ("ṛ", "ृ"), ("ṝ", "ॄ"), ("ḷ", "ॢ"), ("ḹ", "ॣ"),
+    ("a", ""), ("i", "ि"), ("u", "ु"),
+    ("e", "े"), ("o", "ो"),
+]
+_IAST_CONS = [
+    ("kṣ", "क��ष"), ("jñ", "ज्ञ"), ("tr", "त्र"),
+    ("kh", "ख"), ("gh", "घ"), ("ch", "छ"), ("jh", "झ"),
+    ("ṭh", "ठ"), ("ḍh", "ढ"), ("th", "थ"), ("dh", "ध"),
+    ("ph", "फ"), ("bh", "भ"),
+    ("ṅ", "ङ"), ("ñ", "ञ"), ("ṭ", "ट"), ("ḍ", "ड"),
+    ("ṇ", "ण"), ("ś", "श"), ("ṣ", "ष"), ("ḥ", "ः"),
+    ("ṃ", "ं"), ("ṁ", "ं"),
+    ("y", "य"), ("r", "र"), ("l", "ल"), ("v", "व"),
+    ("s", "स"), ("h", "ह"),
+    ("k", "क"), ("g", "ग"), ("c", "च"), ("j", "ज"),
+    ("t", "त"), ("d", "द"), ("n", "न"),
+    ("p", "प"), ("b", "ब"), ("m", "म"),
+]
+_PUNCT = {".": "।", "|": "।", "||": "॥", ",": ",", "?": "?", "!": "!"}
+
+
+def _iast_to_deva(text: str) -> str:
+    s = (text or "").lower()
+    out = []
+    i = 0
+    pending_consonant = False
+
+    def _match_any(pairs, pos):
+        for k, v in pairs:
+            if s.startswith(k, pos):
+                return k, v
+        return None, None
+
+    while i < len(s):
+        if s[i].isspace():
+            pending_consonant = False
+            out.append(s[i])
+            i += 1
+            continue
+        if s[i:i+2] == "||":
+            pending_consonant = False
+            out.append(_PUNCT["||"])
+            i += 2
+            continue
+        if s[i] in _PUNCT:
+            pending_consonant = False
+            out.append(_PUNCT[s[i]])
+            i += 1
+            continue
+
+        v_key, v_deva = _match_any(_IAST_VOWELS, i)
+        if v_key:
+            if pending_consonant:
+                _, v_matra = _match_any(_IAST_MATRAS, i)
+                out[-1] = out[-1] + (v_matra or "")
+                pending_consonant = False
+            else:
+                out.append(v_deva)
+            i += len(v_key)
+            continue
+
+        c_key, c_deva = _match_any(_IAST_CONS, i)
+        if c_key:
+            if pending_consonant:
+                out[-1] = out[-1] + "्"
+            out.append(c_deva)
+            pending_consonant = True
+            i += len(c_key)
+            continue
+
+        out.append(s[i])
+        pending_consonant = False
+        i += 1
+
+    return "".join(out).strip()
 
 
-def _build_tokenizers(cfg):
+def _compute_cer(pred: str, ref: str) -> float:
+    if pred == ref:
+        return 0.0
+    if not pred or not ref:
+        return 1.0
+    m, n = len(pred), len(ref)
+    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]
+            cost = 0 if pred[i - 1] == ref[j - 1] else 1
+            dp[j] = min(dp[j] + 1, dp[j - 1] + 1, prev + cost)
+            prev = temp
+    return dp[n] / max(m, n)
+
+
+def _cleanup_thresholds(temperature: float, top_k: int):
+    temp = float(temperature)
+    k = max(1, int(top_k))
+    t_norm = max(0.0, min((temp - 0.4) / 0.6, 1.0))
+    k_norm = max(0.0, min((k - 20) / 80.0, 1.0))
+    diversity = 0.6 * t_norm + 0.4 * k_norm
+    cer_threshold = 0.10 + 0.18 * diversity
+    deva_ratio_threshold = 0.60 - 0.20 * diversity
+    return cer_threshold, deva_ratio_threshold
+
+
+def _decode_with_cleanup(tgt_tok, ids, src_text: str, inf_cfg: dict):
+    model_out = _decode_clean(tgt_tok, ids)
+    rule_out = _iast_to_deva(src_text.strip())
+    deva_chars = sum(1 for ch in model_out if "\u0900" <= ch <= "\u097F")
+    deva_ratio = deva_chars / max(1, len(model_out))
+    cer = _compute_cer(model_out, rule_out)
+    cer_thr, ratio_thr = _cleanup_thresholds(
+        inf_cfg.get("temperature", 0.8),
+        inf_cfg.get("top_k", 40),
+    )
+    if deva_ratio < ratio_thr or len(model_out) > 2.0 * max(1, len(rule_out)) or cer > cer_thr:
+        return rule_out
+    return model_out
+
+
+# ── Interactive demo ──────────────────────────────────────────────────
+
+def interactive_demo(checkpoint=None, single_text=None):
     from model.tokenizer import SanskritSourceTokenizer, SanskritTargetTokenizer
 
+    cfg    = CONFIG
+    device = _resolve_device(cfg['training'].get('device', 'cpu'))
+
+    model_name = cfg['model_type']
+    has_neg    = cfg['data']['include_negative_examples']
+    ckpt       = checkpoint or 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()
+
     src_tok = SanskritSourceTokenizer(
-        vocab_size=cfg["model"].get("src_vocab_size", 16000),
-        max_len=cfg["model"]["max_seq_len"],
+        vocab_size=cfg['model'].get('src_vocab_size', 16000),
+        max_len=cfg['model']['max_seq_len'],
     )
     tgt_tok = SanskritTargetTokenizer(
-        vocab_size=cfg["model"].get("tgt_vocab_size", 16000),
-        max_len=cfg["model"]["max_seq_len"],
+        vocab_size=cfg['model'].get('tgt_vocab_size', 16000),
+        max_len=cfg['model']['max_seq_len'],
     )
-    return src_tok, tgt_tok
 
+    print("\n" + "="*55)
+    print("Sanskrit D3PM Paraphrase — type verse, get paraphrase")
+    print("="*55 + "\n")
 
-def load_model(ckpt_path: str, base_cfg: dict, device: torch.device):
-    from model.sanskrit_model import SanskritModel
-
-    cfg = copy.deepcopy(base_cfg)
-    state = torch.load(ckpt_path, map_location="cpu")
+    while True:
+        try:
+            text = (single_text if single_text is not None else input("INPUT > ")).strip()
+        except (EOFError, KeyboardInterrupt):
+            break
+        if not text or text.lower() in ('quit', 'exit', 'q'):
+            break
 
-    emb_key = "model.src_embed.token_emb.weight"
-    if emb_key in state:
-        vocab, d_model = state[emb_key].shape
-        cfg["model"]["src_vocab_size"] = vocab
-        cfg["model"]["d_model"] = d_model
-        cfg["model"]["d_ff"] = d_model * 4
+        ids = torch.tensor(
+            [src_tok.encode(text)[:cfg['model']['max_seq_len']]],
+            dtype=torch.long, device=device
+        )
+        out   = run_inference(model, ids, cfg)
+        cleaned = _decode_with_cleanup(tgt_tok, out[0].tolist(), text, cfg["inference"])
+        print(f"PARAPHRASE → {cleaned}\n")
+        if single_text is not None:
+            break
 
-    layer_ids = {int(k.split(".")[2]) for k in state if k.startswith("model.encoder_blocks.")}
-    if layer_ids:
-        cfg["model"]["n_layers"] = max(layer_ids) + 1
 
-    pos_key = "model.src_embed.pos_enc.pe"
-    if pos_key in state:
-        cfg["model"]["max_seq_len"] = state[pos_key].shape[1]
+# ── Batch evaluation ──────────────────────────────────────────────────
 
-    d_model = cfg["model"]["d_model"]
-    n_heads = cfg["model"].get("n_heads", 8)
-    if d_model % n_heads != 0:
-        n_heads = next(h for h in [8, 6, 4, 2, 1] if d_model % h == 0)
-    cfg["model"]["n_heads"] = n_heads
+def batch_evaluate(sample_size=500, checkpoint=None):
+    from data.dataset import OptimizedSanskritDataset
+    from model.tokenizer import SanskritSourceTokenizer, SanskritTargetTokenizer
 
-    model = SanskritModel(cfg).to(device)
-    model.load_state_dict(torch.load(ckpt_path, map_location=device), strict=False)
-    model.eval()
-    return model, cfg
+    cfg    = CONFIG
+    device = _resolve_device(cfg['training'].get('device', 'cpu'))
 
+    model_name = cfg['model_type']
+    has_neg    = cfg['data']['include_negative_examples']
+    exp_dir    = f"results/{model_name}_neg_{has_neg}"
+    ckpt       = checkpoint or f"{exp_dir}/best_model.pt"
 
-def run_inference(model, input_ids, cfg):
-    inf = cfg["inference"]
-    device = input_ids.device
-    bsz, seqlen = input_ids.shape
-    inner = model.model
+    if not os.path.exists(ckpt):
+        raise FileNotFoundError(f"No checkpoint at {ckpt}")
 
-    total_steps = inner.scheduler.num_timesteps
-    steps = int(inf["num_steps"])
-    step_size = max(1, total_steps // max(steps, 1))
-    timesteps = list(range(total_steps - 1, -1, -step_size))
-    if timesteps[-1] != 0:
-        timesteps.append(0)
+    model, cfg = load_model(ckpt, cfg, device)
+    model.eval()
 
-    x0_est = torch.full((bsz, seqlen), inner.mask_token_id, dtype=torch.long, device=device)
-    hint = None
+    src_tok = SanskritSourceTokenizer(
+        vocab_size=cfg['model'].get('src_vocab_size', 16000),
+        max_len=cfg['model']['max_seq_len'],
+    )
+    tgt_tok = SanskritTargetTokenizer(
+        vocab_size=cfg['model'].get('tgt_vocab_size', 16000),
+        max_len=cfg['model']['max_seq_len'],
+    )
 
-    with torch.no_grad():
-        for i, t_val in enumerate(timesteps):
-            is_last = i == len(timesteps) - 1
-            t = torch.full((bsz,), t_val, dtype=torch.long, device=device)
+    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],
+        }
 
-            logits, _ = model(input_ids, x0_est, t, x0_hint=hint, inference_mode=True)
+    dataset = OptimizedSanskritDataset(
+        split='test',
+        max_len=cfg['model']['max_seq_len'],
+        cfg=cfg,
+        src_tokenizer=src_tok,
+        tgt_tokenizer=tgt_tok,
+    )
+    indices = list(range(min(sample_size, len(dataset))))
+    loader  = DataLoader(
+        Subset(dataset, indices),
+        batch_size=cfg['training']['batch_size'],
+        shuffle=False, collate_fn=collate
+    )
 
-            if inf["repetition_penalty"] != 1.0:
-                from model.d3pm_model_cross_attention import _apply_repetition_penalty
+    all_preds, all_refs, all_inputs = [], [], []
+    print(f"⏳ Generating {len(indices)} paraphrases …")
 
-                logits = _apply_repetition_penalty(logits, x0_est, float(inf["repetition_penalty"]))
-            if inf["diversity_penalty"] > 0.0:
-                from model.d3pm_model_cross_attention import _apply_diversity_penalty_fixed
+    for batch in tqdm(loader):
+        ids = batch['input_ids'].to(device)
+        out = run_inference(model, ids, cfg)
+        for i in range(out.size(0)):
+            all_preds.append(_decode_with_cleanup(
+                tgt_tok, out[i].tolist(), batch['input_text'][i], cfg["inference"]
+            ))
+            all_refs.append(batch['target_text'][i].strip())
+            all_inputs.append(batch['input_text'][i].strip())
 
-                logits = _apply_diversity_penalty_fixed(logits, float(inf["diversity_penalty"]))
+    # 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
 
-            logits = logits / max(float(inf["temperature"]), 1e-5)
-            if int(inf["top_k"]) > 0:
-                from model.d3pm_model_cross_attention import _top_k_filter
+    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
 
-                logits = _top_k_filter(logits, int(inf["top_k"]))
+    # Save
+    out_path = f"{exp_dir}/evaluation_results.txt"
+    pred_path = f"{exp_dir}/evaluation_predictions.jsonl"
+    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")
 
-            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
+    with open(pred_path, 'w', encoding='utf-8') as f:
+        for src, ref, pred in zip(all_inputs, all_refs, all_preds):
+            row = {"input": src, "reference": ref, "prediction": pred}
+            f.write(json.dumps(row, ensure_ascii=False) + "\n")
 
-                x0_est = _batch_multinomial(probs)
-            hint = x0_est
+    print(f"\n✅ Results → {out_path}")
+    print(f"🗂️  Saved predictions → {pred_path}")
+    print(f"📊 BLEU: {bleu_score:.4f}  |  BERTScore: {bert_f1:.4f}")
+    return all_preds, all_refs
 
-    return x0_est
 
+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)
+    p.add_argument('--checkpoint', type=str, default=None)
+    p.add_argument('--text', type=str, default=None, help='Run one-shot demo input and exit')
+    args = p.parse_args()
 
-__all__ = [
-    "CONFIG",
-    "_resolve_device",
-    "_build_tokenizers",
-    "load_model",
-    "run_inference",
-]
+    if args.mode == 'demo':
+        interactive_demo(checkpoint=args.checkpoint, single_text=args.text)
+    else:
+        batch_evaluate(args.samples, checkpoint=args.checkpoint)