feat(lm): perplexity + batchnorm, fix NNLM_X init, auto-close workflow

.github/workflows/close-linked-issues.yaml

@@ -0,0 +1,62 @@
+name: Close linked issues
+
+# When a PR is merged, close every issue it references with a closing keyword
+# (close/closes/closed, fix/fixes/fixed, resolve/resolves/resolved).
+#
+# GitHub already does this automatically — but ONLY when the PR merges into the
+# repository's default branch. This workflow extends the behavior to PRs merged
+# into any base branch (e.g. feature/integration branches) so the issue tracker
+# stays in sync regardless of merge target.
+
+on:
+  pull_request:
+    types: [closed]
+
+permissions:
+  issues: write
+  contents: read
+
+jobs:
+  close-issues:
+    # only act on PRs that were actually merged (not just closed)
+    if: github.event.pull_request.merged == true
+    runs-on: ubuntu-latest
+    steps:
+      - name: Close issues referenced by closing keywords
+        uses: actions/github-script@v7
+        with:
+          script: |
+            const pr = context.payload.pull_request;
+            const text = `${pr.title}\n\n${pr.body || ''}`;
+
+            // match: <keyword> #123  (keyword is case-insensitive)
+            const re = /\b(?:close[sd]?|fix(?:e[sd])?|resolve[sd]?)\b[:\s]+#(\d+)/gi;
+            const numbers = new Set();
+            for (const m of text.matchAll(re)) numbers.add(Number(m[1]));
+
+            if (numbers.size === 0) {
+              core.info('No closing-keyword issue references found in this PR.');
+              return;
+            }
+
+            for (const issue_number of numbers) {
+              try {
+                const { data: issue } = await github.rest.issues.get({
+                  ...context.repo, issue_number,
+                });
+                if (issue.state === 'closed') {
+                  core.info(`#${issue_number} already closed — skipping.`);
+                  continue;
+                }
+                await github.rest.issues.createComment({
+                  ...context.repo, issue_number,
+                  body: `Closed by #${pr.number} (merged into \`${pr.base.ref}\`).`,
+                });
+                await github.rest.issues.update({
+                  ...context.repo, issue_number, state: 'closed',
+                });
+                core.info(`Closed #${issue_number}.`);
+              } catch (err) {
+                core.warning(`Could not close #${issue_number}: ${err.message}`);
+              }
+            }

config/text/model/nnlm.yaml

@@ -7,6 +7,7 @@ nnet:
   n_emb: 10
   n_context: 3
   n_h: 100
+  batch_norm: false # set true to add batchnorm on the hidden layer
 
 optimizer:
   _target_: torch.optim.Adam

nbs/models.lm.ipynb

@@ -3092,13 +3092,15 @@
     "    n_emb:int = 10\n",
     "    n_context:int = 3\n",
     "    n_h:int = 100\n",
+    "    batch_norm:bool = False\n",
     "\n",
     "class NNLM(nn.Module):\n",
     "    def __init__(self,\n",
-    "                n_vocab:int = 30, # vocabulary size \n",
+    "                n_vocab:int = 30, # vocabulary size\n",
     "                n_emb:int = 10, # embedding dimension\n",
     "                n_context:int = 3, # context size bigram/trigram, etc.\n",
-    "                n_h:int = 100 # hidden layer size\n",
+    "                n_h:int = 100, # hidden layer size\n",
+    "                batch_norm:bool = False # add batchnorm on the hidden layer\n",
     "                ):\n",
     "\n",
     "        logger.info(f\"NNLM: Init\")\n",
@@ -3108,22 +3110,25 @@
     "        self.n_emb = n_emb\n",
     "        self.n_context = n_context\n",
     "        # we concatenate input of [n_context length, n_emb] into linear layer (T*C):\n",
-    "        self.l1 = nn.Linear(n_context * n_emb, n_h) \n",
+    "        self.l1 = nn.Linear(n_context * n_emb, n_h)\n",
+    "        # optional batchnorm stabilizes hidden activations & speeds up training (c.f. makemore)\n",
+    "        self.bn = nn.BatchNorm1d(n_h) if batch_norm else nn.Identity()\n",
     "        self.l2 = nn.Linear(n_h, n_vocab)\n",
-    "    \n",
+    "\n",
     "    def forward(self, x:torch.Tensor)->torch.Tensor:\n",
     "        # input: (B,T)\n",
     "        embedding = self.embedder(x) # ->(B,T,C)\n",
     "        # we concatenate input of n_context length * n_emb (T*C) into linear layer:\n",
     "        h = self.l1(embedding.view(-1,self.n_context * self.n_emb))\n",
+    "        h = self.bn(h)\n",
     "        h = torch.tanh(h)\n",
     "        logits = self.l2(h)\n",
     "        return(logits)\n",
     "\n",
     "    @torch.no_grad()\n",
     "    def sample(self, prompt:str, vocab:Vocab, max_new_tokens:int=50, temperature:float=1.0):\n",
-    "\n",
-    "        for _ in range(max_new_tokens):            \n",
+    "        self.eval() # use batchnorm running stats (single-token contexts) & disable train-only layers\n",
+    "        for _ in range(max_new_tokens):\n",
     "            # limit prompt to context size\n",
     "            context = prompt[-self.n_context:]\n",
     "            context = vocab.stoi(list(context))\n",
@@ -3902,7 +3907,16 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "#| export \n",
+    "#| export\n",
+    "\n",
+    "def perplexity(\n",
+    "        logits:torch.Tensor, # (N, n_vocab) unnormalized scores\n",
+    "        targets:torch.Tensor, # (N,) target token ids\n",
+    "        ignore_index:int=-100 # target value to ignore (e.g. padding)\n",
+    "        )->torch.Tensor:\n",
+    "    \"Perplexity = exp(mean cross-entropy) of `targets` under `logits`. Lower is better; equals vocab size for a uniform model.\"\n",
+    "    ce = F.cross_entropy(logits, targets, ignore_index=ignore_index)\n",
+    "    return torch.exp(ce)\n",
     "\n",
     "class NNLM_X(Classifier, LightningModule):\n",
     "    def __init__(\n",
@@ -3915,18 +3929,19 @@
     "\n",
     "        logger.info(\"NNLM_X: Init\")\n",
     "        super().__init__(\n",
-    "            num_classes,\n",
-    "            optimizer,\n",
-    "            scheduler,\n",
+    "            nnet=nnet,\n",
+    "            num_classes=num_classes,\n",
+    "            optimizer=optimizer,\n",
+    "            scheduler=scheduler,\n",
     "            )\n",
-    "        self.save_hyperparameters(logger=False)\n",
+    "        self.save_hyperparameters(logger=False, ignore=['nnet'])\n",
     "        # required attribute for lr finder\n",
     "        self.lr = optimizer.keywords['lr']\n",
     "        self.nnet = nnet\n",
-    "    \n",
+    "\n",
     "    def forward(self, x:torch.Tensor)->torch.Tensor:\n",
     "        return self.nnet(x)\n",
-    "    \n",
+    "\n",
     "    def _step(self, batch, batch_idx):\n",
     "        x, y = batch\n",
     "        y = y[:, -1]\n",
@@ -3935,12 +3950,21 @@
     "        preds = y_hat.argmax(dim=1)\n",
     "        return loss, preds, y\n",
     "\n",
+    "    def on_validation_epoch_end(self) -> None:\n",
+    "        # perplexity = exp(epoch-mean cross-entropy); val_loss already tracks mean CE\n",
+    "        super().on_validation_epoch_end()\n",
+    "        self.log(\"val/perplexity\", torch.exp(self.val_loss.compute()), prog_bar=True, sync_dist=True)\n",
+    "\n",
+    "    def on_test_epoch_end(self) -> None:\n",
+    "        super().on_test_epoch_end()\n",
+    "        self.log(\"test/perplexity\", torch.exp(self.test_loss.compute()), prog_bar=True, sync_dist=True)\n",
+    "\n",
     "    def predict_step(self, batch, batch_idx, dataloader_idx=0):\n",
     "        x, y = batch\n",
     "        y = y[:, -1]  # Get the last token as target\n",
     "        y_hat = self.forward(x)\n",
     "        return y_hat.argmax(dim=1)\n",
-    "    \n",
+    "\n",
     "    def sample(self, prompt:str, vocab:Vocab, max_new_tokens:int=50, temperature:float=1.0):\n",
     "        return self.nnet.sample(prompt, vocab, max_new_tokens, temperature)"
    ]
@@ -4030,7 +4054,14 @@
    "outputs": [],
    "source": [
     "y = lm(x)\n",
-    "print(\"Y_hat logits:\", y.shape)"
+    "print(\"Y_hat logits:\", y.shape)\n",
+    "\n",
+    "# perplexity sanity check: a uniform distribution over V classes has perplexity == V\n",
+    "from nimrod.models.lm import perplexity\n",
+    "_V = 50\n",
+    "_ppl = perplexity(torch.zeros(8, _V), torch.randint(0, _V, (8,)))\n",
+    "assert torch.isclose(_ppl, torch.tensor(float(_V)), atol=1e-4), _ppl\n",
+    "print(\"perplexity(uniform, V=50) =\", _ppl.item())"
    ]
   },
   {

nimrod/_modidx.py

@@ -529,8 +529,13 @@
                                   'nimrod.models.lm.NNLM_X.__init__': ('models.lm.html#nnlm_x.__init__', 'nimrod/models/lm.py'),
                                   'nimrod.models.lm.NNLM_X._step': ('models.lm.html#nnlm_x._step', 'nimrod/models/lm.py'),
                                   'nimrod.models.lm.NNLM_X.forward': ('models.lm.html#nnlm_x.forward', 'nimrod/models/lm.py'),
+                                  'nimrod.models.lm.NNLM_X.on_test_epoch_end': ( 'models.lm.html#nnlm_x.on_test_epoch_end',
+                                                                                 'nimrod/models/lm.py'),
+                                  'nimrod.models.lm.NNLM_X.on_validation_epoch_end': ( 'models.lm.html#nnlm_x.on_validation_epoch_end',
+                                                                                       'nimrod/models/lm.py'),
                                   'nimrod.models.lm.NNLM_X.predict_step': ('models.lm.html#nnlm_x.predict_step', 'nimrod/models/lm.py'),
-                                  'nimrod.models.lm.NNLM_X.sample': ('models.lm.html#nnlm_x.sample', 'nimrod/models/lm.py')},
+                                  'nimrod.models.lm.NNLM_X.sample': ('models.lm.html#nnlm_x.sample', 'nimrod/models/lm.py'),
+                                  'nimrod.models.lm.perplexity': ('models.lm.html#perplexity', 'nimrod/models/lm.py')},
             'nimrod.models.mlp': { 'nimrod.models.mlp.MLP': ('models.mlp.html#mlp', 'nimrod/models/mlp.py'),
                                    'nimrod.models.mlp.MLP.__init__': ('models.mlp.html#mlp.__init__', 'nimrod/models/mlp.py'),
                                    'nimrod.models.mlp.MLP.forward': ('models.mlp.html#mlp.forward', 'nimrod/models/mlp.py'),

nimrod/models/lm.py

@@ -1,7 +1,7 @@
 # AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/models.lm.ipynb.
 
 # %% auto 0
-__all__ = ['logger', 'NNLMConfig', 'NNLM', 'NNLM_X', 'NNBigram']
+__all__ = ['logger', 'NNLMConfig', 'NNLM', 'perplexity', 'NNLM_X', 'NNBigram']
 
 # %% ../../nbs/models.lm.ipynb 4
 import torch.nn as nn
@@ -42,13 +42,15 @@ class NNLMConfig:
     n_emb:int = 10
     n_context:int = 3
     n_h:int = 100
+    batch_norm:bool = False
 
 class NNLM(nn.Module):
     def __init__(self,
-                n_vocab:int = 30, # vocabulary size 
+                n_vocab:int = 30, # vocabulary size
                 n_emb:int = 10, # embedding dimension
                 n_context:int = 3, # context size bigram/trigram, etc.
-                n_h:int = 100 # hidden layer size
+                n_h:int = 100, # hidden layer size
+                batch_norm:bool = False # add batchnorm on the hidden layer
                 ):
 
         logger.info(f"NNLM: Init")
@@ -58,22 +60,25 @@ def __init__(self,
         self.n_emb = n_emb
         self.n_context = n_context
         # we concatenate input of [n_context length, n_emb] into linear layer (T*C):
-        self.l1 = nn.Linear(n_context * n_emb, n_h) 
+        self.l1 = nn.Linear(n_context * n_emb, n_h)
+        # optional batchnorm stabilizes hidden activations & speeds up training (c.f. makemore)
+        self.bn = nn.BatchNorm1d(n_h) if batch_norm else nn.Identity()
         self.l2 = nn.Linear(n_h, n_vocab)
-    
+
     def forward(self, x:torch.Tensor)->torch.Tensor:
         # input: (B,T)
         embedding = self.embedder(x) # ->(B,T,C)
         # we concatenate input of n_context length * n_emb (T*C) into linear layer:
         h = self.l1(embedding.view(-1,self.n_context * self.n_emb))
+        h = self.bn(h)
         h = torch.tanh(h)
         logits = self.l2(h)
         return(logits)
 
     @torch.no_grad()
     def sample(self, prompt:str, vocab:Vocab, max_new_tokens:int=50, temperature:float=1.0):
-
-        for _ in range(max_new_tokens):            
+        self.eval() # use batchnorm running stats (single-token contexts) & disable train-only layers
+        for _ in range(max_new_tokens):
             # limit prompt to context size
             context = prompt[-self.n_context:]
             context = vocab.stoi(list(context))
@@ -86,6 +91,15 @@ def sample(self, prompt:str, vocab:Vocab, max_new_tokens:int=50, temperature:flo
         return(prompt)
 
 # %% ../../nbs/models.lm.ipynb 43
+def perplexity(
+        logits:torch.Tensor, # (N, n_vocab) unnormalized scores
+        targets:torch.Tensor, # (N,) target token ids
+        ignore_index:int=-100 # target value to ignore (e.g. padding)
+        )->torch.Tensor:
+    "Perplexity = exp(mean cross-entropy) of `targets` under `logits`. Lower is better; equals vocab size for a uniform model."
+    ce = F.cross_entropy(logits, targets, ignore_index=ignore_index)
+    return torch.exp(ce)
+
 class NNLM_X(Classifier, LightningModule):
     def __init__(
             self,
@@ -97,18 +111,19 @@ def __init__(
 
         logger.info("NNLM_X: Init")
         super().__init__(
-            num_classes,
-            optimizer,
-            scheduler,
+            nnet=nnet,
+            num_classes=num_classes,
+            optimizer=optimizer,
+            scheduler=scheduler,
             )
-        self.save_hyperparameters(logger=False)
+        self.save_hyperparameters(logger=False, ignore=['nnet'])
         # required attribute for lr finder
         self.lr = optimizer.keywords['lr']
         self.nnet = nnet
-    
+
     def forward(self, x:torch.Tensor)->torch.Tensor:
         return self.nnet(x)
-    
+
     def _step(self, batch, batch_idx):
         x, y = batch
         y = y[:, -1]
@@ -117,6 +132,15 @@ def _step(self, batch, batch_idx):
         preds = y_hat.argmax(dim=1)
         return loss, preds, y
 
+    def on_validation_epoch_end(self) -> None:
+        # perplexity = exp(epoch-mean cross-entropy); val_loss already tracks mean CE
+        super().on_validation_epoch_end()
+        self.log("val/perplexity", torch.exp(self.val_loss.compute()), prog_bar=True, sync_dist=True)
+
+    def on_test_epoch_end(self) -> None:
+        super().on_test_epoch_end()
+        self.log("test/perplexity", torch.exp(self.test_loss.compute()), prog_bar=True, sync_dist=True)
+
     def predict_step(self, batch, batch_idx, dataloader_idx=0):
         x, y = batch
         y = y[:, -1]  # Get the last token as target