decoding results

common.py

@@ -0,0 +1,837 @@
+# =============================================================================
+# common.py - shared utilities for ANLP Assignment 2
+#   * device / AMP / seeding helpers
+#   * Hugging Face Hub helpers (token from Kaggle secrets, push, download)
+#   * decoder-only Transformer with pluggable FFN (dense MLP or MoE)
+#   * KV-cache + batched greedy/sampled generation (uses model.forward only)
+# =============================================================================
+import os, math, json, time, random, contextlib, shutil, re, sys, glob, hashlib
+from dataclasses import dataclass, asdict, field, fields
+from typing import Optional, List, Dict
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+# -----------------------------------------------------------------------------
+# basic helpers
+# -----------------------------------------------------------------------------
+def seed_everything(seed: int):
+    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def get_device():
+    return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def amp_dtype(device):
+    """bf16 on Ampere+ (A100/H100/L4), fp16 on Kaggle's T4/P100, None on CPU."""
+    if device.type != "cuda":
+        return None
+    major, _ = torch.cuda.get_device_capability(device)
+    return torch.bfloat16 if major >= 8 else torch.float16
+
+
+def autocast_ctx(device, dtype):
+    if dtype is None:
+        return contextlib.nullcontext()
+    return torch.autocast(device_type="cuda", dtype=dtype)
+
+
+def make_grad_scaler(device, dtype):
+    enabled = (device.type == "cuda" and dtype == torch.float16)
+    try:
+        return torch.amp.GradScaler("cuda", enabled=enabled)
+    except Exception:  # older torch
+        return torch.cuda.amp.GradScaler(enabled=enabled)
+
+
+def fmt_num(n):
+    for unit, div in (("B", 1e9), ("M", 1e6), ("K", 1e3)):
+        if abs(n) >= div:
+            return f"{n/div:.2f}{unit}"
+    return str(n)
+
+
+def save_json(obj, path):
+    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
+    with open(path, "w", encoding="utf-8") as f:
+        json.dump(obj, f, indent=2, ensure_ascii=False, default=_json_default)
+
+
+def load_json(path):
+    with open(path, encoding="utf-8") as f:
+        return json.load(f)
+
+
+def _json_default(o):
+    if isinstance(o, (np.integer,)): return int(o)
+    if isinstance(o, (np.floating,)): return float(o)
+    if isinstance(o, np.ndarray): return o.tolist()
+    if isinstance(o, torch.Tensor): return o.tolist()
+    if isinstance(o, torch.dtype): return str(o)
+    return str(o)
+
+
+def gpu_name():
+    return torch.cuda.get_device_name(0) if torch.cuda.is_available() else "cpu"
+
+
+# -----------------------------------------------------------------------------
+# Hugging Face Hub helpers
+# -----------------------------------------------------------------------------
+TOKENIZER_REPO = "anlp-a2-tokenizers"      # holds mt_spm.model and lm_spm.model
+
+
+def get_hf_token():
+    """HF token from env var HF_TOKEN or Kaggle secret named HF_TOKEN."""
+    tok = os.environ.get("HF_TOKEN")
+    if tok:
+        return tok
+    try:
+        from kaggle_secrets import UserSecretsClient
+        tok = UserSecretsClient().get_secret("HF_TOKEN")
+        if tok:
+            os.environ["HF_TOKEN"] = tok
+            return tok
+    except Exception:
+        pass
+    return None
+
+
+def hf_login(token):
+    if not token:
+        print("[hf] no HF token found - add a Kaggle secret named HF_TOKEN (write access).")
+        return
+    try:
+        from huggingface_hub import login
+        login(token=token, add_to_git_credential=False)
+    except Exception as e:
+        print("[hf] login failed:", e)
+
+
+def hf_whoami(token):
+    from huggingface_hub import HfApi
+    return HfApi(token=token).whoami()["name"]
+
+
+def hf_resolve_user(cfg_user, token):
+    if cfg_user:
+        return cfg_user
+    if token:
+        try:
+            return hf_whoami(token)
+        except Exception as e:
+            print("[hf] whoami failed:", e)
+    return None
+
+
+def _retry(fn, tries=4, wait=10):
+    last = None
+    for i in range(tries):
+        try:
+            return fn()
+        except Exception as e:
+            last = e
+            print(f"[hf] attempt {i+1}/{tries} failed: {e}")
+            time.sleep(wait * (i + 1))
+    raise last
+
+
+def hf_create_repo(repo_id, token, private=False, repo_type="model"):
+    from huggingface_hub import HfApi
+    api = HfApi(token=token)
+    _retry(lambda: api.create_repo(repo_id, private=private, exist_ok=True, repo_type=repo_type))
+    return api
+
+
+def hf_push_folder(folder, repo_id, token, private=False, message="upload",
+                   path_in_repo=None, run_as_future=False, repo_type="model"):
+    api = hf_create_repo(repo_id, token, private, repo_type)
+    kw = dict(folder_path=folder, repo_id=repo_id, repo_type=repo_type,
+              commit_message=message)
+    if path_in_repo:
+        kw["path_in_repo"] = path_in_repo
+    if run_as_future:
+        return api.upload_folder(run_as_future=True, **kw)
+    return _retry(lambda: api.upload_folder(**kw))
+
+
+def hf_push_file(local_path, path_in_repo, repo_id, token, private=False,
+                 message="upload", run_as_future=False, repo_type="model"):
+    api = hf_create_repo(repo_id, token, private, repo_type)
+    kw = dict(path_or_fileobj=local_path, path_in_repo=path_in_repo, repo_id=repo_id,
+              repo_type=repo_type, commit_message=message)
+    if run_as_future:
+        return api.upload_file(run_as_future=True, **kw)
+    return _retry(lambda: api.upload_file(**kw))
+
+
+def hf_try_download(repo_id, filename, token=None, repo_type="model", force=False):
+    """Returns a local path, or None if the repo/file does not exist."""
+    try:
+        from huggingface_hub import hf_hub_download
+        return hf_hub_download(repo_id=repo_id, filename=filename, token=token,
+                               repo_type=repo_type, force_download=force)
+    except Exception as e:
+        print(f"[hf] could not download {repo_id}/{filename}: {type(e).__name__}")
+        return None
+
+
+def hf_list_files(repo_id, token=None, repo_type="model"):
+    try:
+        from huggingface_hub import HfApi
+        return HfApi(token=token).list_repo_files(repo_id, repo_type=repo_type)
+    except Exception:
+        return []
+
+
+# -----------------------------------------------------------------------------
+# Weights & Biases (assignment: upload training logs to WandB, public links in README/report)
+# Fail-soft: if wandb is missing / not logged in, training continues and logs only to JSON.
+# Auth: `wandb login` or WANDB_API_KEY (Kaggle secret of that name is picked up too).
+# WANDB_MODE=offline logs locally (sync later with `wandb sync`); cfg["wandb"]=False disables.
+# -----------------------------------------------------------------------------
+WANDB_PROJECT = "anlp-a2"
+
+
+def _wandb_key():
+    key = os.environ.get("WANDB_API_KEY")
+    if key:
+        return key
+    try:
+        from kaggle_secrets import UserSecretsClient
+        key = UserSecretsClient().get_secret("WANDB_API_KEY")
+        if key:
+            os.environ["WANDB_API_KEY"] = key
+        return key
+    except Exception:
+        return None
+
+
+def wandb_init(cfg, name, config, tags=(), group=None, notes=None, stable_id=False):
+    """Start (or resume, with the same stable id -> requeued jobs continue one run) a W&B run.
+    Returns the run, or None when disabled / unavailable."""
+    if not cfg.get("wandb", True):
+        return None
+    try:
+        import wandb
+        _wandb_key()
+        run = wandb.init(project=cfg.get("wandb_project") or WANDB_PROJECT, entity=cfg.get("wandb_entity"),
+                         name=name, config=config,
+                         **(dict(id=re.sub(r"[^A-Za-z0-9_-]", "-", name), resume="allow") if stable_id else {}),
+                         tags=list(tags), group=group, notes=notes,
+                         settings=wandb.Settings(init_timeout=120, console="off"))
+        print("[wandb] run:", getattr(run, "url", None) or "(offline - `wandb sync <dir>` later)")
+        return run
+    except Exception as e:
+        print(f"[wandb] disabled ({type(e).__name__}: {str(e)[:120]}) - logging to JSON only")
+        return None
+
+
+def wandb_log(run, data, step=None):
+    if run is None:
+        return
+    try:
+        run.log({k: v for k, v in data.items() if v is not None}, step=step)
+    except Exception as e:
+        print("[wandb] log failed:", e)
+
+
+def wandb_images(run, paths, step=None):
+    """paths: {key: png_path}."""
+    if run is None:
+        return
+    try:
+        import wandb
+        wandb_log(run, {k: wandb.Image(v) for k, v in paths.items() if os.path.exists(v)}, step)
+    except Exception as e:
+        print("[wandb] image log failed:", e)
+
+
+def wandb_finish(run, summary=None):
+    if run is None:
+        return
+    try:
+        for k, v in (summary or {}).items():
+            run.summary[k] = v
+        run.finish()
+    except Exception as e:
+        print("[wandb] finish failed:", e)
+
+
+def expert_load_scalars(prefix, loads):
+    """loads: [layer][expert] fractions -> flat {prefix/L{l}_E{e}: x} for W&B."""
+    return {f"{prefix}/L{l}_E{e}": float(x) for l, row in enumerate(loads) for e, x in enumerate(row)}
+
+
+# -----------------------------------------------------------------------------
+# Model
+# -----------------------------------------------------------------------------
+@dataclass
+class ModelConfig:
+    vocab_size: int = 32000
+    block_size: int = 256
+    n_layer: int = 6
+    n_head: int = 8
+    n_embd: int = 512
+    dropout: float = 0.0
+    bias: bool = False
+    # ---- feed-forward ----
+    ffn_type: str = "dense"        # "dense" | "moe"
+    mlp_hidden: int = 0            # dense hidden size (0 -> 4 * n_embd)
+    n_experts: int = 4             # number of *routed* experts
+    n_shared: int = 0              # number of always-on shared experts
+    top_k: int = 1                 # routed experts active per token
+    expert_hidden: int = 0         # hidden size of each expert (routed & shared)
+    aux_loss_coef: float = 0.01    # Switch-style load-balancing loss
+    z_loss_coef: float = 1e-3      # router z-loss (ST-MoE)
+
+    def to_dict(self):
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, d):
+        names = {f.name for f in fields(cls)}
+        return cls(**{k: v for k, v in d.items() if k in names})
+
+
+# The five FFN variants of Part 1.  H = 4 * d is the dense hidden size.
+#   dense            : 1 MLP (d -> H -> d)                  total = active = 2dH
+#   moe_4e_top1      : 4 experts of hidden H/4, top-1         total = 2dH, active = 2dH/4
+#   moe_4e_top2      : 4 experts of hidden H/4, top-2         total = 2dH, active = 2dH/2
+#   moe_1s3r_top1    : 1 shared + 3 routed (hidden H/4), top-1 total = 2dH, active = 2dH/2
+#   moe_4e_top2_active: 4 experts of hidden H/2, top-2         total = 4dH, active = 2dH (= dense)
+FFN_VARIANTS = ["dense", "moe_4e_top1", "moe_4e_top2", "moe_1s3r_top1", "moe_4e_top2_active"]
+
+FFN_VARIANT_DESCRIPTIONS = {
+    "dense": "Standard 2-layer MLP (d -> 4d -> d)",
+    "moe_4e_top1": "MoE: 4 experts (hidden d), top-1 routing - same total params as dense",
+    "moe_4e_top2": "MoE: 4 experts (hidden d), top-2 routing - same total params as dense",
+    "moe_1s3r_top1": "MoE: 1 shared + 3 routed experts (hidden d), top-1 routed - same total params as dense",
+    "moe_4e_top2_active": "MoE: 4 experts (hidden 2d), top-2 routing - same ACTIVE params as dense",
+}
+
+
+def ffn_variant_kwargs(name, n_embd):
+    H = 4 * n_embd
+    if name == "dense":
+        return dict(ffn_type="dense", mlp_hidden=H)
+    if name == "moe_4e_top1":
+        return dict(ffn_type="moe", n_experts=4, n_shared=0, top_k=1, expert_hidden=H // 4)
+    if name == "moe_4e_top2":
+        return dict(ffn_type="moe", n_experts=4, n_shared=0, top_k=2, expert_hidden=H // 4)
+    if name == "moe_1s3r_top1":
+        return dict(ffn_type="moe", n_experts=3, n_shared=1, top_k=1, expert_hidden=H // 4)
+    if name == "moe_4e_top2_active":
+        return dict(ffn_type="moe", n_experts=4, n_shared=0, top_k=2, expert_hidden=H // 2)
+    raise ValueError(f"unknown FFN variant {name}")
+
+
+class LayerNorm(nn.Module):
+    def __init__(self, ndim, bias):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(ndim))
+        self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None
+
+    def forward(self, x):
+        return F.layer_norm(x, self.weight.shape, self.weight, self.bias, 1e-5)
+
+
+class KVCache:
+    """Simple per-layer key/value cache. Tensors are (B, n_head, T, head_dim)."""
+    def __init__(self, n_layer):
+        self.k = [None] * n_layer
+        self.v = [None] * n_layer
+
+    @property
+    def length(self):
+        return 0 if self.k[0] is None else self.k[0].shape[2]
+
+    def update(self, i, k, v):
+        if self.k[i] is None:
+            self.k[i], self.v[i] = k, v
+        else:
+            self.k[i] = torch.cat([self.k[i], k], dim=2)
+            self.v[i] = torch.cat([self.v[i], v], dim=2)
+        return self.k[i], self.v[i]
+
+    def reorder(self, idx):
+        for i in range(len(self.k)):
+            if self.k[i] is not None:
+                self.k[i] = self.k[i].index_select(0, idx)
+                self.v[i] = self.v[i].index_select(0, idx)
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, cfg: ModelConfig):
+        super().__init__()
+        assert cfg.n_embd % cfg.n_head == 0
+        self.c_attn = nn.Linear(cfg.n_embd, 3 * cfg.n_embd, bias=cfg.bias)
+        self.c_proj = nn.Linear(cfg.n_embd, cfg.n_embd, bias=cfg.bias)
+        self.resid_dropout = nn.Dropout(cfg.dropout)
+        self.n_head, self.n_embd, self.dropout = cfg.n_head, cfg.n_embd, cfg.dropout
+
+    def forward(self, x, attn_mask=None, cache=None, layer_idx=0):
+        B, T, C = x.size()
+        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
+        hs = C // self.n_head
+        q = q.view(B, T, self.n_head, hs).transpose(1, 2)
+        k = k.view(B, T, self.n_head, hs).transpose(1, 2)
+        v = v.view(B, T, self.n_head, hs).transpose(1, 2)
+        if cache is not None:
+            k, v = cache.update(layer_idx, k, v)
+        dp = self.dropout if self.training else 0.0
+        if attn_mask is None:
+            # plain causal attention; valid only when queries and keys are aligned
+            assert T == 1 or k.shape[2] == T, "pass an explicit attn_mask when prefilling a non-empty cache"
+            y = F.scaled_dot_product_attention(q, k, v, dropout_p=dp, is_causal=(T > 1))
+        else:
+            y = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=dp)
+        y = y.transpose(1, 2).contiguous().view(B, T, C)
+        return self.resid_dropout(self.c_proj(y))
+
+
+class MLP(nn.Module):
+    """Standard 2-layer GELU MLP: d -> hidden -> d."""
+    def __init__(self, n_embd, hidden, bias=False, dropout=0.0):
+        super().__init__()
+        self.c_fc = nn.Linear(n_embd, hidden, bias=bias)
+        self.gelu = nn.GELU()
+        self.c_proj = nn.Linear(hidden, n_embd, bias=bias)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        return self.dropout(self.c_proj(self.gelu(self.c_fc(x))))
+
+
+class MoE(nn.Module):
+    """Token-choice top-k Mixture-of-Experts, drop-in replacement for MLP.
+
+    * router: linear d -> n_experts, softmax over routed experts
+    * top-1: output = p_e * E_e(x)             (Switch Transformer; keeps router gradient)
+    * top-k (k>1): gates renormalised over the chosen k (Mixtral-style)
+    * shared experts (DeepSeek-MoE style) are applied to every token and added
+    * aux loss = coef * E * sum_e f_e * P_e  (+ router z-loss), computed on real tokens only
+    * optional tracking of which experts are chosen, per token group (language / segment)
+    """
+    def __init__(self, cfg: ModelConfig):
+        super().__init__()
+        h = cfg.expert_hidden or (4 * cfg.n_embd // max(cfg.n_experts + cfg.n_shared, 1))
+        self.n_experts, self.top_k, self.n_shared = cfg.n_experts, cfg.top_k, cfg.n_shared
+        self.aux_coef, self.z_coef = cfg.aux_loss_coef, cfg.z_loss_coef
+        self.experts = nn.ModuleList([MLP(cfg.n_embd, h, cfg.bias, cfg.dropout) for _ in range(cfg.n_experts)])
+        self.shared = nn.ModuleList([MLP(cfg.n_embd, h, cfg.bias, cfg.dropout) for _ in range(cfg.n_shared)])
+        self.router = nn.Linear(cfg.n_embd, cfg.n_experts, bias=False)
+        self.aux_loss = torch.zeros(())
+        # tracking state
+        self.track = False
+        self.n_groups = 0
+        self.counts = None        # (n_groups, n_experts) long
+        self.token_groups = None  # (N,) long, -1 = ignore
+        self.load_accum = None    # (n_experts,) running fraction for training logs
+        self.load_steps = 0
+
+    def reset_tracking(self, n_groups):
+        self.n_groups = n_groups
+        self.counts = torch.zeros(n_groups, self.n_experts, dtype=torch.long, device=self.router.weight.device)
+
+    def forward(self, x, tok_mask=None):
+        B, T, C = x.shape
+        xf = x.reshape(-1, C)
+        N = xf.shape[0]
+        if tok_mask is not None:
+            sel = tok_mask.reshape(-1).nonzero(as_tuple=True)[0]
+            xs = xf.index_select(0, sel)
+        else:
+            sel, xs = None, xf
+
+        logits = self.router(xs).float()                        # (n, E)
+        probs = F.softmax(logits, dim=-1)
+        topv, topi = probs.topk(self.top_k, dim=-1)             # (n, k)
+        if self.top_k > 1:
+            topv = topv / topv.sum(-1, keepdim=True)
+
+        out = torch.zeros(xs.shape, dtype=torch.float32, device=xs.device)
+        for e, expert in enumerate(self.experts):
+            tok, slot = (topi == e).nonzero(as_tuple=True)
+            if tok.numel() == 0:
+                continue
+            y = expert(xs.index_select(0, tok)).float() * topv[tok, slot].unsqueeze(-1)
+            out.index_add_(0, tok, y)
+        for s in self.shared:
+            out = out + s(xs).float()
+
+        if self.training and xs.shape[0] > 0:
+            E = self.n_experts
+            f = F.one_hot(topi, E).sum(1).float().mean(0) / self.top_k   # fraction of assignments
+            P = probs.mean(0)                                            # mean router prob
+            aux = E * (f * P).sum()
+            z = torch.logsumexp(logits, dim=-1).pow(2).mean()
+            self.aux_loss = self.aux_coef * aux + self.z_coef * z
+            with torch.no_grad():
+                self.load_accum = f.detach() if self.load_accum is None else self.load_accum + f.detach()
+                self.load_steps += 1
+        else:
+            self.aux_loss = torch.zeros((), device=x.device)
+
+        if self.track and self.token_groups is not None:
+            with torch.no_grad():
+                g = self.token_groups.reshape(-1)
+                if sel is not None:
+                    g = g.index_select(0, sel)
+                g = g.unsqueeze(1).expand_as(topi)
+                ok = g >= 0
+                flat = (g[ok] * self.n_experts + topi[ok]).reshape(-1)
+                self.counts += torch.bincount(flat, minlength=self.n_groups * self.n_experts).view(
+                    self.n_groups, self.n_experts)
+
+        out = out.to(x.dtype)
+        if sel is not None:
+            full = torch.zeros_like(xf)
+            full.index_copy_(0, sel, out)
+            out = full
+        return out.view(B, T, C)
+
+
+class Block(nn.Module):
+    def __init__(self, cfg: ModelConfig):
+        super().__init__()
+        self.ln_1 = LayerNorm(cfg.n_embd, cfg.bias)
+        self.attn = CausalSelfAttention(cfg)
+        self.ln_2 = LayerNorm(cfg.n_embd, cfg.bias)
+        if cfg.ffn_type == "moe":
+            self.ffn = MoE(cfg)
+        else:
+            self.ffn = MLP(cfg.n_embd, cfg.mlp_hidden or 4 * cfg.n_embd, cfg.bias, cfg.dropout)
+        self.is_moe = cfg.ffn_type == "moe"
+
+    def forward(self, x, attn_mask=None, cache=None, layer_idx=0, tok_mask=None):
+        x = x + self.attn(self.ln_1(x), attn_mask, cache, layer_idx)
+        h = self.ln_2(x)
+        x = x + (self.ffn(h, tok_mask) if self.is_moe else self.ffn(h))
+        return x
+
+
+class Transformer(nn.Module):
+    def __init__(self, cfg: ModelConfig):
+        super().__init__()
+        self.config = cfg
+        self.transformer = nn.ModuleDict(dict(
+            wte=nn.Embedding(cfg.vocab_size, cfg.n_embd),
+            wpe=nn.Embedding(cfg.block_size, cfg.n_embd),
+            drop=nn.Dropout(cfg.dropout),
+            h=nn.ModuleList([Block(cfg) for _ in range(cfg.n_layer)]),
+            ln_f=LayerNorm(cfg.n_embd, cfg.bias),
+        ))
+        self.lm_head = nn.Linear(cfg.n_embd, cfg.vocab_size, bias=False)
+        self.transformer.wte.weight = self.lm_head.weight   # weight tying
+        self.apply(self._init_weights)
+        for pn, p in self.named_parameters():
+            if pn.endswith("c_proj.weight"):
+                torch.nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * cfg.n_layer))
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            torch.nn.init.normal_(m.weight, mean=0.0, std=0.02)
+            if m.bias is not None:
+                torch.nn.init.zeros_(m.bias)
+        elif isinstance(m, nn.Embedding):
+            torch.nn.init.normal_(m.weight, mean=0.0, std=0.02)
+
+    # ---- MoE helpers ----
+    def moe_layers(self):
+        return [b.ffn for b in self.transformer.h if isinstance(b.ffn, MoE)]
+
+    # ---- forward ----
+    def hidden_states(self, idx, pos=None, attn_mask=None, cache=None, tok_mask=None):
+        B, T = idx.shape
+        if pos is None:
+            start = cache.length if cache is not None else 0
+            assert start + T <= self.config.block_size, f"sequence length {start+T} > block_size"
+            pos = torch.arange(start, start + T, device=idx.device).unsqueeze(0)
+        x = self.transformer.drop(self.transformer.wte(idx) + self.transformer.wpe(pos))
+        for i, block in enumerate(self.transformer.h):
+            x = block(x, attn_mask, cache, i, tok_mask)
+        return self.transformer.ln_f(x)
+
+    def aux_loss(self):
+        layers = self.moe_layers()
+        if not layers:
+            return None
+        return sum(m.aux_loss for m in layers)
+
+    def forward(self, idx, targets=None, pos=None, attn_mask=None, cache=None,
+                tok_mask=None, last_only=False):
+        """targets: (B,T) with -1 = ignore. Returns (logits, loss, aux_loss).
+        When targets are given, logits are only computed on supervised positions
+        (saves the big vocab projection on source / padding tokens) and None is returned."""
+        x = self.hidden_states(idx, pos, attn_mask, cache, tok_mask)
+        aux = self.aux_loss()
+        if targets is not None:
+            valid = targets != -1
+            logits = self.lm_head(x[valid])
+            loss = F.cross_entropy(logits.float(), targets[valid])
+            return None, loss, aux
+        if last_only:
+            x = x[:, -1:, :]
+        return self.lm_head(x), None, aux
+
+    @torch.no_grad()
+    def token_nll(self, idx, targets, tok_mask=None):
+        """Per-token NLL (B,T), 0 where targets == -1."""
+        x = self.hidden_states(idx, tok_mask=tok_mask)
+        valid = targets != -1
+        out = torch.zeros(targets.shape, dtype=torch.float32, device=idx.device)
+        logits = self.lm_head(x[valid]).float()
+        out[valid] = F.cross_entropy(logits, targets[valid], reduction="none")
+        return out
+
+
+# ---- parameter accounting ----
+def _numel(m):
+    return sum(p.numel() for p in m.parameters())
+
+
+def param_report(model: Transformer):
+    cfg = model.config
+    seen, total = set(), 0
+    for p in model.parameters():
+        if id(p) not in seen:
+            seen.add(id(p)); total += p.numel()
+    emb = model.transformer.wte.weight.numel() + model.transformer.wpe.weight.numel()
+    ffn_total = ffn_active = router = 0
+    for b in model.transformer.h:
+        f = b.ffn
+        if isinstance(f, MoE):
+            per_exp = _numel(f.experts[0])
+            r = _numel(f.router)
+            sh = sum(_numel(s) for s in f.shared)
+            ffn_total += per_exp * f.n_experts + sh + r
+            ffn_active += per_exp * f.top_k + sh + r
+            router += r
+        else:
+            ffn_total += _numel(f); ffn_active += _numel(f)
+    non_ffn = total - ffn_total
+    return {
+        "total_params": total,
+        "active_params_per_token": non_ffn + ffn_active,
+        "embedding_params": emb,
+        "total_non_embedding": total - emb,
+        "active_non_embedding": non_ffn + ffn_active - emb,
+        "ffn_total_params": ffn_total,
+        "ffn_active_params": ffn_active,
+        "router_params": router,
+    }
+
+
+def print_param_report(rep, title=""):
+    print(f"--- parameters {title} ---")
+    for k, v in rep.items():
+        print(f"  {k:26s} {v:>12,d}  ({fmt_num(v)})")
+
+
+# ---- save / load (safetensors handles the tied embedding) ----
+def save_model_dir(model, out_dir, extra_config=None):
+    os.makedirs(out_dir, exist_ok=True)
+    try:
+        from safetensors.torch import save_model
+        save_model(model, os.path.join(out_dir, "model.safetensors"))
+    except Exception as e:
+        print("[save] safetensors failed, falling back to torch.save:", e)
+        torch.save(model.state_dict(), os.path.join(out_dir, "pytorch_model.bin"))
+    cfg = model.config.to_dict()
+    if extra_config:
+        cfg.update(extra_config)
+    save_json(cfg, os.path.join(out_dir, "config.json"))
+
+
+def load_model_dir(path_or_repo, token=None, device="cpu"):
+    """Load a model saved by save_model_dir from a local dir or a HF repo id."""
+    if os.path.isdir(path_or_repo):
+        d = path_or_repo
+        cfg_path = os.path.join(d, "config.json")
+        st = os.path.join(d, "model.safetensors")
+    else:
+        cfg_path = hf_try_download(path_or_repo, "config.json", token)
+        st = hf_try_download(path_or_repo, "model.safetensors", token)
+    cfg = ModelConfig.from_dict(load_json(cfg_path))
+    model = Transformer(cfg)
+    from safetensors.torch import load_model
+    load_model(model, st, device="cpu")
+    return model.to(device)
+
+
+# -----------------------------------------------------------------------------
+# Generation (KV cache, left padding, batched). Uses model.forward only.
+# -----------------------------------------------------------------------------
+@torch.no_grad()
+def generate(model, prompts: List[List[int]], max_new_tokens: int, eos_id: Optional[int],
+             pad_id: int = 0, temperature: float = 0.0, top_k: Optional[int] = None,
+             device=None, amp=None, generator=None):
+    """Batched autoregressive generation. temperature=0 -> greedy.
+    Returns list of generated token lists (EOS stripped)."""
+    device = device or next(model.parameters()).device
+    model.eval()
+    B = len(prompts)
+    L = max(len(p) for p in prompts)
+    max_new_tokens = min(max_new_tokens, model.config.block_size - L)
+    idx = torch.full((B, L), pad_id, dtype=torch.long)
+    valid = torch.zeros((B, L), dtype=torch.bool)
+    for i, p in enumerate(prompts):
+        idx[i, L - len(p):] = torch.tensor(p, dtype=torch.long)
+        valid[i, L - len(p):] = True
+    idx, valid = idx.to(device), valid.to(device)
+    pos = (valid.long().cumsum(1) - 1).clamp(min=0)
+    causal = torch.tril(torch.ones(L, L, dtype=torch.bool, device=device))
+    eye = torch.eye(L, dtype=torch.bool, device=device)
+    # key must be valid (not left padding); always allow self-attention so pad rows never go NaN
+    mask = (causal[None, None] & valid[:, None, None, :]) | eye[None, None]
+    cache = KVCache(model.config.n_layer)
+    with autocast_ctx(device, amp):
+        logits, _, _ = model(idx, pos=pos, attn_mask=mask, cache=cache, last_only=True)
+    next_pos = pos[:, -1] + 1
+    finished = torch.zeros(B, dtype=torch.bool, device=device)
+    outs = []
+    for _ in range(max_new_tokens):
+        lg = logits[:, -1, :].float()
+        if temperature and temperature > 0:
+            lg = lg / temperature
+            if top_k:
+                v, _ = torch.topk(lg, min(top_k, lg.size(-1)))
+                lg[lg < v[:, [-1]]] = -float("inf")
+            nxt = torch.multinomial(F.softmax(lg, -1), 1, generator=generator).squeeze(1)
+        else:
+            nxt = lg.argmax(-1)
+        nxt = torch.where(finished, torch.full_like(nxt, pad_id), nxt)
+        outs.append(nxt)
+        if eos_id is not None:
+            finished |= nxt == eos_id
+        if bool(finished.all()):
+            break
+        valid = torch.cat([valid, torch.ones(B, 1, dtype=torch.bool, device=device)], 1)
+        with autocast_ctx(device, amp):
+            logits, _, _ = model(nxt[:, None], pos=next_pos[:, None], attn_mask=valid[:, None, None, :],
+                                 cache=cache, last_only=True)
+        next_pos = next_pos + 1
+    if not outs:
+        return [[] for _ in range(B)]
+    gen = torch.stack(outs, 1).tolist()
+    res = []
+    for row in gen:
+        if eos_id is not None and eos_id in row:
+            row = row[:row.index(eos_id)]
+        res.append(row)
+    return res
+
+
+@torch.no_grad()
+def generate_sorted(model, prompts, max_new_tokens, eos_id, pad_id=0, batch_size=64,
+                    device=None, amp=None, progress=False, **kw):
+    """Sort prompts by length for efficient batching, restore the original order."""
+    order = sorted(range(len(prompts)), key=lambda i: len(prompts[i]))
+    out = [None] * len(prompts)
+    t0 = time.time()
+    for bi, s in enumerate(range(0, len(order), batch_size)):
+        ids = order[s:s + batch_size]
+        res = generate(model, [prompts[i] for i in ids], max_new_tokens, eos_id, pad_id,
+                       device=device, amp=amp, **kw)
+        for i, r in zip(ids, res):
+            out[i] = r
+        if progress and bi % 20 == 0:
+            print(f"  generated {min(s+batch_size, len(order))}/{len(order)}  ({time.time()-t0:.0f}s)")
+    return out
+
+
+# -----------------------------------------------------------------------------
+# HAP-E split (Part 2) - shared by the tokenizer notebook and the LM runs
+# -----------------------------------------------------------------------------
+def hape_split(doc_ids, seed=0, val_frac=0.05, test_frac=0.05):
+    """Deterministic split by base document id. Returns dict base -> 'train'|'val'|'test'."""
+    bases = sorted({d.split("@")[0] for d in doc_ids})
+    rng = np.random.default_rng(seed)
+    perm = rng.permutation(len(bases))
+    n_val, n_test = int(len(bases) * val_frac), int(len(bases) * test_frac)
+    split = {}
+    for rank, i in enumerate(perm):
+        split[bases[i]] = "val" if rank < n_val else ("test" if rank < n_val + n_test else "train")
+    return split
+
+
+# -----------------------------------------------------------------------------
+# LR schedule
+# -----------------------------------------------------------------------------
+def lr_factor(step, total_steps, warmup_steps, min_ratio=0.1):
+    """Linear warmup then cosine decay to min_ratio."""
+    if step < warmup_steps:
+        return (step + 1) / max(1, warmup_steps)
+    prog = (step - warmup_steps) / max(1, total_steps - warmup_steps)
+    prog = min(max(prog, 0.0), 1.0)
+    return min_ratio + (1 - min_ratio) * 0.5 * (1 + math.cos(math.pi * prog))
+
+
+def set_lr(optimizer, factor):
+    for g in optimizer.param_groups:
+        if "base_lr" not in g:
+            g["base_lr"] = g["lr"]
+        g["lr"] = g["base_lr"] * factor
+
+
+# -----------------------------------------------------------------------------
+# run several jobs at once, one per GPU (e.g. Kaggle "GPU T4 x2")
+# -----------------------------------------------------------------------------
+def launch_parallel(jobs, poll_s=180, tail=3):
+    """jobs: list of (name, python_code). Job i runs in its own process on GPU i.
+    Output goes to <name>.log; the last lines of every log are printed every poll_s seconds."""
+    import subprocess
+    get_hf_token()                       # puts HF_TOKEN into os.environ for the children
+    n_gpu = max(torch.cuda.device_count(), 1)
+    assert len(jobs) <= n_gpu, f"{len(jobs)} jobs but only {n_gpu} GPU(s)"
+    procs = []
+    for gpu, (name, code) in enumerate(jobs):
+        env = dict(os.environ, CUDA_VISIBLE_DEVICES=str(gpu), PYTHONUNBUFFERED="1")
+        logf = open(f"{name}.log", "w")
+        procs.append((name, subprocess.Popen([sys.executable, "-u", "-c", code], env=env,
+                                             stdout=logf, stderr=subprocess.STDOUT), logf))
+        print(f"started {name} on GPU {gpu} (log: {name}.log)")
+    while True:
+        time.sleep(poll_s)
+        alive = False
+        for name, p, _ in procs:
+            alive |= p.poll() is None
+            try:
+                lines = open(f"{name}.log").read().splitlines()[-tail:]
+            except Exception:
+                lines = []
+            print(f"--- {name} ({'running' if p.poll() is None else 'exit ' + str(p.returncode)}) ---")
+            for l in lines:
+                print("   ", l[:200])
+        if not alive:
+            break
+    for _, _, f in procs:
+        f.close()
+    codes = {name: p.returncode for name, p, _ in procs}
+    print("exit codes:", codes)
+    return codes
+
+
+# -----------------------------------------------------------------------------
+# Plot style (validated categorical palette, fixed order)
+# -----------------------------------------------------------------------------
+PALETTE = ["#2a78d6", "#eb6834", "#1baf7a", "#eda100", "#e87ba4", "#008300", "#4a3aa7", "#e34948"]
+
+
+def setup_plot_style():
+    import matplotlib as mpl
+    mpl.rcParams.update({
+        "figure.dpi": 110, "savefig.dpi": 150, "savefig.bbox": "tight",
+        "axes.spines.top": False, "axes.spines.right": False,
+        "axes.grid": True, "grid.color": "#e4e3df", "grid.linewidth": 0.8,
+        "axes.edgecolor": "#8a8984", "axes.labelcolor": "#2b2a27",
+        "xtick.color": "#52514e", "ytick.color": "#52514e",
+        "axes.prop_cycle": mpl.cycler(color=PALETTE),
+        "lines.linewidth": 2.0, "legend.frameon": False, "font.size": 10,
+    })

decoding.py

@@ -0,0 +1,558 @@
+# =============================================================================
+# decoding.py - Part 3: decoding strategies implemented from scratch on top of
+# EleutherAI/pythia-160m.  Only model.forward() (i.e. model(...)) is used -
+# never model.generate().
+#
+#   greedy, beam search (width 1/2/4/...), top-k sampling, top-p (nucleus) sampling
+#
+# Evaluated on hamishivi/ROCStories (test split): prompt = first sentence,
+# reference = the remaining four sentences.
+# =============================================================================
+import os, re, math, time, json, string
+from collections import Counter
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+DECODING_DEFAULTS = dict(
+    model_name="EleutherAI/pythia-160m",
+    dataset="hamishivi/ROCStories", split="test", n_samples=1000, seed=0,
+    max_new_tokens=70, batch_size=32, beam_batch_size=8, dtype="float32",
+    length_penalty=1.0,
+    configs=[
+        {"name": "greedy", "strategy": "greedy"},
+        {"name": "beam_w1", "strategy": "beam", "width": 1},
+        {"name": "beam_w2", "strategy": "beam", "width": 2},
+        {"name": "beam_w4", "strategy": "beam", "width": 4},
+        {"name": "topk_k5", "strategy": "top_k", "k": 5},
+        {"name": "topk_k20", "strategy": "top_k", "k": 20},
+        {"name": "topk_k50", "strategy": "top_k", "k": 50},
+        {"name": "topk_k50_t0.7", "strategy": "top_k", "k": 50, "temperature": 0.7},
+        {"name": "topp_p0.5", "strategy": "top_p", "p": 0.5},
+        {"name": "topp_p0.8", "strategy": "top_p", "p": 0.8},
+        {"name": "topp_p0.95", "strategy": "top_p", "p": 0.95},
+        {"name": "topp_p0.9_t0.7", "strategy": "top_p", "p": 0.9, "temperature": 0.7},
+    ],
+    timing_widths=[1, 2, 4, 8], timing_samples=100,
+    judge_model="EleutherAI/pythia-410m",   # external LM for an extra perplexity column (None to skip)
+    use_bertscore=True,
+    out_dir="/kaggle/working/outputs/decoding" if os.path.isdir("/kaggle/working") else "outputs/decoding",
+    hf_user=None, repo_name="anlp-a2-decoding", push_to_hub=True, private=False,
+)
+
+
+# -----------------------------------------------------------------------------
+# model wrapper: one forward step with KV cache
+# -----------------------------------------------------------------------------
+def load_lm(name, device, dtype="float32"):
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    tok = AutoTokenizer.from_pretrained(name)
+    dt = {"float32": torch.float32, "float16": torch.float16, "bfloat16": torch.bfloat16}[dtype]
+    try:
+        model = AutoModelForCausalLM.from_pretrained(name, dtype=dt)
+    except TypeError:
+        model = AutoModelForCausalLM.from_pretrained(name, torch_dtype=dt)
+    model.to(device).eval()
+    if tok.pad_token_id is None:
+        tok.pad_token = tok.eos_token
+    return model, tok
+
+
+def reorder_cache(past, idx):
+    """Select rows (beams) of a KV cache - works for legacy tuples and Cache objects."""
+    if past is None:
+        return None
+    if isinstance(past, (tuple, list)):
+        return tuple(tuple(t.index_select(0, idx) for t in layer) for layer in past)
+    if hasattr(past, "reorder_cache"):
+        past.reorder_cache(idx)
+        return past
+    if hasattr(past, "layers"):          # transformers >= 4.56 DynamicCache
+        for layer in past.layers:
+            layer.keys = layer.keys.index_select(0, idx)
+            layer.values = layer.values.index_select(0, idx)
+        return past
+    if hasattr(past, "key_cache"):
+        for i in range(len(past.key_cache)):
+            past.key_cache[i] = past.key_cache[i].index_select(0, idx)
+            past.value_cache[i] = past.value_cache[i].index_select(0, idx)
+        return past
+    raise TypeError(f"cannot reorder cache of type {type(past)}")
+
+
+@torch.no_grad()
+def forward_step(model, input_ids, attention_mask, position_ids, past):
+    out = model(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids,
+                past_key_values=past, use_cache=True)
+    return out.logits[:, -1, :].float(), out.past_key_values
+
+
+def left_pad(prompts, pad_id, device):
+    L = max(len(p) for p in prompts)
+    ids = torch.full((len(prompts), L), pad_id, dtype=torch.long)
+    mask = torch.zeros((len(prompts), L), dtype=torch.long)
+    for i, p in enumerate(prompts):
+        ids[i, L - len(p):] = torch.tensor(p)
+        mask[i, L - len(p):] = 1
+    pos = (mask.cumsum(1) - 1).clamp(min=0)
+    return ids.to(device), mask.to(device), pos.to(device)
+
+
+# -----------------------------------------------------------------------------
+# logit processors
+# -----------------------------------------------------------------------------
+def top_k_filter(logits, k):
+    k = min(k, logits.size(-1))
+    kth = torch.topk(logits, k, dim=-1).values[:, -1:]
+    return logits.masked_fill(logits < kth, float("-inf"))
+
+
+def top_p_filter(logits, p):
+    """Keep the smallest set of tokens whose cumulative probability >= p (always >= 1 token)."""
+    sorted_logits, sorted_idx = torch.sort(logits, descending=True, dim=-1)
+    probs = F.softmax(sorted_logits, dim=-1)
+    cum = probs.cumsum(-1)
+    remove = (cum - probs) >= p          # mass *before* this token already reaches p
+    remove[:, 0] = False
+    sorted_logits = sorted_logits.masked_fill(remove, float("-inf"))
+    return torch.full_like(logits, float("-inf")).scatter(-1, sorted_idx, sorted_logits)
+
+
+# -----------------------------------------------------------------------------
+# greedy / top-k / top-p (batched)
+# -----------------------------------------------------------------------------
+@torch.no_grad()
+def sample_decode(model, prompts, max_new_tokens, eos_id, pad_id, strategy="greedy", k=None, p=None,
+                  temperature=1.0, generator=None):
+    """Returns (list of generated id lists, list of sum log-prob under the *unmodified*
+    model distribution, list of lengths)."""
+    device = next(model.parameters()).device
+    B = len(prompts)
+    ids, mask, pos = left_pad(prompts, pad_id, device)
+    logits, past = forward_step(model, ids, mask, pos, None)
+    next_pos = pos[:, -1:] + 1
+    finished = torch.zeros(B, dtype=torch.bool, device=device)
+    out, lp_sum, lengths = [], torch.zeros(B, device=device), torch.zeros(B, dtype=torch.long, device=device)
+    for _ in range(max_new_tokens):
+        logp = F.log_softmax(logits, -1)
+        if strategy == "greedy":
+            nxt = logits.argmax(-1)
+        else:
+            lg = logits / max(temperature, 1e-6)
+            if strategy == "top_k":
+                lg = top_k_filter(lg, k)
+            elif strategy == "top_p":
+                lg = top_p_filter(lg, p)
+            nxt = torch.multinomial(F.softmax(lg, -1), 1, generator=generator).squeeze(1)
+        nxt = torch.where(finished, torch.full_like(nxt, pad_id), nxt)
+        lp_sum += torch.where(finished, torch.zeros_like(lp_sum), logp.gather(1, nxt[:, None]).squeeze(1))
+        lengths += (~finished).long()
+        out.append(nxt)
+        finished |= nxt == eos_id
+        if bool(finished.all()):
+            break
+        mask = torch.cat([mask, torch.ones(B, 1, dtype=mask.dtype, device=device)], 1)
+        logits, past = forward_step(model, nxt[:, None], mask, next_pos, past)
+        next_pos = next_pos + 1
+    gen = torch.stack(out, 1).tolist() if out else [[] for _ in range(B)]
+    res = []
+    for row in gen:
+        if eos_id in row:
+            row = row[:row.index(eos_id)]
+        res.append(row)
+    return res, lp_sum.tolist(), lengths.tolist()
+
+
+# -----------------------------------------------------------------------------
+# beam search (batched over prompts, width W per prompt)
+# -----------------------------------------------------------------------------
+@torch.no_grad()
+def beam_search(model, prompts, width, max_new_tokens, eos_id, pad_id, length_penalty=1.0):
+    """Standard beam search with a KV cache.
+    * each step: scores(B,W) + log p(B,W,V) -> top 2W candidates per prompt
+    * candidates ending in EOS become finished hypotheses (score / len^alpha)
+    * a prompt is done once it has W finished hypotheses (early stopping)
+    * at max length, open beams are added to the finished pool.
+    Returns (ids, sum log-prob, length) of the best hypothesis per prompt."""
+    device = next(model.parameters()).device
+    B, W = len(prompts), width
+    rep = [p for p in prompts for _ in range(W)]
+    ids, mask, pos = left_pad(rep, pad_id, device)
+    logits, past = forward_step(model, ids, mask, pos, None)
+    V = logits.size(-1)
+    next_pos = pos[:, -1:] + 1
+    beam_scores = torch.zeros(B, W, device=device)
+    beam_scores[:, 1:] = float("-inf")       # all beams identical at t=0 -> only expand beam 0
+    seqs = torch.zeros(B * W, 0, dtype=torch.long, device=device)
+    finished = [[] for _ in range(B)]        # (normalised score, raw score, token list)
+    done = [False] * B
+    for t in range(max_new_tokens):
+        logp = F.log_softmax(logits, -1).view(B, W, V)
+        cand = (beam_scores.unsqueeze(-1) + logp).view(B, W * V)
+        top_s, top_i = cand.topk(2 * W, dim=-1)
+        top_s_l, top_i_l = top_s.tolist(), top_i.tolist()
+        new_scores = torch.zeros(B, W, device=device)
+        new_beam = torch.zeros(B, W, dtype=torch.long)
+        new_tok = torch.full((B, W), pad_id, dtype=torch.long)
+        seqs_l = None
+        for b in range(B):
+            if done[b]:
+                new_scores[b] = float("-inf"); new_scores[b, 0] = 0.0
+                continue
+            j = 0
+            for r in range(2 * W):
+                s, ci = top_s_l[b][r], top_i_l[b][r]
+                bi, ti = ci // V, ci % V
+                if s == float("-inf"):
+                    break
+                if ti == eos_id:
+                    if r < W:
+                        if seqs_l is None:
+                            seqs_l = seqs.tolist()
+                        toks = seqs_l[b * W + bi]
+                        n = len(toks) + 1
+                        finished[b].append((s / (n ** length_penalty), s, toks, n))
+                    continue
+                new_scores[b, j] = s; new_beam[b, j] = bi; new_tok[b, j] = ti
+                j += 1
+                if j == W:
+                    break
+            while j < W:                      # not enough live candidates
+                new_scores[b, j] = float("-inf"); j += 1
+            if len(finished[b]) >= W:
+                done[b] = True
+        if all(done):
+            break
+        gidx = (torch.arange(B).unsqueeze(1) * W + new_beam).view(-1).to(device)
+        seqs = torch.cat([seqs.index_select(0, gidx), new_tok.view(-1, 1).to(device)], 1)
+        beam_scores = new_scores
+        past = reorder_cache(past, gidx)
+        mask = torch.cat([mask.index_select(0, gidx), torch.ones(B * W, 1, dtype=mask.dtype, device=device)], 1)
+        next_pos = next_pos.index_select(0, gidx)
+        logits, past = forward_step(model, new_tok.view(-1, 1).to(device), mask, next_pos, past)
+        next_pos = next_pos + 1
+    seqs_l = seqs.tolist(); sc = beam_scores.tolist()
+    res, lps, lens = [], [], []
+    for b in range(B):
+        if not done[b]:
+            for w in range(W):
+                if sc[b][w] > float("-inf"):
+                    toks = seqs_l[b * W + w]
+                    n = max(len(toks), 1)
+                    finished[b].append((sc[b][w] / (n ** length_penalty), sc[b][w], toks, n))
+        best = max(finished[b], key=lambda z: z[0])
+        res.append(best[2]); lps.append(best[1]); lens.append(best[3])
+    return res, lps, lens
+
+
+# -----------------------------------------------------------------------------
+# metrics
+# -----------------------------------------------------------------------------
+_PUNCT = re.compile(f"[{re.escape(string.punctuation)}]")
+
+
+def norm_words(s):
+    return _PUNCT.sub(" ", s.lower()).split()
+
+
+def clean_generation(text, max_sentences=4):
+    """Post-processing applied before scoring: cut at the first newline (Pythia tends to
+    drift into unrelated text after a line break) and keep at most 4 sentences
+    (ROCStories continuations are exactly 4 sentences)."""
+    text = text.lstrip()
+    text = text.split("\n")[0]
+    ends = [m.end() for m in re.finditer(r"[.!?](?=\s|$)", text)]
+    if len(ends) >= max_sentences:
+        text = text[:ends[max_sentences - 1]]
+    return " " + text.strip()
+
+
+def token_prf_acc(gen, ref):
+    g, r = norm_words(gen), norm_words(ref)
+    common = Counter(g) & Counter(r)
+    ov = sum(common.values())
+    P = ov / len(g) if g else 0.0
+    R = ov / len(r) if r else 0.0
+    F1 = 2 * P * R / (P + R) if P + R > 0 else 0.0
+    acc = sum(1 for i in range(len(r)) if i < len(g) and g[i] == r[i]) / len(r) if r else 0.0
+    return P, R, F1, acc
+
+
+def distinct_n(texts, n):
+    grams, total = set(), 0
+    for t in texts:
+        w = norm_words(t)
+        ng = [tuple(w[i:i + n]) for i in range(len(w) - n + 1)]
+        grams.update(ng); total += len(ng)
+    return len(grams) / max(total, 1)
+
+
+def repetition_rate(text, n=4):
+    w = norm_words(text)
+    ng = [tuple(w[i:i + n]) for i in range(len(w) - n + 1)]
+    return 1 - len(set(ng)) / len(ng) if ng else 0.0
+
+
+@torch.no_grad()
+def continuation_ppl(model, prompt_ids, cont_ids, pad_id, batch_size=32):
+    """Per-sample perplexity of cont given prompt (teacher forced)."""
+    device = next(model.parameters()).device
+    out = []
+    for s in range(0, len(prompt_ids), batch_size):
+        P = prompt_ids[s:s + batch_size]; C = cont_ids[s:s + batch_size]
+        seqs = [p + c for p, c in zip(P, C)]
+        L = max(len(x) for x in seqs)
+        ids = torch.full((len(seqs), L), pad_id, dtype=torch.long)
+        att = torch.zeros((len(seqs), L), dtype=torch.long)
+        lab = torch.full((len(seqs), L), -100, dtype=torch.long)
+        for i, (p, c) in enumerate(zip(P, C)):
+            x = p + c
+            ids[i, :len(x)] = torch.tensor(x); att[i, :len(x)] = 1
+            lab[i, len(p):len(x)] = torch.tensor(c) if c else lab[i, len(p):len(x)]
+        ids, att, lab = ids.to(device), att.to(device), lab.to(device)
+        logits = model(input_ids=ids, attention_mask=att).logits.float()
+        lp = F.cross_entropy(logits[:, :-1].reshape(-1, logits.size(-1)), lab[:, 1:].reshape(-1),
+                             ignore_index=-100, reduction="none").view(len(seqs), -1)
+        n = (lab[:, 1:] != -100).sum(1).clamp(min=1)
+        out += torch.exp(lp.sum(1) / n).tolist()
+    return out
+
+
+def score_config(gens, refs, self_lp, self_len):
+    import sacrebleu
+    prf = np.array([token_prf_acc(g, r) for g, r in zip(gens, refs)])
+    ppl_self = [math.exp(-lp / max(n, 1)) for lp, n in zip(self_lp, self_len)]
+    return {
+        "precision": float(prf[:, 0].mean()), "recall": float(prf[:, 1].mean()),
+        "f1": float(prf[:, 2].mean()), "accuracy": float(prf[:, 3].mean()),
+        "bleu": sacrebleu.corpus_bleu(gens, [refs]).score,
+        "self_ppl_mean": float(np.mean(ppl_self)), "self_ppl_median": float(np.median(ppl_self)),
+        "distinct_1": distinct_n(gens, 1), "distinct_2": distinct_n(gens, 2),
+        "rep_4gram": float(np.mean([repetition_rate(g) for g in gens])),
+        "avg_words": float(np.mean([len(norm_words(g)) for g in gens])),
+    }
+
+
+# -----------------------------------------------------------------------------
+# driver
+# -----------------------------------------------------------------------------
+def run_config(model, tok, prompts, conf, cfg, seed=0):
+    device = next(model.parameters()).device
+    eos, pad = tok.eos_token_id, tok.pad_token_id
+    gen_ids, lps, lens = [None] * len(prompts), [None] * len(prompts), [None] * len(prompts)
+    order = sorted(range(len(prompts)), key=lambda i: len(prompts[i]))
+    g = torch.Generator(device=device); g.manual_seed(seed)
+    bs = cfg["beam_batch_size"] if conf["strategy"] == "beam" else cfg["batch_size"]
+    if device.type == "cuda": torch.cuda.synchronize()
+    t0 = time.time()
+    for s in range(0, len(order), bs):
+        idx = order[s:s + bs]
+        P = [prompts[i] for i in idx]
+        if conf["strategy"] == "beam":
+            r, lp, ln = beam_search(model, P, conf["width"], cfg["max_new_tokens"], eos, pad, cfg["length_penalty"])
+        else:
+            r, lp, ln = sample_decode(model, P, cfg["max_new_tokens"], eos, pad, conf["strategy"],
+                                      conf.get("k"), conf.get("p"), conf.get("temperature", 1.0), g)
+        for i, a, b, c in zip(idx, r, lp, ln):
+            gen_ids[i], lps[i], lens[i] = a, b, c
+    if device.type == "cuda": torch.cuda.synchronize()
+    return gen_ids, lps, lens, time.time() - t0
+
+
+def beam_timing(model, tok, prompts, widths, max_new_tokens):
+    """Wall time of beam search with batch size 1 (one prompt at a time) for each width.
+    Two modes: 'natural' (stops early once W hypotheses hit EOS) and 'fixed' (EOS disabled,
+    always max_new_tokens steps) - the fixed mode isolates the per-step cost as a function of W."""
+    device = next(model.parameters()).device
+    res = []
+    beam_search(model, prompts[:2], 2, 5, tok.eos_token_id, tok.pad_token_id)   # warm-up
+    for mode, eos in (("natural", tok.eos_token_id), ("fixed", -1)):
+        for w in widths:
+            times, steps = [], []
+            for p in prompts:
+                if device.type == "cuda": torch.cuda.synchronize()
+                t0 = time.time()
+                r, _, ln = beam_search(model, [p], w, max_new_tokens, eos, tok.pad_token_id)
+                if device.type == "cuda": torch.cuda.synchronize()
+                times.append(time.time() - t0); steps.append(max(len(r[0]), 1))
+            rec = {"mode": mode, "width": w, "mean_s_per_prompt": float(np.mean(times)), "std_s": float(np.std(times)),
+                   "mean_ms_per_token": float(np.mean([t / s for t, s in zip(times, steps)]) * 1e3),
+                   "mean_output_tokens": float(np.mean(steps)), "total_s": float(np.sum(times)),
+                   "n_prompts": len(prompts)}
+            res.append(rec)
+            print(f"  [{mode}] width {w}: {rec['mean_s_per_prompt']*1e3:.1f} ms/prompt, "
+                  f"{rec['mean_ms_per_token']:.2f} ms/token, {rec['mean_output_tokens']:.1f} tokens")
+    return res
+
+
+def run_decoding_eval(user_cfg=None):
+    import sys
+    sys.path.insert(0, ".")
+    from common import get_hf_token, hf_login, hf_resolve_user, hf_push_folder, save_json
+    cfg = dict(DECODING_DEFAULTS); cfg.update(user_cfg or {})
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    os.makedirs(cfg["out_dir"], exist_ok=True)
+    torch.manual_seed(cfg["seed"])
+
+    token = get_hf_token(); hf_login(token)
+    model, tok = load_lm(cfg["model_name"], device, cfg["dtype"])
+    print(f"loaded {cfg['model_name']} on {device}; eos={tok.eos_token_id} pad={tok.pad_token_id}")
+
+    if cfg.get("dataset_override") is not None:
+        rows = cfg["dataset_override"]
+    else:
+        from datasets import load_dataset
+        rows = load_dataset(cfg["dataset"], split=cfg["split"])
+    prompts_txt = list(rows["prompt"]); refs_txt = list(rows["continuation"])
+    keep = [i for i, r in enumerate(refs_txt) if r and r.strip()]
+    rng = np.random.default_rng(cfg["seed"])
+    keep = sorted(rng.choice(keep, min(cfg["n_samples"], len(keep)), replace=False).tolist())
+    prompts_txt = [prompts_txt[i] for i in keep]; refs_txt = [refs_txt[i] for i in keep]
+    refs_txt = [" " + r.strip() for r in refs_txt]
+    prompts = [tok(p)["input_ids"] for p in prompts_txt]
+    print(f"{len(prompts)} test prompts")
+
+    # reference perplexity (independent of the decoding strategy)
+    ref_ids = [tok(r)["input_ids"] for r in refs_txt]
+    ref_ppl = continuation_ppl(model, prompts, ref_ids, tok.pad_token_id)
+    summary = {"reference_ppl_mean": float(np.mean(ref_ppl)), "reference_ppl_median": float(np.median(ref_ppl)),
+               "n_samples": len(prompts), "model": cfg["model_name"], "max_new_tokens": cfg["max_new_tokens"]}
+    print(f"reference continuation ppl: mean {summary['reference_ppl_mean']:.2f}, "
+          f"median {summary['reference_ppl_median']:.2f}")
+
+    results, generations = [], {"prompt": prompts_txt, "reference": refs_txt}
+    for conf in cfg["configs"]:
+        print(f"== {conf['name']} ==")
+        ids, lps, lens, dt = run_config(model, tok, prompts, conf, cfg, cfg["seed"])
+        raw = [tok.decode(x, skip_special_tokens=True) for x in ids]
+        gens = [clean_generation(t) for t in raw]
+        m = score_config(gens, refs_txt, lps, lens)
+        m.update({"name": conf["name"], "strategy": conf["strategy"], "params": conf,
+                  "time_s": dt, "ms_per_prompt": dt / len(prompts) * 1e3,
+                  "gen_tokens_mean": float(np.mean([len(x) for x in ids]))})
+        results.append(m)
+        generations[conf["name"]] = gens
+        generations[conf["name"] + "__raw"] = raw
+        print(f"   P {m['precision']:.3f} R {m['recall']:.3f} F1 {m['f1']:.3f} Acc {m['accuracy']:.3f} "
+              f"BLEU {m['bleu']:.2f} self-ppl {m['self_ppl_median']:.2f} dist-2 {m['distinct_2']:.3f} "
+              f"rep4 {m['rep_4gram']:.3f} time {dt:.1f}s")
+
+    # sanity: beam width 1 must equal greedy
+    if "greedy" in generations and "beam_w1" in generations:
+        same = np.mean([a == b for a, b in zip(generations["greedy__raw"], generations["beam_w1__raw"])])
+        summary["beam_w1_equals_greedy_frac"] = float(same)
+        print(f"beam(w=1) identical to greedy on {same*100:.1f}% of prompts")
+
+    # external judge perplexity
+    if cfg["judge_model"]:
+        try:
+            judge, _ = load_lm(cfg["judge_model"], device, cfg["dtype"])
+            for m in results:
+                gen_ids = [tok(g)["input_ids"] for g in generations[m["name"]]]
+                pp = continuation_ppl(judge, prompts, gen_ids, tok.pad_token_id)
+                m["judge_ppl_median"] = float(np.median(pp)); m["judge_ppl_mean"] = float(np.mean(pp))
+            jr = continuation_ppl(judge, prompts, ref_ids, tok.pad_token_id)
+            summary["reference_judge_ppl_median"] = float(np.median(jr))
+            del judge
+        except Exception as e:
+            print("[judge] skipped:", e)
+
+    if cfg["use_bertscore"]:
+        try:
+            from bert_score import score as bscore
+            for m in results:
+                P, R, F1 = bscore(generations[m["name"]], refs_txt, lang="en", verbose=False,
+                                  device=str(device), batch_size=64)
+                m["bertscore_p"], m["bertscore_r"], m["bertscore_f1"] = (float(P.mean()), float(R.mean()),
+                                                                         float(F1.mean()))
+        except Exception as e:
+            print("[bertscore] skipped:", e)
+
+    print("beam-search timing (batch size 1):")
+    timing = beam_timing(model, tok, prompts[:cfg["timing_samples"]], cfg["timing_widths"], cfg["max_new_tokens"])
+
+    save_json({"summary": summary, "results": results, "beam_timing": timing, "config":
+               {k: v for k, v in cfg.items() if k != "dataset_override"}},
+              os.path.join(cfg["out_dir"], "decoding_results.json"))
+    save_json(generations, os.path.join(cfg["out_dir"], "generations.json"))
+    shutil_copy_code(cfg["out_dir"])
+    user = hf_resolve_user(cfg["hf_user"], token)
+    if cfg["push_to_hub"] and user:
+        hf_push_folder(cfg["out_dir"], f"{user}/{cfg['repo_name']}", token, cfg["private"], "decoding results")
+        print(f"pushed to https://huggingface.co/{user}/{cfg['repo_name']}")
+    return summary, results, timing, generations
+
+
+def shutil_copy_code(out_dir):
+    import shutil
+    for f in ("decoding.py", "common.py"):
+        if os.path.exists(f):
+            shutil.copy(f, os.path.join(out_dir, f))
+
+
+# -----------------------------------------------------------------------------
+# plots / tables for the write-up
+# -----------------------------------------------------------------------------
+PALETTE = ["#2a78d6", "#eb6834", "#1baf7a", "#eda100", "#e87ba4", "#008300", "#4a3aa7", "#e34948"]
+STRATEGY_COLOR = {"greedy": PALETTE[0], "beam": PALETTE[6], "top_k": PALETTE[1], "top_p": PALETTE[2]}
+
+
+def decoding_report(res_json_path, out_dir):
+    import pandas as pd
+    import matplotlib.pyplot as plt
+    from common import setup_plot_style
+    setup_plot_style()
+    R = json.load(open(res_json_path))
+    df = pd.DataFrame(R["results"])
+    cols = ["name", "accuracy", "precision", "recall", "f1", "bleu", "self_ppl_median", "judge_ppl_median",
+            "bertscore_f1", "distinct_2", "rep_4gram", "avg_words", "ms_per_prompt"]
+    table = df[[c for c in cols if c in df.columns]].copy()
+    table.to_csv(os.path.join(out_dir, "decoding_table.csv"), index=False)
+    colors = [STRATEGY_COLOR[s] for s in df["strategy"]]
+
+    metrics = [("f1", "token F1"), ("accuracy", "positional accuracy"), ("bleu", "BLEU"),
+               ("self_ppl_median", "self perplexity (median)"), ("distinct_2", "distinct-2"),
+               ("rep_4gram", "repeated 4-gram rate")]
+    fig, axes = plt.subplots(2, 3, figsize=(15, 7.5), constrained_layout=True)
+    for ax, (k, lab) in zip(axes.flat, metrics):
+        ax.bar(range(len(df)), df[k], color=colors, width=0.7)
+        ax.set_xticks(range(len(df))); ax.set_xticklabels(df["name"], rotation=60, ha="right", fontsize=8)
+        ax.set_title(lab); ax.grid(axis="x", visible=False)
+    handles = [plt.Rectangle((0, 0), 1, 1, color=c) for c in STRATEGY_COLOR.values()]
+    fig.legend(handles, STRATEGY_COLOR.keys(), loc="upper right", ncol=4)
+    fig.savefig(os.path.join(out_dir, "decoding_metrics.png"))
+
+    # quality vs diversity
+    fig2, ax = plt.subplots(figsize=(6.5, 4.5), constrained_layout=True)
+    for _, r in df.iterrows():
+        ax.scatter(r["distinct_2"], r["f1"], color=STRATEGY_COLOR[r["strategy"]], s=60, zorder=3,
+                   edgecolor="white", linewidth=1.5)
+        ax.annotate(r["name"], (r["distinct_2"], r["f1"]), fontsize=7, xytext=(4, 3), textcoords="offset points")
+    ax.set_xlabel("distinct-2 (diversity)"); ax.set_ylabel("token F1 vs reference")
+    ax.set_title("quality / diversity trade-off")
+    fig2.savefig(os.path.join(out_dir, "quality_vs_diversity.png"))
+
+    # beam timing
+    tm = pd.DataFrame(R["beam_timing"])
+    fig3, axes3 = plt.subplots(1, 2, figsize=(11, 3.8), constrained_layout=True)
+    for i, (mode, g) in enumerate(tm.groupby("mode", sort=False)):
+        axes3[0].plot(g["width"], g["mean_s_per_prompt"] * 1e3, marker="o", label=mode, color=PALETTE[i])
+        axes3[1].plot(g["width"], g["mean_ms_per_token"], marker="o", label=mode, color=PALETTE[i])
+    for ax, lab in zip(axes3, ["ms per prompt", "ms per generated token"]):
+        ax.set_xlabel("beam width"); ax.set_ylabel(lab); ax.set_xticks(sorted(tm["width"].unique())); ax.legend()
+    axes3[0].set_title("beam search wall time (batch size 1)")
+    axes3[1].set_title("per-step cost")
+    fig3.savefig(os.path.join(out_dir, "beam_timing.png"))
+    return table, tm
+
+
+def show_examples(gen_json_path, n=5, names=None, seed=0):
+    G = json.load(open(gen_json_path))
+    names = names or [k for k in G if k not in ("prompt", "reference") and not k.endswith("__raw")]
+    rng = np.random.default_rng(seed)
+    for i in rng.choice(len(G["prompt"]), n, replace=False):
+        print("=" * 100)
+        print("PROMPT   :", G["prompt"][i])
+        print("REFERENCE:", G["reference"][i].strip())
+        for k in names:
+            print(f"{k:>16s}: {G[k][i].strip()}")

decoding_results.json

@@ -1,40 +1,40 @@
 {
   "summary": {
-    "reference_ppl_mean": 27.426975273132324,
-    "reference_ppl_median": 24.039149284362793,
+    "reference_ppl_mean": 27.426984312057495,
+    "reference_ppl_median": 24.0395565032959,
     "n_samples": 1000,
     "model": "EleutherAI/pythia-160m",
     "max_new_tokens": 70,
-    "beam_w1_equals_greedy_frac": 0.998,
-    "reference_judge_ppl_median": 16.60109519958496
+    "beam_w1_equals_greedy_frac": 0.997,
+    "reference_judge_ppl_median": 16.601248741149902
   },
   "results": [
     {
-      "precision": 0.15680400502615022,
+      "precision": 0.156788215552466,
       "recall": 0.10761634910100995,
-      "f1": 0.11683061986311886,
-      "accuracy": 0.013858261701571426,
-      "bleu": 0.45979362180213906,
-      "self_ppl_mean": 2.129511192101869,
-      "self_ppl_median": 1.977763465028044,
-      "distinct_1": 0.03781326236416057,
-      "distinct_2": 0.11821601289629231,
+      "f1": 0.11682790615891263,
+      "accuracy": 0.013836985105826744,
+      "bleu": 0.45975591061692944,
+      "self_ppl_mean": 2.129623557711418,
+      "self_ppl_median": 1.9777633712182454,
+      "distinct_1": 0.037811864720014784,
+      "distinct_2": 0.11828823642295144,
       "rep_4gram": 0.3296551592928797,
-      "avg_words": 27.054,
+      "avg_words": 27.055,
       "name": "greedy",
       "strategy": "greedy",
       "params": {
         "name": "greedy",
         "strategy": "greedy"
       },
-      "time_s": 24.53210139274597,
-      "ms_per_prompt": 24.53210139274597,
+      "time_s": 26.33022141456604,
+      "ms_per_prompt": 26.33022141456604,
       "gen_tokens_mean": 69.995,
-      "judge_ppl_median": 3.7136964797973633,
-      "judge_ppl_mean": 3.8388092192411425,
-      "bertscore_p": 0.8487208485603333,
-      "bertscore_r": 0.8463250994682312,
-      "bertscore_f1": 0.8473963141441345
+      "judge_ppl_median": 3.7136934995651245,
+      "judge_ppl_mean": 3.8431873861551287,
+      "bertscore_p": 0.8487190008163452,
+      "bertscore_r": 0.8463276028633118,
+      "bertscore_f1": 0.847396731376648
     },
     {
       "precision": 0.1568385962211188,
@@ -42,8 +42,8 @@
       "f1": 0.11688205526981743,
       "accuracy": 0.013886833130142854,
       "bleu": 0.4603869452185167,
-      "self_ppl_mean": 2.1297556273252223,
-      "self_ppl_median": 1.9777561918932842,
+      "self_ppl_mean": 2.1293727944369447,
+      "self_ppl_median": 1.977726615504968,
       "distinct_1": 0.03780348102435239,
       "distinct_2": 0.11810751697939449,
       "rep_4gram": 0.3296463873630552,
@@ -55,11 +55,11 @@
         "strategy": "beam",
         "width": 1
       },
-      "time_s": 89.72963857650757,
-      "ms_per_prompt": 89.72963857650757,
+      "time_s": 110.271555185318,
+      "ms_per_prompt": 110.271555185318,
       "gen_tokens_mean": 69.995,
-      "judge_ppl_median": 3.7136964797973633,
-      "judge_ppl_mean": 3.839159709095955,
+      "judge_ppl_median": 3.7136934995651245,
+      "judge_ppl_mean": 3.839161422610283,
       "bertscore_p": 0.8486897945404053,
       "bertscore_r": 0.8463218808174133,
       "bertscore_f1": 0.8473787903785706
@@ -70,8 +70,8 @@
       "f1": 0.11012341441029258,
       "accuracy": 0.010045733414015445,
       "bleu": 0.29556628142570907,
-      "self_ppl_mean": 1.8974248541153338,
-      "self_ppl_median": 1.8301979892743205,
+      "self_ppl_mean": 1.8974243870210339,
+      "self_ppl_median": 1.8302103567597376,
       "distinct_1": 0.045782342657342656,
       "distinct_2": 0.13661581137309292,
       "rep_4gram": 0.20882754144393487,
@@ -83,27 +83,27 @@
         "strategy": "beam",
         "width": 2
       },
-      "time_s": 93.67243909835815,
-      "ms_per_prompt": 93.67243909835815,
+      "time_s": 107.95216870307922,
+      "ms_per_prompt": 107.95216870307922,
       "gen_tokens_mean": 69.827,
-      "judge_ppl_median": 4.0844643115997314,
-      "judge_ppl_mean": 5.812758823871612,
+      "judge_ppl_median": 4.08447003364563,
+      "judge_ppl_mean": 5.812772500276566,
       "bertscore_p": 0.8609820604324341,
-      "bertscore_r": 0.846573531627655,
+      "bertscore_r": 0.8465735912322998,
       "bertscore_f1": 0.8535703420639038
     },
     {
-      "precision": 0.197717653773812,
+      "precision": 0.1976862812247924,
       "recall": 0.07404912467040367,
-      "f1": 0.09570392741330655,
+      "f1": 0.09568651928951545,
       "accuracy": 0.008065531697039307,
-      "bleu": 0.3401904436011399,
-      "self_ppl_mean": 1.7576851325086995,
-      "self_ppl_median": 1.739890192921894,
-      "distinct_1": 0.04027860266724259,
-      "distinct_2": 0.11694537982991839,
-      "rep_4gram": 0.23772850694366054,
-      "avg_words": 18.521,
+      "bleu": 0.340259807839598,
+      "self_ppl_mean": 1.7576852720858818,
+      "self_ppl_median": 1.739890666494589,
+      "distinct_1": 0.040265558374264585,
+      "distinct_2": 0.11690534603754207,
+      "rep_4gram": 0.23782374503889867,
+      "avg_words": 18.527,
       "name": "beam_w4",
       "strategy": "beam",
       "params": {
@@ -111,27 +111,27 @@
         "strategy": "beam",
         "width": 4
       },
-      "time_s": 106.44298911094666,
-      "ms_per_prompt": 106.44298911094666,
+      "time_s": 118.95717024803162,
+      "ms_per_prompt": 118.95717024803162,
       "gen_tokens_mean": 69.253,
-      "judge_ppl_median": 3.722070574760437,
-      "judge_ppl_mean": 549.2117587163449,
-      "bertscore_p": 0.8594975471496582,
-      "bertscore_r": 0.8445566892623901,
-      "bertscore_f1": 0.8517829775810242
+      "judge_ppl_median": 3.7220958471298218,
+      "judge_ppl_mean": 549.211596192956,
+      "bertscore_p": 0.8594744801521301,
+      "bertscore_r": 0.8445523977279663,
+      "bertscore_f1": 0.851769208908081
     },
     {
-      "precision": 0.20247790936126098,
-      "recall": 0.16240848409656986,
-      "f1": 0.16326782514510493,
-      "accuracy": 0.015387661973703906,
-      "bleu": 0.6321177804113657,
-      "self_ppl_mean": 5.788990184998834,
-      "self_ppl_median": 5.7832007131546685,
-      "distinct_1": 0.07138813761871816,
-      "distinct_2": 0.35044483408013505,
-      "rep_4gram": 0.03276803420190723,
-      "avg_words": 31.798,
+      "precision": 0.20828856855773425,
+      "recall": 0.16455016427602215,
+      "f1": 0.16573516044843375,
+      "accuracy": 0.015756478767379786,
+      "bleu": 0.5589188575413241,
+      "self_ppl_mean": 5.8339501863641905,
+      "self_ppl_median": 5.838575948334771,
+      "distinct_1": 0.07212438521339044,
+      "distinct_2": 0.35802064558413893,
+      "rep_4gram": 0.028893427447938293,
+      "avg_words": 31.515,
       "name": "topk_k5",
       "strategy": "top_k",
       "params": {
@@ -139,27 +139,27 @@
         "strategy": "top_k",
         "k": 5
       },
-      "time_s": 24.93057608604431,
-      "ms_per_prompt": 24.93057608604431,
-      "gen_tokens_mean": 69.711,
-      "judge_ppl_median": 8.327152252197266,
-      "judge_ppl_mean": 10.00695155286789,
-      "bertscore_p": 0.8572027683258057,
-      "bertscore_r": 0.8528223037719727,
-      "bertscore_f1": 0.8548986911773682
+      "time_s": 28.06767702102661,
+      "ms_per_prompt": 28.06767702102661,
+      "gen_tokens_mean": 69.697,
+      "judge_ppl_median": 8.458185195922852,
+      "judge_ppl_mean": 9.329122072458267,
+      "bertscore_p": 0.8579514026641846,
+      "bertscore_r": 0.8530383706092834,
+      "bertscore_f1": 0.8553735613822937
     },
     {
-      "precision": 0.18981835217131812,
-      "recall": 0.17505740033246367,
-      "f1": 0.16608450826242258,
-      "accuracy": 0.013520378544886719,
-      "bleu": 0.5262730257755575,
-      "self_ppl_mean": 10.425129562862075,
-      "self_ppl_median": 10.068886870395898,
-      "distinct_1": 0.08757376684110901,
-      "distinct_2": 0.4628797205600252,
-      "rep_4gram": 0.005379857536448407,
-      "avg_words": 35.924,
+      "precision": 0.1919000986912784,
+      "recall": 0.1684283187083047,
+      "f1": 0.1622514196784391,
+      "accuracy": 0.013794441209097598,
+      "bleu": 0.44804585179427503,
+      "self_ppl_mean": 10.859201539564154,
+      "self_ppl_median": 9.966973243847942,
+      "distinct_1": 0.09244218671549662,
+      "distinct_2": 0.4732998630762209,
+      "rep_4gram": 0.004260111395545372,
+      "avg_words": 33.859,
       "name": "topk_k20",
       "strategy": "top_k",
       "params": {
@@ -167,24 +167,24 @@
         "strategy": "top_k",
         "k": 20
       },
-      "time_s": 24.941646337509155,
-      "ms_per_prompt": 24.941646337509155,
-      "gen_tokens_mean": 69.067,
-      "judge_ppl_median": 14.630431175231934,
-      "judge_ppl_mean": 24.414003211975096
+      "time_s": 27.953922510147095,
+      "ms_per_prompt": 27.953922510147095,
+      "gen_tokens_mean": 69.165,
+      "judge_ppl_median": 14.761709213256836,
+      "judge_ppl_mean": 50.361770061254504
     },
     {
-      "precision": 0.18100045337501575,
-      "recall": 0.16393725762403472,
-      "f1": 0.15556260822373996,
-      "accuracy": 0.010887085233937758,
-      "bleu": 0.40375922428576166,
-      "self_ppl_mean": 14.773549448869355,
-      "self_ppl_median": 13.841562918287725,
-      "distinct_1": 0.10454026270702456,
-      "distinct_2": 0.5344305657604702,
-      "rep_4gram": 0.0027486902217292155,
-      "avg_words": 35.02,
+      "precision": 0.17549255996014995,
+      "recall": 0.16199041954465634,
+      "f1": 0.15203213911919755,
+      "accuracy": 0.011857919480724241,
+      "bleu": 0.44096547944277925,
+      "self_ppl_mean": 15.653030294977679,
+      "self_ppl_median": 13.604209219791796,
+      "distinct_1": 0.106938636169668,
+      "distinct_2": 0.5367388758782201,
+      "rep_4gram": 0.002577424448111978,
+      "avg_words": 35.151,
       "name": "topk_k50",
       "strategy": "top_k",
       "params": {
@@ -192,24 +192,24 @@
         "strategy": "top_k",
         "k": 50
       },
-      "time_s": 24.886651039123535,
-      "ms_per_prompt": 24.886651039123535,
-      "gen_tokens_mean": 68.922,
-      "judge_ppl_median": 21.331963539123535,
-      "judge_ppl_mean": 501.72633791828156
+      "time_s": 27.92880129814148,
+      "ms_per_prompt": 27.92880129814148,
+      "gen_tokens_mean": 68.635,
+      "judge_ppl_median": 20.48440647125244,
+      "judge_ppl_mean": 72.74511349225044
     },
     {
-      "precision": 0.20526678420165223,
-      "recall": 0.15309069901051064,
-      "f1": 0.15840480805376173,
-      "accuracy": 0.014216155649643127,
-      "bleu": 0.50328331736766,
-      "self_ppl_mean": 7.068113830843848,
-      "self_ppl_median": 6.912323629738417,
-      "distinct_1": 0.09306409130816505,
-      "distinct_2": 0.43222086318364494,
-      "rep_4gram": 0.02030865920372129,
-      "avg_words": 29.614,
+      "precision": 0.20069336987864045,
+      "recall": 0.14639978279849664,
+      "f1": 0.1519355828806502,
+      "accuracy": 0.01273002667646736,
+      "bleu": 0.5404882348932311,
+      "self_ppl_mean": 7.576500820452772,
+      "self_ppl_median": 6.75727633465444,
+      "distinct_1": 0.09244264507422402,
+      "distinct_2": 0.43183742591024554,
+      "rep_4gram": 0.021025283581782708,
+      "avg_words": 28.158,
       "name": "topk_k50_t0.7",
       "strategy": "top_k",
       "params": {
@@ -218,24 +218,24 @@
         "k": 50,
         "temperature": 0.7
       },
-      "time_s": 24.90457844734192,
-      "ms_per_prompt": 24.90457844734192,
-      "gen_tokens_mean": 69.718,
-      "judge_ppl_median": 10.698809146881104,
-      "judge_ppl_mean": 15.45937124478817
+      "time_s": 27.997068405151367,
+      "ms_per_prompt": 27.997068405151367,
+      "gen_tokens_mean": 69.563,
+      "judge_ppl_median": 10.548301219940186,
+      "judge_ppl_mean": 1104.2365723810196
     },
     {
-      "precision": 0.20739706440864433,
-      "recall": 0.1499013118788447,
-      "f1": 0.15918566183599664,
-      "accuracy": 0.014822320849886555,
-      "bleu": 0.4478536424658865,
-      "self_ppl_mean": 6.073593265666012,
-      "self_ppl_median": 5.982498789525279,
-      "distinct_1": 0.09500420993544766,
-      "distinct_2": 0.41255817335660266,
-      "rep_4gram": 0.03136431125812165,
-      "avg_words": 28.504,
+      "precision": 0.21224675405441254,
+      "recall": 0.14685590502611967,
+      "f1": 0.15687408460192412,
+      "accuracy": 0.013971389777601836,
+      "bleu": 0.5397175501288309,
+      "self_ppl_mean": 6.038022332522205,
+      "self_ppl_median": 5.858416877967823,
+      "distinct_1": 0.1009201217053411,
+      "distinct_2": 0.4334855403348554,
+      "rep_4gram": 0.02484618782785406,
+      "avg_words": 27.279,
       "name": "topp_p0.5",
       "strategy": "top_p",
       "params": {
@@ -243,24 +243,24 @@
         "strategy": "top_p",
         "p": 0.5
       },
-      "time_s": 27.650500059127808,
-      "ms_per_prompt": 27.650500059127808,
-      "gen_tokens_mean": 69.866,
-      "judge_ppl_median": 9.648499011993408,
-      "judge_ppl_mean": 13.219340204000472
+      "time_s": 31.63517665863037,
+      "ms_per_prompt": 31.635176658630375,
+      "gen_tokens_mean": 69.823,
+      "judge_ppl_median": 9.4617280960083,
+      "judge_ppl_mean": 1095.229723884821
     },
     {
-      "precision": 0.1786455293948896,
-      "recall": 0.15898556161207772,
-      "f1": 0.15182183169853083,
-      "accuracy": 0.01125190065874159,
-      "bleu": 0.22743045693599412,
-      "self_ppl_mean": 16.296575027955598,
-      "self_ppl_median": 14.50375505361637,
-      "distinct_1": 0.13724910579680155,
-      "distinct_2": 0.5839322694603484,
-      "rep_4gram": 0.003038437207161451,
-      "avg_words": 33.829,
+      "precision": 0.1754558860783798,
+      "recall": 0.15128935460540213,
+      "f1": 0.14627172467982152,
+      "accuracy": 0.010996779233892341,
+      "bleu": 0.3233499868741911,
+      "self_ppl_mean": 17.14533789994608,
+      "self_ppl_median": 14.304176726370176,
+      "distinct_1": 0.13880592419250545,
+      "distinct_2": 0.5860382392174299,
+      "rep_4gram": 0.003378257483535917,
+      "avg_words": 32.477,
       "name": "topp_p0.8",
       "strategy": "top_p",
       "params": {
@@ -268,24 +268,24 @@
         "strategy": "top_p",
         "p": 0.8
       },
-      "time_s": 27.59784722328186,
-      "ms_per_prompt": 27.59784722328186,
-      "gen_tokens_mean": 69.023,
-      "judge_ppl_median": 23.389395713806152,
-      "judge_ppl_mean": 4002.7205463142395
+      "time_s": 31.48725962638855,
+      "ms_per_prompt": 31.48725962638855,
+      "gen_tokens_mean": 69.009,
+      "judge_ppl_median": 23.211604118347168,
+      "judge_ppl_mean": 74.37143414640427
     },
     {
-      "precision": 0.15578726318266997,
-      "recall": 0.14156256433661007,
-      "f1": 0.13438127406575798,
-      "accuracy": 0.008718096877912114,
-      "bleu": 0.16851247100645653,
-      "self_ppl_mean": 40.52525353706895,
-      "self_ppl_median": 34.312852488083806,
-      "distinct_1": 0.18865485737786786,
-      "distinct_2": 0.7045688350983359,
-      "rep_4gram": 0.0005719175842843666,
-      "avg_words": 34.041,
+      "precision": 0.15602106748659775,
+      "recall": 0.14042348676499086,
+      "f1": 0.13287360175699992,
+      "accuracy": 0.00936425787467936,
+      "bleu": 0.16536685767069867,
+      "self_ppl_mean": 39.594243798258624,
+      "self_ppl_median": 33.30959383748847,
+      "distinct_1": 0.18891279361056262,
+      "distinct_2": 0.701587687076895,
+      "rep_4gram": 0.0007256760411977545,
+      "avg_words": 33.931,
       "name": "topp_p0.95",
       "strategy": "top_p",
       "params": {
@@ -293,24 +293,24 @@
         "strategy": "top_p",
         "p": 0.95
       },
-      "time_s": 27.591065883636475,
-      "ms_per_prompt": 27.591065883636475,
-      "gen_tokens_mean": 68.753,
-      "judge_ppl_median": 58.844425201416016,
-      "judge_ppl_mean": 4085.3288768868447
+      "time_s": 31.551427125930786,
+      "ms_per_prompt": 31.551427125930786,
+      "gen_tokens_mean": 68.636,
+      "judge_ppl_median": 54.274309158325195,
+      "judge_ppl_mean": 297.4285540444851
     },
     {
-      "precision": 0.20651589086843022,
-      "recall": 0.1484358064877369,
-      "f1": 0.15685438845860386,
-      "accuracy": 0.014285365431144705,
-      "bleu": 0.6053996783877227,
-      "self_ppl_mean": 5.729693244652217,
-      "self_ppl_median": 5.577883722137798,
-      "distinct_1": 0.09334971047552203,
-      "distinct_2": 0.4059283506092017,
-      "rep_4gram": 0.03475423802329605,
-      "avg_words": 28.495,
+      "precision": 0.19970724391003492,
+      "recall": 0.14364941747564147,
+      "f1": 0.15108641777377704,
+      "accuracy": 0.01395932104782926,
+      "bleu": 0.5577671411688147,
+      "self_ppl_mean": 5.73871449144444,
+      "self_ppl_median": 5.589571006013263,
+      "distinct_1": 0.09601266733841946,
+      "distinct_2": 0.408413901228116,
+      "rep_4gram": 0.034714608915333796,
+      "avg_words": 27.788,
       "name": "topp_p0.9_t0.7",
       "strategy": "top_p",
       "params": {
@@ -319,92 +319,92 @@
         "p": 0.9,
         "temperature": 0.7
       },
-      "time_s": 27.645728826522827,
-      "ms_per_prompt": 27.645728826522827,
-      "gen_tokens_mean": 69.939,
-      "judge_ppl_median": 9.250245571136475,
-      "judge_ppl_mean": 11.254222588062287
+      "time_s": 31.56927514076233,
+      "ms_per_prompt": 31.569275140762326,
+      "gen_tokens_mean": 69.865,
+      "judge_ppl_median": 9.213471412658691,
+      "judge_ppl_mean": 1094.454825058937
     }
   ],
   "beam_timing": [
     {
       "mode": "natural",
       "width": 1,
-      "mean_s_per_prompt": 0.6848450946807861,
-      "std_s": 0.03387934526472246,
-      "mean_ms_per_token": 9.78350135258266,
+      "mean_s_per_prompt": 0.781981463432312,
+      "std_s": 0.01558768772931337,
+      "mean_ms_per_token": 11.171163763318747,
       "mean_output_tokens": 70.0,
-      "total_s": 68.48450946807861,
+      "total_s": 78.1981463432312,
       "n_prompts": 100
     },
     {
       "mode": "natural",
       "width": 2,
-      "mean_s_per_prompt": 0.6798255443572998,
-      "std_s": 0.009905260321141293,
-      "mean_ms_per_token": 9.7313968216702,
+      "mean_s_per_prompt": 0.793917977809906,
+      "std_s": 0.022438927991426688,
+      "mean_ms_per_token": 11.364886324126998,
       "mean_output_tokens": 69.86,
-      "total_s": 67.98255443572998,
+      "total_s": 79.3917977809906,
       "n_prompts": 100
     },
     {
       "mode": "natural",
       "width": 4,
-      "mean_s_per_prompt": 0.6822133040428162,
-      "std_s": 0.027737611702376378,
-      "mean_ms_per_token": 9.803958435890708,
+      "mean_s_per_prompt": 0.8327497005462646,
+      "std_s": 0.033384079158313755,
+      "mean_ms_per_token": 11.970075382877122,
       "mean_output_tokens": 69.59,
-      "total_s": 68.22133040428162,
+      "total_s": 83.27497005462646,
       "n_prompts": 100
     },
     {
       "mode": "natural",
       "width": 8,
-      "mean_s_per_prompt": 0.7037396359443665,
-      "std_s": 0.03137733007915173,
-      "mean_ms_per_token": 10.1683868747084,
+      "mean_s_per_prompt": 0.8716441249847412,
+      "std_s": 0.043158742726592765,
+      "mean_ms_per_token": 12.592297802702353,
       "mean_output_tokens": 69.34,
-      "total_s": 70.37396359443665,
+      "total_s": 87.16441249847412,
       "n_prompts": 100
     },
     {
       "mode": "fixed",
       "width": 1,
-      "mean_s_per_prompt": 0.6784237551689148,
-      "std_s": 0.0016587773214447714,
-      "mean_ms_per_token": 9.691767930984499,
+      "mean_s_per_prompt": 0.7738243556022644,
+      "std_s": 0.01421080819544866,
+      "mean_ms_per_token": 11.054633651460922,
       "mean_output_tokens": 70.0,
-      "total_s": 67.84237551689148,
+      "total_s": 77.38243556022644,
       "n_prompts": 100
     },
     {
       "mode": "fixed",
       "width": 2,
-      "mean_s_per_prompt": 0.6811458969116211,
-      "std_s": 0.001066027939126381,
-      "mean_ms_per_token": 9.730655670166016,
+      "mean_s_per_prompt": 0.7971447992324829,
+      "std_s": 0.02233234077874458,
+      "mean_ms_per_token": 11.387782846178329,
       "mean_output_tokens": 70.0,
-      "total_s": 68.11458969116211,
+      "total_s": 79.71447992324829,
       "n_prompts": 100
     },
     {
       "mode": "fixed",
       "width": 4,
-      "mean_s_per_prompt": 0.6865559673309326,
-      "std_s": 0.0012468601240148754,
-      "mean_ms_per_token": 9.807942390441896,
+      "mean_s_per_prompt": 0.8367985081672669,
+      "std_s": 0.013029678554817858,
+      "mean_ms_per_token": 11.954264402389526,
       "mean_output_tokens": 70.0,
-      "total_s": 68.65559673309326,
+      "total_s": 83.67985081672668,
       "n_prompts": 100
     },
     {
       "mode": "fixed",
       "width": 8,
-      "mean_s_per_prompt": 0.7092588877677918,
-      "std_s": 0.0009664558877848735,
-      "mean_ms_per_token": 10.132269825254168,
+      "mean_s_per_prompt": 0.8789880871772766,
+      "std_s": 0.012277801042479799,
+      "mean_ms_per_token": 12.556972673961095,
       "mean_output_tokens": 70.0,
-      "total_s": 70.92588877677917,
+      "total_s": 87.89880871772766,
       "n_prompts": 100
     }
   ],
@@ -491,13 +491,10 @@
     "timing_samples": 100,
     "judge_model": "EleutherAI/pythia-410m",
     "use_bertscore": true,
-    "out_dir": "/home2/ojas.k/a2/outputs/decoding",
+    "out_dir": "/kaggle/working/outputs/decoding",
     "hf_user": null,
     "repo_name": "anlp-a2-decoding",
     "push_to_hub": true,
-    "private": false,
-    "wandb": true,
-    "wandb_project": "anlp-a2",
-    "wandb_entity": null
+    "private": false
   }
 }