Add files using upload-large-folder tool

.gitattributes

@@ -1,35 +1,3 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,65 @@
+---
+license: mit
+language:
+- sa
+- en
+tags:
+- sanskrit
+- paraphrase
+- diffusion
+- d3pm
+- pytorch
+pipeline_tag: text2text-generation
+---
+
+# Sanskrit D3PM Paraphrase Model
+
+Roman/IAST Sanskrit input to Devanagari output using a D3PM cross-attention model.
+
+## Files Included
+
+- `best_model.pt` — trained checkpoint
+- `config.py` — runtime config
+- `inference.py` — model loading + generation loop
+- `inference_api.py` — simple Python API (`predict`)
+- `handler.py` — Hugging Face Endpoint handler
+- `model/`, `diffusion/` — architecture modules
+- `sanskrit_src_tokenizer.json`, `sanskrit_tgt_tokenizer.json` — tokenizers
+
+## Quick Local Test
+
+```python
+from inference_api import predict
+print(predict("dharmo rakṣati rakṣitaḥ")["output"])
+```
+
+## Endpoint Payload
+
+```json
+{
+  "inputs": "yadā mano nivarteta viṣayebhyaḥ svabhāvataḥ",
+  "parameters": {
+    "temperature": 0.7,
+    "top_k": 40,
+    "repetition_penalty": 1.2,
+    "diversity_penalty": 0.0,
+    "num_steps": 64,
+    "clean_output": true
+  }
+}
+```
+
+## Push This Folder To Model Hub
+
+```bash
+huggingface-cli login
+huggingface-cli repo create <your-username>/sanskrit-d3pm --type model
+cd hf_model_repo
+git init
+git lfs install
+git remote add origin https://huggingface.co/<your-username>/sanskrit-d3pm
+git add .
+git commit -m "Initial model release"
+git push -u origin main
+```
+

best_model.pt

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

config.py

@@ -0,0 +1,33 @@
+import torch
+
+CONFIG = {
+    "model_type": "d3pm_cross_attention",
+    "data": {
+        "include_negative_examples": True,
+        "dataset_size": 60000,
+    },
+    "diffusion": {
+        "mask_token_id": 0,
+    },
+    "model": {
+        "src_vocab_size": 16000,
+        "tgt_vocab_size": 16000,
+        "d_model": 384,
+        "n_heads": 8,
+        "d_ff": 1536,
+        "n_layers": 6,
+        "dropout": 0.1,
+        "max_seq_len": 80,
+        "diffusion_steps": 64,
+    },
+    "training": {
+        "device": "cuda" if torch.cuda.is_available() else "cpu",
+    },
+    "inference": {
+        "num_steps": 64,
+        "temperature": 0.7,
+        "top_k": 40,
+        "repetition_penalty": 1.2,
+        "diversity_penalty": 0.0,
+    },
+}

diffusion/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

diffusion/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)

diffusion/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

diffusion/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)

diffusion/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()]

handler.py

@@ -0,0 +1,30 @@
+from typing import Any, Dict
+
+from inference_api import predict
+
+
+class EndpointHandler:
+    """
+    Hugging Face Inference Endpoint handler.
+    Expects payload:
+      {
+        "inputs": "dharmo rakṣati rakṣitaḥ",
+        "parameters": {"temperature": 0.7, ...}
+      }
+    """
+
+    def __init__(self, path: str = ""):
+        self.path = path
+
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        text = data.get("inputs", "")
+        params = data.get("parameters", {}) or {}
+        return predict(
+            text=text,
+            temperature=params.get("temperature", 0.7),
+            top_k=params.get("top_k", 40),
+            repetition_penalty=params.get("repetition_penalty", 1.2),
+            diversity_penalty=params.get("diversity_penalty", 0.0),
+            num_steps=params.get("num_steps", 64),
+            clean_output=params.get("clean_output", True),
+        )

inference.py

@@ -0,0 +1,122 @@
+import copy
+
+import torch
+import torch.nn.functional as F
+
+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)
+
+
+def _build_tokenizers(cfg):
+    from model.tokenizer import SanskritSourceTokenizer, SanskritTargetTokenizer
+
+    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"],
+    )
+    return src_tok, tgt_tok
+
+
+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")
+
+    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
+
+    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]
+
+    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
+
+    model = SanskritModel(cfg).to(device)
+    model.load_state_dict(torch.load(ckpt_path, map_location=device), strict=False)
+    model.eval()
+    return model, cfg
+
+
+def run_inference(model, input_ids, cfg):
+    inf = cfg["inference"]
+    device = input_ids.device
+    bsz, seqlen = input_ids.shape
+    inner = model.model
+
+    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)
+
+    x0_est = torch.full((bsz, seqlen), inner.mask_token_id, dtype=torch.long, device=device)
+    hint = None
+
+    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)
+
+            logits, _ = model(input_ids, x0_est, t, x0_hint=hint, inference_mode=True)
+
+            if inf["repetition_penalty"] != 1.0:
+                from model.d3pm_model_cross_attention import _apply_repetition_penalty
+
+                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
+
+                logits = _apply_diversity_penalty_fixed(logits, float(inf["diversity_penalty"]))
+
+            logits = logits / max(float(inf["temperature"]), 1e-5)
+            if int(inf["top_k"]) > 0:
+                from model.d3pm_model_cross_attention import _top_k_filter
+
+                logits = _top_k_filter(logits, int(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
+
+
+__all__ = [
+    "CONFIG",
+    "_resolve_device",
+    "_build_tokenizers",
+    "load_model",
+    "run_inference",
+]

inference_api.py

@@ -0,0 +1,103 @@
+import copy
+from typing import Dict, Any
+
+import torch
+
+from config import CONFIG
+from inference import _build_tokenizers, _resolve_device, load_model, run_inference
+
+
+_STATE = {
+    "loaded": False,
+    "model": None,
+    "cfg": None,
+    "device": None,
+    "src_tok": None,
+    "tgt_tok": None,
+}
+
+
+def _load_once() -> None:
+    if _STATE["loaded"]:
+        return
+
+    cfg = copy.deepcopy(CONFIG)
+    cfg["model_type"] = "d3pm_cross_attention"
+    cfg["data"]["include_negative_examples"] = True
+    device = _resolve_device(cfg)
+
+    model, cfg = load_model("best_model.pt", cfg, device)
+    src_tok, tgt_tok = _build_tokenizers(cfg)
+
+    _STATE["model"] = model
+    _STATE["cfg"] = cfg
+    _STATE["device"] = device
+    _STATE["src_tok"] = src_tok
+    _STATE["tgt_tok"] = tgt_tok
+    _STATE["loaded"] = True
+
+
+def _clean_text(text: str) -> str:
+    text = " ".join(text.split())
+    if not text:
+        return text
+    toks = text.split()
+    out = []
+    prev = None
+    run = 0
+    for tok in toks:
+        if tok == prev:
+            run += 1
+        else:
+            prev = tok
+            run = 1
+        if run <= 2:
+            out.append(tok)
+    s = " ".join(out)
+    s = s.replace(" ।", "।").replace(" ॥", "॥")
+    return " ".join(s.split())
+
+
+def predict(
+    text: str,
+    temperature: float = 0.7,
+    top_k: int = 40,
+    repetition_penalty: float = 1.2,
+    diversity_penalty: float = 0.0,
+    num_steps: int = 64,
+    clean_output: bool = True,
+) -> Dict[str, Any]:
+    _load_once()
+    if not text or not text.strip():
+        return {"error": "empty input", "output": ""}
+
+    cfg = copy.deepcopy(_STATE["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 = _STATE["src_tok"]
+    tgt_tok = _STATE["tgt_tok"]
+    device = _STATE["device"]
+
+    input_ids = torch.tensor([src_tok.encode(text.strip())], dtype=torch.long, device=device)
+    out = run_inference(_STATE["model"], input_ids, cfg)
+    decoded_ids = [x for x in out[0].tolist() if x > 4]
+    raw = tgt_tok.decode(decoded_ids).strip()
+    output = _clean_text(raw) if clean_output else raw
+
+    return {
+        "input": text,
+        "output": output,
+        "raw_output": raw,
+        "config": {
+            "temperature": float(temperature),
+            "top_k": int(top_k),
+            "repetition_penalty": float(repetition_penalty),
+            "diversity_penalty": float(diversity_penalty),
+            "num_steps": int(num_steps),
+            "clean_output": bool(clean_output),
+        },
+    }

model/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)

model/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:]

model/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)

model/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

model/tokenizers.py

@@ -0,0 +1,112 @@
+"""
+tokenizer.py — FINAL
+=====================
+Uses the original sanskrit_tokenizer_m4pro.json — the exact one the model
+was trained with. Hard-coded absolute path as primary, with fallbacks.
+
+This tokenizer has NO </w> end-of-word markers and NO decoder set.
+decode() returns space-separated BPE pieces — this is the format the
+model was trained and evaluated on (BERTScore 0.71). Do NOT add a decoder
+or retrain: that would break alignment with the checkpoint.
+"""
+
+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
+import os
+
+# Hard-coded absolute path — update if you move the project
+TOKENIZER_PATH = "/Users/bhsingh/Documents/Final_Paraphrase/sanskrit_tokenizer_m4pro.json"
+
+
+def build_tokenizer(texts, vocab_size=16000):
+    tokenizer = Tokenizer(BPE(unk_token="[UNK]"))
+    tokenizer.pre_tokenizer = Whitespace()
+    trainer = BpeTrainer(
+        vocab_size=vocab_size,
+        special_tokens=["[MASK]", "[PAD]", "[UNK]", "[CLS]", "[SEP]"],
+        min_frequency=2,
+    )
+    tokenizer.train_from_iterator(texts, trainer)
+    return tokenizer
+
+
+class SanskritTokenizer:
+    def __init__(self, vocab_size=16000, max_len=80):
+        self.vocab_size = vocab_size
+        self.max_len    = max_len
+        self.mask_token_id = 0
+
+        script_dir = Path(__file__).resolve().parent
+        candidates = [
+            os.environ.get("SANSKRIT_TOKENIZER_PATH", ""),
+            TOKENIZER_PATH,
+            str(script_dir.parent / "sanskrit_tokenizer_m4pro.json"),
+            str(script_dir / "sanskrit_tokenizer_m4pro.json"),
+            str(Path.cwd() / "sanskrit_tokenizer_m4pro.json"),
+        ]
+
+        self.model_path = None
+        for c in candidates:
+            if c and Path(c).exists():
+                self.model_path = c
+                break
+
+        if self.model_path:
+            print(f"📖 Loading tokenizer from: {self.model_path}")
+            self.tokenizer = Tokenizer.from_file(self.model_path)
+            self._validate_mask_token()
+        else:
+            print(f"⚠️  Tokenizer not found at any candidate path.")
+            print(f"    Expected: {TOKENIZER_PATH}")
+            print("    Retraining — WARNING: output will not match existing checkpoint!")
+            self.model_path = TOKENIZER_PATH
+            self._train_tokenizer()
+
+    def _validate_mask_token(self):
+        mask_id = self.tokenizer.token_to_id("[MASK]")
+        assert mask_id == 0, f"[MASK] must be ID 0, got {mask_id}"
+        print("✅ [MASK] token confirmed at ID=0")
+
+    def _train_tokenizer(self):
+        dataset = load_dataset("paws/sanskrit-verses-gretil", split="train")
+        texts = []
+        for sample in dataset.select(range(50000)):
+            texts.extend([sample["quote_text"], sample["quote_devanagari"]])
+        tokenizer = build_tokenizer(texts, self.vocab_size)
+        tokenizer.save(self.model_path)
+        self.tokenizer = tokenizer
+        self._validate_mask_token()
+        print(f"✅ Tokenizer saved to: {self.model_path}")
+
+    def encode(self, text):
+        encoded   = self.tokenizer.encode(text)
+        token_ids = encoded.ids[:self.max_len]
+        pad_id    = self.tokenizer.token_to_id("[PAD]")
+        if len(token_ids) < self.max_len:
+            token_ids += [pad_id] * (self.max_len - len(token_ids))
+        return token_ids[:self.max_len]
+
+    def decode(self, ids):
+        if isinstance(ids, list) and len(ids) > 0 and isinstance(ids[0], list):
+            raise TypeError("decode() expects 1D list of IDs, not 2D.")
+        # Filter special tokens: 0=MASK 1=PAD 2=UNK 3=CLS 4=SEP
+        clean = [i for i in ids if isinstance(i, int) and i > 4]
+        if not clean:
+            return ""
+        return self.tokenizer.decode(clean, skip_special_tokens=True).strip()
+
+    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

requirements.txt

@@ -0,0 +1,6 @@
+torch>=2.2
+numpy>=1.24
+tqdm>=4.66
+datasets>=2.19
+tokenizers>=0.15
+scikit-learn>=1.3