Upload folder using huggingface_hub

.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

.gitignore

@@ -0,0 +1,2 @@
+__pycache__/
+*.pyc

README.md

@@ -0,0 +1,122 @@
+---
+license: mit
+language:
+- sa
+- en
+tags:
+- sanskrit
+- paraphrase
+- diffusion
+- d3pm
+- pytorch
+pipeline_tag: text2text-generation
+---
+
+# Sanskrit D3PM Encoder-Decoder Model
+
+Roman/IAST Sanskrit input to Devanagari output using a custom D3PM checkpoint.
+This package is configured for the `d3pm_encoder_decoder` checkpoint stored in
+`best_model.pt`.
+Hugging Face model repo: `bhsinghgrid/devflow2`
+
+## Files Included
+
+- `best_model.pt` — trained checkpoint
+- `model_settings.json` — packaged runtime metadata
+- `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"])
+```
+
+## Runtime Settings
+
+For local/API usage, the runtime first reads `model_settings.json`, then allows
+optional environment variable overrides:
+
+- `HF_MODEL_TYPE` = `d3pm_cross_attention` or `d3pm_encoder_decoder`
+- `HF_INCLUDE_NEG` = `true` or `false`
+- `HF_NUM_STEPS` = diffusion step count for the packaged checkpoint
+
+Packaged settings for this repo:
+
+```bash
+export HF_MODEL_TYPE=d3pm_encoder_decoder
+export HF_INCLUDE_NEG=false
+export HF_NUM_STEPS=4
+```
+
+## Use This Model In A Hugging Face Space
+
+In your Space settings, set:
+
+- `HF_CHECKPOINT_REPO=bhsinghgrid/devflow2`
+- `HF_CHECKPOINT_FILE=best_model.pt`
+
+If your Space reads model metadata automatically, no extra model-type variables
+are required. If it does not, also set:
+
+```bash
+HF_DEFAULT_MODEL_TYPE=d3pm_encoder_decoder
+HF_DEFAULT_INCLUDE_NEG=false
+HF_DEFAULT_NUM_STEPS=4
+```
+
+## Transformer-Style Usage (Custom Runtime)
+
+This checkpoint is a custom D3PM architecture (`.pt`), not a native `transformers`
+`AutoModel` format. Use it via the provided runtime:
+
+```python
+import torch
+from config import CONFIG
+from inference import load_model, run_inference, _decode_clean
+from model.tokenizer import SanskritSourceTokenizer, SanskritTargetTokenizer
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model, cfg = load_model("best_model.pt", CONFIG, device)
+
+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"])
+
+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()))
+```
+
+If you need full `transformers` compatibility (`AutoModel.from_pretrained`),
+export weights to a Hugging Face Transformers model format first.
+
+## 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": 4,
+    "clean_output": true
+  }
+}
+```
+
+## Push This Folder To Model Hub
+
+```bash
+cd hf_model_repo_encoder_decoder
+git add .
+git commit -m "Add encoder-decoder T4 model package"
+git push -u hf main
+```

analysis_reports/outputs_all_models_20260325/T16/task1_kv_cache.txt

@@ -0,0 +1,15 @@
+TASK 1 — KV CACHE BENCHMARK
+========================================
+
+has_generate_cached=True
+memory_profile=Torch CPU mem-event reduction: 30.4% @ src_len=64 (std=2143.0MB, cache=1492.1MB)
+
+ src_len   standard(s)   cached(s)   speedup   encoder%
+      16         0.893       0.571     1.56x      40.0%
+      32         0.751       0.509     1.48x      42.3%
+      64         1.141       0.822     1.39x      40.7%
+
+Saved graphs:
+  - task1_time_comparison.png
+  - task1_speedup.png
+  - task1_encoder_cost.png

analysis_reports/outputs_all_models_20260325/T16/task2_report.txt

@@ -0,0 +1,35 @@
+TASK 2 — ATTENTION + DRIFT REPORT
+==================================================
+
+Input : dharmo rakṣati rakṣitaḥ
+Output: धर्मो रक्षति रक्षितः
+
+Captured steps: 16
+Analysis quality: WEAK
+Final output uniq-ratio: 1.000
+Degenerate output: NO
+Multi-sample semantic score (n<=8): 0.1471
+Lock-in step (CER<=0.05): t=0
+Locked tokens: 38  Flexible tokens: 42
+TF-IDF vs attention stability corr: 0.9294
+TF-IDF status: OK
+
+Saved graphs:
+  - task2_attn_t*.png / task2_all_layers_t0.png
+  - task2_attn_evolution.png
+  - task2_semantic_drift.png
+  - task2_source_alignment.png
+  - task2_tfidf_vs_attention.png
+
+Step trajectory (first 10 rows)
+------------------------------------------------------------
+t= 15  bert=0.0475  drift=0.9525  text=धर्मो ति रक्ष रक्षि तः तः तः तः ितः तः धर्मो धर्मो धर्मो धर्
+t= 14  bert=0.0478  drift=0.9522  text=धर्मो ति रक्ष रक्षि तः तः तः तः ितः तः धर्मो धर्मो धर्मो धर्
+t= 13  bert=0.0478  drift=0.9522  text=धर्मो ति रक्ष रक्षि तः तः तः तः ितः तः धर्मो धर्मो धर्मो धर्
+t= 12  bert=0.0478  drift=0.9522  text=धर्मो ति रक्ष रक्षि तः तः तः तः ितः तः धर्मो धर्मो धर्मो धर्
+t= 11  bert=0.0478  drift=0.9522  text=धर्मो ति रक्ष रक्षि तः तः तः तः ितः तः धर्मो धर्मो धर्मो धर्
+t= 10  bert=0.0478  drift=0.9522  text=धर्मो ति रक्ष रक्षि तः तः तः तः ितः तः धर्मो धर्मो धर्मो धर्
+t=  9  bert=0.0478  drift=0.9522  text=धर्मो ति रक्ष रक्षि तः तः तः तः ितः तः धर्मो धर्मो धर्मो धर्
+t=  8  bert=0.0478  drift=0.9522  text=धर्मो ति रक्ष रक्षि तः तः तः तः ितः तः धर्मो धर्मो धर्मो धर्
+t=  7  bert=0.0478  drift=0.9522  text=धर्मो ति रक्ष रक्षि तः तः तः तः ितः तः धर्मो धर्मो धर्मो धर्
+t=  6  bert=0.0478  drift=0.9522  text=धर्मो ति रक्ष रक्षि तः तः तः तः ितः तः धर्मो धर्मो धर्मो धर्

analysis_reports/outputs_all_models_20260325/T16/task3_report.txt

@@ -0,0 +1,21 @@
+TASK 3 — CONCEPT VECTORS + PCA STEERING
+==================================================
+
+PCA: 50 components, 74.8% variance
+Diversity PC: 0  (|r|=0.325 with diversity proxy)
+
+Direction validity: WEAK
+Spectrum unique ratio (mean over 5 seeds): 1.000
+Spectrum semantic stability (mean over 5 seeds): 0.312
+
+Saved graphs:
+  - task3_concept_space.png
+  - task3_pca_explained_variance.png
+  - task3_diversity_curve.png
+
+Diversity spectrum:
+  alpha=-2.0  →  बले वेध विवर् धान वीर्य वीर्य धिं सिंहा भि̱ सन वस्तु वेध वै वेध वस्तु सन सन सिंहा सिंहा वीर्य वीर्य वस्तु सन रुते प्रभवति मन वेध बले बले र्वृ प्रपूजयेत् युगा मलि धान तुल वीर्य वीर्य वीर्य वीर्य वीर्य वीर्य धान तुल कालेन युगा वेध बले वेध वेध च्छे ष्मस् यस्या काष्ठा ज्ञप्त अर्णव धिं धिं वस्तु धिं सन तया सन सन देवाः देवाः स्वातन्त्र अर्णव मह वस्तु मुष् सन धिं धिं धिं विक्र त्र मह हस्ते च्छे मह
+  alpha=-1.0  →  बले र् अ तुल वीर्य वीर्य गुरु सिंहा सन सन विलेप वै वै वै गतस्य वेध सन सिंहा सिंहा स्य स्य । सन वै वै वै बले बले बले बले र् अ अ तुल तुल वीर्य वीर्य वीर्य वीर्य वीर्य वीर्य तुल तुल तुल ् बले वेध दिव्यां मान वै अप्सु सन ॥ ॥ वस्तु सिंहा सन सन विक्र सन स काष्ठा सन सन सन कार सन सन सन सन भ बल ु सिंहा सन सिंहा सन म् म् सन
+  alpha=+0.0  →  बले र् अ तुल वीर्य वीर्य स्य सिंहा सन सन पितो वै वै वै दक्षिणां सन सन सिंहा सिंहा स्य स्य स्य सन गतस्य वै वै ॥ बले बले र् र् अ अ । तुल वीर्य वीर्य वीर्य वीर्य वीर्य तुल तुल तुल तुल अ स बले बले वै वै ॥ ॥ ॥ सन सन सिंहा स सन सन सन सन सन सन सन सन सन सन ॥ ॥ सन सन शतैः ॥ सिंहा सिंहा द सिंहा सन त् सन
+  alpha=+1.0  →  बले र् अ अ विशुद्धं स्य स्य सिंहा सिंहा सन गतस्य वै वै वै वेत्ति सन सन सिंहा स्य स्य स्य स्य सन वै वै स मल बले बले र् र् व अ अ तुल वीर्य वीर्य वीर्य स्य वीर्य स्य तुल ानु अ अ । र् व ॥ वै वै सन द ॥ ॥ सिंहा सिंहा ॥ सं सन ॥ ॥ व ॥ ॥ हेम सन सन व ॥ ै ॥ वै भ न न ॥ मित्रो सिंहा सन
+  alpha=+2.0  →  आविश र् अ किंचिद् वर स्य स्य सिंहा सं निमे ञ् सं वै वै ञ् सन कृपा सिंहा स्य स्य स्य स्य फणा ञ् वै ौ जिह्व बले मानाः र् र् वराय अ माने वर विशुद्धं स्य स्य स्य – वर विशुद्धं व वर अ कृपा ॥ परम् ॥ कश्चि वै ॥ ञ् ञ् सं स्य स्य तम् व प्रवर्तन्ते कर्मसु परम् वर ते ॥ व ञ् ॥ ॥ सं द ॥ ॥ वर न्द ̱व ॥ व व ै

analysis_reports/outputs_all_models_20260325/T16/task4_report.txt

@@ -0,0 +1,14 @@
+TASK 4 — SEMANTIC ROBUSTNESS ABLATION
+==================================================
+
+Optimal diffusion steps = 16
+
+     T     BERT-F1     SEM_SIM      BLEU    sec/sample
+  --------------------------------------------------------
+    16      0.2574      0.0580    0.0007        0.9068
+
+Marginal gains (BERT-F1):
+
+Saved plots/files:
+  - task4_3d.png
+  - task4_raw_results.json

analysis_reports/outputs_all_models_20260325/T16/task5_report.txt

@@ -0,0 +1,15 @@
+TASK 5 — CLASSIFIER-FREE GUIDANCE
+==================================================
+
+Classifier params: 139521
+Training samples : 40
+
+Guidance scale sweep:
+       λ       CER   diversity      d2     sBLEU
+  ----------------------------------------------------
+     0.0    0.8336       0.808   0.624     0.007 ← optimal
+     0.5    0.8362       0.800   0.606     0.007
+     1.0    0.8390       0.798   0.601     0.005
+     1.5    0.8458       0.813   0.631     0.004
+     2.0    0.8531       0.828   0.660     0.004
+     3.0    0.8773       0.830   0.669     0.009

analysis_reports/outputs_all_models_20260325/T32/task1_kv_cache.txt

@@ -0,0 +1,15 @@
+TASK 1 — KV CACHE BENCHMARK
+========================================
+
+has_generate_cached=True
+memory_profile=Torch CPU mem-event reduction: 31.1% @ src_len=64 (std=4287.2MB, cache=2953.9MB)
+
+ src_len   standard(s)   cached(s)   speedup   encoder%
+      16         1.914       1.165     1.64x      39.6%
+      32         1.542       0.891     1.73x      42.1%
+      64         2.096       1.475     1.42x      42.7%
+
+Saved graphs:
+  - task1_time_comparison.png
+  - task1_speedup.png
+  - task1_encoder_cost.png

analysis_reports/outputs_all_models_20260325/T32/task2_report.txt

@@ -0,0 +1,35 @@
+TASK 2 — ATTENTION + DRIFT REPORT
+==================================================
+
+Input : dharmo rakṣati rakṣitaḥ
+Output: धर्मो रक्षति रक्षितः
+
+Captured steps: 32
+Analysis quality: WEAK
+Final output uniq-ratio: 1.000
+Degenerate output: NO
+Multi-sample semantic score (n<=8): 0.0627
+Lock-in step (CER<=0.05): t=0
+Locked tokens: 75  Flexible tokens: 5
+TF-IDF vs attention stability corr: -0.0869
+TF-IDF status: WEAK
+
+Saved graphs:
+  - task2_attn_t*.png / task2_all_layers_t0.png
+  - task2_attn_evolution.png
+  - task2_semantic_drift.png
+  - task2_source_alignment.png
+  - task2_tfidf_vs_attention.png
+
+Step trajectory (first 10 rows)
+------------------------------------------------------------
+t= 31  bert=0.0167  drift=0.9833  text=तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ 
+t= 30  bert=0.0167  drift=0.9833  text=तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ 
+t= 29  bert=0.0167  drift=0.9833  text=तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ 
+t= 28  bert=0.0167  drift=0.9833  text=तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ 
+t= 27  bert=0.0167  drift=0.9833  text=तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ 
+t= 26  bert=0.0167  drift=0.9833  text=तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ 
+t= 25  bert=0.0167  drift=0.9833  text=तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ 
+t= 24  bert=0.0167  drift=0.9833  text=तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ 
+t= 23  bert=0.0167  drift=0.9833  text=तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ 
+t= 22  bert=0.0167  drift=0.9833  text=तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ तृ 

analysis_reports/outputs_all_models_20260325/T32/task3_report.txt

@@ -0,0 +1,21 @@
+TASK 3 — CONCEPT VECTORS + PCA STEERING
+==================================================
+
+PCA: 50 components, 94.6% variance
+Diversity PC: 0  (|r|=-0.530 with diversity proxy)
+
+Direction validity: WEAK
+Spectrum unique ratio (mean over 5 seeds): 0.840
+Spectrum semantic stability (mean over 5 seeds): 0.234
+
+Saved graphs:
+  - task3_concept_space.png
+  - task3_pca_explained_variance.png
+  - task3_diversity_curve.png
+
+Diversity spectrum:
+  alpha=-2.0  →  ेन श्रे श्रे ेन श्रे अण्ड व्याः श्रे तन्त्रा ॥ ॥ ॥ व्याः व्याः व्याः तद्वद् तद्वद् तद्वद् ॥ ॥ ॥ ॥ ॥ ॥ ॥ ॥ ॥ ॥ ॥ ॥ ॥ ॥ तद्वद् ॥ ॥ ॥ ॥ ॥ व्याः व्याः व्याः ॥ ॥ राजन्य व्याः व्याः व्याः ॥ व्याः व्याः ॥ ॥ काम्य ॥ ॥ ॥ व्याः ॥ तद्वद् ॥ ॥ ॥ ॥ ॥ तन्त्रा तन्त्रा ॥ ॥ ॥ ॥ व्याः ॥ ॥ ॥ ॥ ॥ युधम् तद्वद् युधम् ॥
+  alpha=-1.0  →  श्रे श्रे श्रे ेन श्रे श्रे श्रे श्रे अण्ड तन्त्रा व्याः ॥ अण्ड अण्ड तन्त्रा व्याः तद्वद् ॥ व्याः ॥ ॥ ॥ ॥ ॥ ॥ ॥ ॥ ॥ ॥ ॥ ॥ ॥ ॥ अण्ड ॥ ॥ ॥ व्याः ॥ व्याः नो̍ ॥ ॥ ॥ ॥ ॥ व्याः व्याः अण्ड ॥ ॥ तन्त्रा ॥ ॥ तद्वद् युधम् रोमा शम्भु ॥ धूमं तन्त्रा ॥ तन्त्रा ॥ व्याः ॥ ॥ ॥ ॥ ॥ ॥ ॥ ॥ ॥ ॥ ॥ ॥ ॥ ॥ ॥
+  alpha=+0.0  →  अण्ड श्रे करः श्रे तन्त्रा करः करः तन्त्रा श्रे अण्ड अण्ड अण्ड ॥ श्रे तद्वद् अण्ड ॥ ॥ अण्ड ॥ ॥ ॥ ॥ ॥ ॥ ॥ अण्ड ॥ ॥ ॥ ॥ अण्ड ॥ ॥ ॥ ॥ ॥ ॥ राजन्य तन्त्रा नो̍ ॥ ॥ ॥ ॥ ॥ व्याः ॥ अण्ड ॥ काम्य ॥ ॥ ॥ ॥ ॥ शम्भु धूमं तन्त्रा तन्त्रा ेन ॥ काम्य ॥ ॥ करः तन्त्रा ॥ अण्ड ॥ अण्ड ॥ विनिर्जित्य ॥ ॥ ॥ तन्त्रा अण्ड तद्वद् करः
+  alpha=+1.0  →  माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण
+  alpha=+2.0  →  माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण माण

analysis_reports/outputs_all_models_20260325/T32/task4_report.txt

@@ -0,0 +1,14 @@
+TASK 4 — SEMANTIC ROBUSTNESS ABLATION
+==================================================
+
+Optimal diffusion steps = 32
+
+     T     BERT-F1     SEM_SIM      BLEU    sec/sample
+  --------------------------------------------------------
+    32      0.0422      0.0012    0.0000        1.8451
+
+Marginal gains (BERT-F1):
+
+Saved plots/files:
+  - task4_3d.png
+  - task4_raw_results.json

analysis_reports/outputs_all_models_20260325/T32/task5_report.txt

@@ -0,0 +1,15 @@
+TASK 5 — CLASSIFIER-FREE GUIDANCE
+==================================================
+
+Classifier params: 139521
+Training samples : 40
+
+Guidance scale sweep:
+       λ       CER   diversity      d2     sBLEU
+  ----------------------------------------------------
+     0.0    0.9357       0.239   0.011     0.533 ← optimal
+     0.5    0.9372       0.251   0.015     0.512
+     1.0    0.9467       0.164   0.018     0.690
+     1.5    0.9528       0.137   0.017     0.743
+     2.0    0.9525       0.144   0.013     0.725
+     3.0    0.9496       0.181   0.018     0.656

analysis_reports/outputs_all_models_20260325/T4/task1_kv_cache.txt

@@ -0,0 +1,15 @@
+TASK 1 — KV CACHE BENCHMARK
+========================================
+
+has_generate_cached=True
+memory_profile=Torch CPU mem-event reduction: 24.6% @ src_len=64 (std=525.8MB, cache=396.4MB)
+
+ src_len   standard(s)   cached(s)   speedup   encoder%
+      16         0.267       0.173     1.54x      43.2%
+      32         0.197       0.153     1.29x      40.7%
+      64         0.353       0.265     1.33x      42.0%
+
+Saved graphs:
+  - task1_time_comparison.png
+  - task1_speedup.png
+  - task1_encoder_cost.png

analysis_reports/outputs_all_models_20260325/T4/task2_report.txt

@@ -0,0 +1,29 @@
+TASK 2 — ATTENTION + DRIFT REPORT
+==================================================
+
+Input : dharmo rakṣati rakṣitaḥ
+Output: धर्मो रक्षति रक्षितः
+
+Captured steps: 4
+Analysis quality: VALID
+Final output uniq-ratio: 1.000
+Degenerate output: NO
+Multi-sample semantic score (n<=8): 0.1568
+Lock-in step (CER<=0.05): t=0
+Locked tokens: 79  Flexible tokens: 1
+TF-IDF vs attention stability corr: 0.9472
+TF-IDF status: OK
+
+Saved graphs:
+  - task2_attn_t*.png / task2_all_layers_t0.png
+  - task2_attn_evolution.png
+  - task2_semantic_drift.png
+  - task2_source_alignment.png
+  - task2_tfidf_vs_attention.png
+
+Step trajectory (first 10 rows)
+------------------------------------------------------------
+t=  3  bert=0.0603  drift=0.9397  text=ति ति ति रक्षि तः तः तः तः तः धर्मो धर्मो धर्मो धर्मो धर्मो 
+t=  2  bert=0.0597  drift=0.9403  text=ति ति ति रक्षि तः तः तः तः तः धर्मो धर्मो धर्मो धर्मो धर्मो 
+t=  1  bert=0.0597  drift=0.9403  text=ति ति ति रक्षि तः तः तः तः तः धर्मो धर्मो धर्मो धर्मो धर्मो 
+t=  0  bert=0.0597  drift=0.9403  text=ति ति ति रक्षि तः तः तः तः तः धर्मो धर्मो धर्मो धर्मो धर्मो 

analysis_reports/outputs_all_models_20260325/T4/task3_report.txt

@@ -0,0 +1,21 @@
+TASK 3 — CONCEPT VECTORS + PCA STEERING
+==================================================
+
+PCA: 50 components, 72.0% variance
+Diversity PC: 0  (|r|=-0.349 with diversity proxy)
+
+Direction validity: WEAK
+Spectrum unique ratio (mean over 5 seeds): 1.000
+Spectrum semantic stability (mean over 5 seeds): 0.325
+
+Saved graphs:
+  - task3_concept_space.png
+  - task3_pca_explained_variance.png
+  - task3_diversity_curve.png
+
+Diversity spectrum:
+  alpha=-2.0  →  बले र् अपश्य येहि ऌ वीर्य ऌ सिंहा सन सन ̍त̱ ज्ज्वा माम् वै वै महर्द्धि महर्द्धि ऌ सिंहा कू दिक्षु ऌ दश्य वै क्रमं बले र् दश्य स्वस्थ तुल तुल वीर्य वीर्य वी ऌ सिंहा राज कू वीर्य वीर्य वीर्य वीर्य ऌ वी निरुद्धा ̍त̱ बले बले साध्व उपशान्त वी वी दाक्षि हतः महर्द्धि साध्व तु वी वी ऌ दिक्षु दिक्षु पूष माम् पुरं ऌ दिक्षु वी पूष ̍त̱ ोद् दिक्षु पुरं स्त्रं मनोरथ अस्मा ऌ वाहि राजान वी
+  alpha=-1.0  →  बले बले अ तुल तुल वीर्य स्य सिंहा सन सन गतस्य गतस्य वै वै वै गतस्य सन पाता सिंहा दिता । ज्ज्वा वै वै बले बले र् अ अ तुल तुल वीर्य वीर्य स्य सिंहा सिंहा ध्रा स्य वीर्य वीर्य वीर्य तुल तुल ̍त̱ अ र् र् बले दिक्षु वै वै । वै संस्थिता रतं सन गतस्य पूष । वक्त्र सन सन सन सन सन व गतस्य व सन ॥ ति मनो हतः मातु ̍त̱ व कू कू सन सन
+  alpha=+0.0  →  बले र् अ तुल तुल वीर्य स्य सिंहा सन सन गतस्य गतस्य वै वै वै गतस्य सन सन सिंहा सिंहा । व वै वै त्वम् बले र् अ अ तुल तुल वीर्य वीर्य स्य स्य स्य सिंहा स्य स्य वीर्य वीर्य तुल तुल तुल अ र् र् र् त्ते वै वै गतस्य सन सन सन सन सन गतस्य निःसृ गतस्य सन गतस्य सन सन सन सन सन वि सन वि स्रव सिंहा सन सन सन सन सन सन सन गतस्य
+  alpha=+1.0  →  बले र् अ अ तुल वीर्य स्य सिंहा सन सन गतस्य गतस्य वै वै वै गतस्य सन सन सिंहा सिंहा षण् स्य ै वै बले बले र् अ अ तुल कान्ते षण् वीर्य स्य स्य सिंहा स्य स्य स्य वीर्य षण् वीर्य षण् अ अ र् र् र् षण् ेष गतस्य गतस्य गतस्य गतस्य सन सन षण् षण् गतस्य सन गतस्य गतस्य सन सन सन सन ष्णु गतस्य नो षण् नो - सन सन सन सन सन सन सन सन
+  alpha=+2.0  →  षण् र् अ षण् षण् षण् स्य षण् षण् षण् षण् षण् वै षण् षण् गतस्य षण् षण् षण् षण् षण् षण् षण् स षण् षण् र् षण् अ षण् षण् षण् षण् स्य स्य षण् स्य स्य स्य षण् षण् षण् षण् षण् अ र् षण् षण् षण् षण् षण् षण् षण् षण् षण् षण् षण् षण् षण् षण् षण् षण् षण् गतस्य षण् षण् षण् षण् षण् षण् षण् षण् षण् षण् षण् षण् षण् षण् षण् षण्

analysis_reports/outputs_all_models_20260325/T4/task4_report.txt

@@ -0,0 +1,14 @@
+TASK 4 — SEMANTIC ROBUSTNESS ABLATION
+==================================================
+
+Optimal diffusion steps = 4
+
+     T     BERT-F1     SEM_SIM      BLEU    sec/sample
+  --------------------------------------------------------
+     4      0.2644      0.0574    0.0000        0.2782
+
+Marginal gains (BERT-F1):
+
+Saved plots/files:
+  - task4_3d.png
+  - task4_raw_results.json

analysis_reports/outputs_all_models_20260325/T4/task5_report.txt

@@ -0,0 +1,15 @@
+TASK 5 — CLASSIFIER-FREE GUIDANCE
+==================================================
+
+Classifier params: 139521
+Training samples : 40
+
+Guidance scale sweep:
+       λ       CER   diversity      d2     sBLEU
+  ----------------------------------------------------
+     0.0    0.8366       0.815   0.635     0.005
+     0.5    0.8356       0.797   0.599     0.004 ← optimal
+     1.0    0.8369       0.791   0.588     0.006
+     1.5    0.8367       0.783   0.571     0.006
+     2.0    0.8367       0.774   0.553     0.005
+     3.0    0.8363       0.769   0.542     0.005

analysis_reports/outputs_all_models_20260325/T64/task1_kv_cache.txt

@@ -0,0 +1,15 @@
+TASK 1 — KV CACHE BENCHMARK
+========================================
+
+has_generate_cached=True
+memory_profile=Torch CPU mem-event reduction: 31.4% @ src_len=64 (std=8592.3MB, cache=5890.5MB)
+
+ src_len   standard(s)   cached(s)   speedup   encoder%
+      16         4.206       3.584     1.17x      74.4%
+      32         4.647       3.371     1.38x      37.6%
+      64         8.403       4.593     1.83x      49.6%
+
+Saved graphs:
+  - task1_time_comparison.png
+  - task1_speedup.png
+  - task1_encoder_cost.png

analysis_reports/outputs_all_models_20260325/T64/task2_report.txt

@@ -0,0 +1,35 @@
+TASK 2 — ATTENTION + DRIFT REPORT
+==================================================
+
+Input : dharmo rakṣati rakṣitaḥ
+Output: धर्मो रक्षति रक्षितः
+
+Captured steps: 64
+Analysis quality: VALID
+Final output uniq-ratio: 1.000
+Degenerate output: NO
+Multi-sample semantic score (n<=8): 0.1490
+Lock-in step (CER<=0.05): t=0
+Locked tokens: 59  Flexible tokens: 21
+TF-IDF vs attention stability corr: 0.7804
+TF-IDF status: OK
+
+Saved graphs:
+  - task2_attn_t*.png / task2_all_layers_t0.png
+  - task2_attn_evolution.png
+  - task2_semantic_drift.png
+  - task2_source_alignment.png
+  - task2_tfidf_vs_attention.png
+
+Step trajectory (first 10 rows)
+------------------------------------------------------------
+t= 63  bert=0.0552  drift=0.9448  text=धर्मो ति काम्य तः तः तः तः तः तः धर्मो धर्मो धर्मो धर्मो धर्
+t= 62  bert=0.0548  drift=0.9452  text=धर्मो ति काम्य तः तः तः तः तः तः धर्मो धर्मो धर्मो धर्मो धर्
+t= 61  bert=0.0548  drift=0.9452  text=धर्मो ति काम्य तः तः तः तः तः तः धर्मो धर्मो धर्मो धर्मो धर्
+t= 60  bert=0.0548  drift=0.9452  text=धर्मो ति काम्य तः तः तः तः तः तः धर्मो धर्मो धर्मो धर्मो धर्
+t= 59  bert=0.0548  drift=0.9452  text=धर्मो ति काम्य तः तः तः तः तः तः धर्मो धर्मो धर्मो धर्मो धर्
+t= 58  bert=0.0548  drift=0.9452  text=धर्मो ति काम्य तः तः तः तः तः तः धर्मो धर्मो धर्मो धर्मो धर्
+t= 57  bert=0.0548  drift=0.9452  text=धर्मो ति काम्य तः तः तः तः तः तः धर्मो धर्मो धर्मो धर्मो धर्
+t= 56  bert=0.0546  drift=0.9454  text=धर्मो ति काम्य तः तः तः तः तः तः धर्मो धर्मो धर्मो धर्मो धर्
+t= 55  bert=0.0546  drift=0.9454  text=धर्मो ति काम्य तः तः तः तः तः तः धर्मो धर्मो धर्मो धर्मो धर्
+t= 54  bert=0.0546  drift=0.9454  text=धर्मो ति काम्य तः तः तः तः तः तः धर्मो धर्मो धर्मो धर्मो धर्

analysis_reports/outputs_all_models_20260325/T64/task3_report.txt

@@ -0,0 +1,21 @@
+TASK 3 — CONCEPT VECTORS + PCA STEERING
+==================================================
+
+PCA: 39 components, 100.0% variance
+Diversity PC: 0  (|r|=0.314 with diversity proxy)
+
+Direction validity: WEAK
+Spectrum unique ratio (mean over 5 seeds): 1.000
+Spectrum semantic stability (mean over 5 seeds): 0.302
+
+Saved graphs:
+  - task3_concept_space.png
+  - task3_pca_explained_variance.png
+  - task3_diversity_curve.png
+
+Diversity spectrum:
+  alpha=-2.0  →  बले बले अ तुल तुल वीर्य अ̱स्या भूयः सन ान्ते लब्ध्वा अर्थ वै भूयः त्ति ान्त सन भूयः भूयः तुल अ अ वीरो बले बले बले अ̱स्या भूयः ॥ ॥ अ अ अ̱स्या ॥ वे̱ध सन सन सन सन ह् ह् ान ह् स्य ह् ानु ह् यो बो दादि ह् मतां ह् सान्त्व ह् ( मतां ॥ धीमान् भूयः ॥ अ̱स्या पान होमयेत् सारथि ( ॥ भूयः ॥ ॥ ॥ ॥ ॥ ॥ गोप लज्ज अ̱स्या मतां लज्ज यो
+  alpha=-1.0  →  बले बले अ तुल तुल वीर्य स्य सिंहा सिंहा सन त्ति वै वै वै द् सन सिंहा स्य स्य तुल तुल अ प्रयाति र् बले बले बले ॥ ॥ ॥ अ ॥ पि महा सन सन सन सन सन सन सिंहा स्रव सन स्य मीं गोप स्य स्य स्य स्य भूयः तुल सान्त्व यो अ ह् ान ान तव वेग ( यो भूषणम् ( ानु ॥ ॥ ॥ ॥ ॥ ॥ अ̱स्या ॥ ॥ ॥ ॥ यो पि ॥ म
+  alpha=+0.0  →  बले र् अ तुल तुल वीर्य स्य सिंहा सिंहा सन ध्या वै वै वै गतस्य भ सिंहा भ स्य तुल तुल अ अ र् बले बले बले ॥ बले र् । । ॥ वै वै सन सन सन सन सिंहा सिंहा सिंहा स्य स्य स्य स्य भ स्य स्य स्य स्य ानु स्य ् ता यो स्य फल ॥ म तुल च सि ॥ ् ॥ ॥ न् ॥ ॥ ॥ ॥ ॥ ॥ ॥ महा सन सन क्ष ॥
+  alpha=+1.0  →  बले र् । तुल यु वीर्य स्य सिंहा सिंहा सन ध्या वै वै वै । सन सिंहा स्य स्य तुल तुल । र् र् बले बले बले ॥ ॥ र् अ स् सन ते सन ीं सन सन त्र सिंहा यु सिंहा स्थल स्य स्य रौद्र स्य स्य न्दा ता यु स्य यु त्र क्ष ।। ीं स्य म्र कल्प यत् स् क्ष क्ष ॥ स्य यु मण्डलं यु ॥ ॥ ीं ॥ ॥ भ्यः ीं ीं ॥ ॥ ॥
+  alpha=+2.0  →  र् र् तुरङ्ग आहुः ितो । स्य सिंहा सिंहा सिंहा ब्रह्मा वै वै & ते तस् तुरङ्ग नो स्तम्भ ीं यु संच र् र् बले ीं स्तम्भ ते । तस् न्तं मण्डलं यु । स्तम्भ स्तम्भ सन आहुः सिंहा यु सिंहा सिंहा स्य मण्डलं यु स्य स्य स्य एव कल्प स्तम्भ ̱र स्तम्भ अमु आहुः यु ̱र कल्प यु तुरङ्ग यु तुरङ्ग ̱र तुरङ्ग रणम् मण्डलं यु ीं मण्डलं दिनं ̱र यु ॥ तुरङ्ग ितः आहुः ॥ मण्डलं आहुः क्षमः

analysis_reports/outputs_all_models_20260325/T64/task4_report.txt

@@ -0,0 +1,14 @@
+TASK 4 — SEMANTIC ROBUSTNESS ABLATION
+==================================================
+
+Optimal diffusion steps = 64
+
+     T     BERT-F1     SEM_SIM      BLEU    sec/sample
+  --------------------------------------------------------
+    64      0.2482      0.0580    0.0007        5.6116
+
+Marginal gains (BERT-F1):
+
+Saved plots/files:
+  - task4_3d.png
+  - task4_raw_results.json

analysis_reports/outputs_all_models_20260325/T64/task5_report.txt

@@ -0,0 +1,15 @@
+TASK 5 — CLASSIFIER-FREE GUIDANCE
+==================================================
+
+Classifier params: 139521
+Training samples : 20
+
+Guidance scale sweep:
+       λ       CER   diversity      d2     sBLEU
+  ----------------------------------------------------
+     0.0    0.8451       0.838   0.689     0.013 ← optimal
+     0.5    0.8490       0.818   0.650     0.013
+     1.0    0.8509       0.838   0.684     0.007
+     1.5    0.8622       0.857   0.720     0.005
+     2.0    0.8761       0.869   0.744     0.005
+     3.0    0.9056       0.814   0.642     0.013

analysis_reports/outputs_all_models_20260325/T8/task1_kv_cache.txt

@@ -0,0 +1,15 @@
+TASK 1 — KV CACHE BENCHMARK
+========================================
+
+has_generate_cached=True
+memory_profile=Torch CPU mem-event reduction: 25.8% @ src_len=64 (std=1168.9MB, cache=866.9MB)
+
+ src_len   standard(s)   cached(s)   speedup   encoder%
+      16         0.582       0.400     1.45x      35.9%
+      32         0.511       0.402     1.27x      37.7%
+      64         0.666       0.490     1.36x      35.6%
+
+Saved graphs:
+  - task1_time_comparison.png
+  - task1_speedup.png
+  - task1_encoder_cost.png

analysis_reports/outputs_all_models_20260325/T8/task2_report.txt

@@ -0,0 +1,33 @@
+TASK 2 — ATTENTION + DRIFT REPORT
+==================================================
+
+Input : dharmo rakṣati rakṣitaḥ
+Output: धर्मो रक्षति रक्षितः
+
+Captured steps: 8
+Analysis quality: WEAK
+Final output uniq-ratio: 1.000
+Degenerate output: NO
+Multi-sample semantic score (n<=8): 0.0915
+Lock-in step (CER<=0.05): t=0
+Locked tokens: 79  Flexible tokens: 1
+TF-IDF vs attention stability corr: 0.8905
+TF-IDF status: OK
+
+Saved graphs:
+  - task2_attn_t*.png / task2_all_layers_t0.png
+  - task2_attn_evolution.png
+  - task2_semantic_drift.png
+  - task2_source_alignment.png
+  - task2_tfidf_vs_attention.png
+
+Step trajectory (first 10 rows)
+------------------------------------------------------------
+t=  7  bert=0.0219  drift=0.9781  text=ितः ितः ितः ितः ितः ितः ितः ितः ितः ितः ितः ितः ितः ितः ितः 
+t=  6  bert=0.0225  drift=0.9775  text=ं ं ं ं ितः ितः ितः ितः ितः ितः ितः ितः ितः ितः ितः ितः ितः 
+t=  5  bert=0.0225  drift=0.9775  text=ं ं ं ं ितः ितः ितः ितः ितः ितः ितः ितः ितः ितः ितः ितः ितः 
+t=  4  bert=0.0225  drift=0.9775  text=ं ं ं ं ितः ितः ितः ितः ितः ितः ितः ितः ितः ितः ितः ितः ितः 
+t=  3  bert=0.0225  drift=0.9775  text=ं ं ं ं ितः ितः ितः ितः ितः ितः ितः ितः ितः ितः ितः ितः ितः 
+t=  2  bert=0.0227  drift=0.9773  text=ं ं ं ं ं ितः ितः ितः ितः ितः ितः ितः ितः ितः ितः ितः ितः ित
+t=  1  bert=0.0228  drift=0.9772  text=ं ं ं ं ं ितः ितः ितः ितः ितः ितः ितः ितः ितः ितः ितः ितः ित
+t=  0  bert=0.0228  drift=0.9772  text=ं ं ं ं ं ितः ितः ितः ितः ितः ितः ितः ितः ितः ितः ितः ितः ित

analysis_reports/outputs_all_models_20260325/T8/task3_report.txt

@@ -0,0 +1,21 @@
+TASK 3 — CONCEPT VECTORS + PCA STEERING
+==================================================
+
+PCA: 50 components, 75.9% variance
+Diversity PC: 0  (|r|=-0.344 with diversity proxy)
+
+Direction validity: WEAK
+Spectrum unique ratio (mean over 5 seeds): 1.000
+Spectrum semantic stability (mean over 5 seeds): 0.341
+
+Saved graphs:
+  - task3_concept_space.png
+  - task3_pca_explained_variance.png
+  - task3_diversity_curve.png
+
+Diversity spectrum:
+  alpha=-2.0  →  मनसः श्चक्र स्य स्य स्य अ स्य तैः तैः तैः स्य श्चक्र तैः गतभी स्य स्य स्य श्चक्र तैः तैः श्चक्र स्य तैः स्त्वं श्चक्र श्चक्र स्त्र तैः तैः कुण्ठ तैः तैः स्य तैः तैः तैः स्य तैः तैः गतभी तैः तैः णि̍ स्य तैः तैः तैः अ तैः तैः ह्वये मनसः ॥ तैः तैः गतभी ॥ श्चक्र तैः तैः तैः तैः तैः तैः तैः तैः स्य तैः करिष्या तैः स्त्वं तैः तैः श्चक्र तैः तैः श्चक्र तैः तैः ह्वये
+  alpha=-1.0  →  स्य अ तैः वै तैः अ वेद् मनसः स्य । । तैः तैः स्य गतभी स्य अ स्त्वं सीद् स्य तैः स्य तैः सु̱म् सीद् र्ध कृतानि गतभी गतभी तैः तैः स्य तैः तैः तैः मनसः तैः कृतानि तैः तैः सु̱म् अ तैः मनसः अ मनसः स्य अ तैः ॥ ॥ स्य गतभी गतभी ॥ ॥ वै तैः तैः मनसः तैः अ तैः तैः च वर स्य तैः या वात् स्य तैः सीद् तैः स्य तैः स्य अ तैः तैः
+  alpha=+0.0  →  अ अ वै ज्ञ स्य अ ज्ञ गतभी वर द शिख मन्त्र गतभी सु̱म् । द द स्य मन्त्र वा यो सीद् ज्ञ वै अ स्य स्य मन्त्र स्य मन्त्र स्य गतभी । गतभी गतभी तैः गतभी कृत तैः स्य तैः ॥ वै तैः ॥ वै अ कृतानि स्य वर वै ॥ ॥ वै ॥ अ ॥ स्य ॥ वै स्य ज्ञ ॥ स्य तैः तैः वै स्य स्य अ स्य तैः वै स्य तैः प्रण तीरे स्य । सीद्
+  alpha=+1.0  →  पम वै तुल्य शत्रू पम शिख वर अ णाः परा णाः स्य कृत प्रिय । भिन् णाः ज्ञ वै विराज वै गणो वै ्या अ वै पम ्या भिन् वै लब्ध शोभ स्य च श वर वै ॥ वै क्षिप्य शिख भिर् ॥ सन वा मन्त्र मृ ॥ ॥ ॥ वै ॥ मन्त्र ॥ पम सङ् वर शोभ क्षिप्य भिर् स्य क्षिप्य वै सन वर शिख वै शिख वर दर्श शिख कलं पम ौ वर कलं भिर् कलं वै शिख
+  alpha=+2.0  →  पम पम लब्ध पम शोभ पम परे भिर् अन्य णाः रसा लब्ध पम शोभ लब्ध शोभ पम शत्रू शिख भिन् पम पम पम णाः शोभ णाः पम शोभ शोभ परे णाः णाः पम पम पम पम शोभ पम शोभ शोभ पम डा णाः शोभ वै पम वै ॥ पम ॥ पम णाः ॥ ॥ परे अन्य णाः पम शोभ शोभ शोभ शोभ णाः वै परे कलं परे वै पम णाः पम शोभ णाः णाः कलं परे शोभ णाः कलं शत्रू

analysis_reports/outputs_all_models_20260325/T8/task4_report.txt

@@ -0,0 +1,14 @@
+TASK 4 — SEMANTIC ROBUSTNESS ABLATION
+==================================================
+
+Optimal diffusion steps = 8
+
+     T     BERT-F1     SEM_SIM      BLEU    sec/sample
+  --------------------------------------------------------
+     8      0.1210      0.0400    0.0000        0.6194
+
+Marginal gains (BERT-F1):
+
+Saved plots/files:
+  - task4_3d.png
+  - task4_raw_results.json

analysis_reports/outputs_all_models_20260325/T8/task5_report.txt

@@ -0,0 +1,15 @@
+TASK 5 — CLASSIFIER-FREE GUIDANCE
+==================================================
+
+Classifier params: 139521
+Training samples : 40
+
+Guidance scale sweep:
+       λ       CER   diversity      d2     sBLEU
+  ----------------------------------------------------
+     0.0    0.8834       0.796   0.596     0.004 ← optimal
+     0.5    0.8881       0.781   0.568     0.005
+     1.0    0.8876       0.767   0.540     0.007
+     1.5    0.8921       0.757   0.517     0.004
+     2.0    0.8929       0.734   0.474     0.006
+     3.0    0.8970       0.724   0.453     0.005

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,554 @@
+"""
+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 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
+
+
+# ── 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 _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'],
+    )
+    tgt_tok = SanskritTargetTokenizer(
+        vocab_size=cfg['model'].get('tgt_vocab_size', 16000),
+        max_len=cfg['model']['max_seq_len'],
+    )
+
+    print("\n" + "="*55)
+    print("Sanskrit D3PM Paraphrase — type verse, get paraphrase")
+    print("="*55 + "\n")
+
+    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
+
+        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
+
+
+# ── Batch evaluation ──────────────────────────────────────────────────
+
+def batch_evaluate(sample_size=500, checkpoint=None):
+    from data.dataset import OptimizedSanskritDataset
+    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']
+    exp_dir    = f"results/{model_name}_neg_{has_neg}"
+    ckpt       = checkpoint or f"{exp_dir}/best_model.pt"
+
+    if not os.path.exists(ckpt):
+        raise FileNotFoundError(f"No checkpoint at {ckpt}")
+
+    model, cfg = load_model(ckpt, cfg, device)
+    model.eval()
+
+    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'],
+    )
+
+    def collate(batch):
+        return {
+            'input_ids':   torch.stack([b['input_ids'].long() for b in batch]),
+            'target_text': [b['target_text'] for b in batch],
+            'input_text':  [b['input_text']  for b in batch],
+        }
+
+    dataset = OptimizedSanskritDataset(
+        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
+    )
+
+    all_preds, all_refs, all_inputs = [], [], []
+    print(f"⏳ Generating {len(indices)} paraphrases …")
+
+    for batch in tqdm(loader):
+        ids = batch['input_ids'].to(device)
+        out = run_inference(model, ids, cfg)
+        for i in range(out.size(0)):
+            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())
+
+    # Metrics
+    bleu_score, bert_f1 = 0.0, 0.0
+    try:
+        from nltk.translate.bleu_score import corpus_bleu
+        bleu_score = corpus_bleu(
+            [[r.split()] for r in all_refs],
+            [p.split() for p in all_preds]
+        )
+    except Exception:
+        pass
+
+    try:
+        import evaluate as hf_eval
+        res    = hf_eval.load('bertscore').compute(
+            predictions=all_preds, references=all_refs, lang='hi'
+        )
+        bert_f1 = sum(res['f1']) / len(res['f1'])
+    except Exception:
+        pass
+
+    # Save
+    out_path = f"{exp_dir}/evaluation_results.txt"
+    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")
+
+    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")
+
+    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
+
+
+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()
+
+    if args.mode == 'demo':
+        interactive_demo(checkpoint=args.checkpoint, single_text=args.text)
+    else:
+        batch_evaluate(args.samples, checkpoint=args.checkpoint)

inference_api.py

@@ -0,0 +1,131 @@
+import copy
+import json
+import os
+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 _read_model_settings() -> Dict[str, Any]:
+    if not os.path.exists("model_settings.json"):
+        return {}
+    try:
+        with open("model_settings.json", "r", encoding="utf-8") as f:
+            data = json.load(f)
+        return data if isinstance(data, dict) else {}
+    except Exception:
+        return {}
+
+
+def _load_once() -> None:
+    if _STATE["loaded"]:
+        return
+
+    settings = _read_model_settings()
+    cfg = copy.deepcopy(CONFIG)
+    cfg["model_type"] = os.environ.get(
+        "HF_MODEL_TYPE",
+        settings.get("model_type", "d3pm_cross_attention"),
+    )
+    cfg["data"]["include_negative_examples"] = (
+        os.environ.get(
+            "HF_INCLUDE_NEG",
+            str(settings.get("include_negative_examples", True)).lower(),
+        ).lower()
+        == "true"
+    )
+    num_steps_raw = os.environ.get("HF_NUM_STEPS", settings.get("num_steps"))
+    if num_steps_raw is not None:
+        num_steps = int(num_steps_raw)
+        cfg["model"]["diffusion_steps"] = num_steps
+        cfg["inference"]["num_steps"] = num_steps
+    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

model_settings.json

@@ -0,0 +1,5 @@
+{
+  "model_type": "d3pm_encoder_decoder",
+  "include_negative_examples": false,
+  "num_steps": 4
+}

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