From f82ba37957fb36b3db8ce7e567e0a230e02b815b Mon Sep 17 00:00:00 2001 From: Sylvain Le Groux Date: Tue, 16 Jun 2026 16:59:55 -0700 Subject: [PATCH] feat(lm): add perplexity & batchnorm; fix NNLM_X init; auto-close workflow models.lm: - add reusable `perplexity(logits, targets)` = exp(mean cross-entropy) and log val/test perplexity per-epoch in NNLM_X (closes #64) - add optional BatchNorm1d on the NNLM hidden layer via `batch_norm` flag on NNLM/NNLMConfig (default off = backward compatible); sample() now eval()s so BN uses running stats on single-token contexts (closes #16) - fix NNLM_X.__init__: it forwarded args to Classifier positionally in the wrong order (nnet got the int num_classes), so NNLM_X could not be instantiated at all and its notebook test cell was failing CI; now passes keyword args - surface batch_norm in config/text/model/nnlm.yaml CI: - add close-linked-issues workflow: on merged PRs, close issues referenced by closing keywords on any base branch (GitHub only does this for the default branch) (closes #42) Verified: nbdev_test models.lm.ipynb passes; perplexity(uniform, V)=V and ~1.0 for near-perfect; full Lightning train/val/test logs test/perplexity=exp(loss). Co-Authored-By: Claude Opus 4.8 (1M context) --- .github/workflows/close-linked-issues.yaml | 62 ++++++++++++++++++++++ config/text/model/nnlm.yaml | 1 + nbs/models.lm.ipynb | 61 +++++++++++++++------ nimrod/_modidx.py | 7 ++- nimrod/models/lm.py | 50 ++++++++++++----- 5 files changed, 152 insertions(+), 29 deletions(-) create mode 100644 .github/workflows/close-linked-issues.yaml diff --git a/.github/workflows/close-linked-issues.yaml b/.github/workflows/close-linked-issues.yaml new file mode 100644 index 0000000..3b73450 --- /dev/null +++ b/.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: #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}`); + } + } diff --git a/config/text/model/nnlm.yaml b/config/text/model/nnlm.yaml index 9e5ee46..ae3d5ca 100644 --- a/config/text/model/nnlm.yaml +++ b/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 diff --git a/nbs/models.lm.ipynb b/nbs/models.lm.ipynb index c668011..d89bf49 100644 --- a/nbs/models.lm.ipynb +++ b/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())" ] }, { diff --git a/nimrod/_modidx.py b/nimrod/_modidx.py index cc8a7b2..99bbd8c 100644 --- a/nimrod/_modidx.py +++ b/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'), diff --git a/nimrod/models/lm.py b/nimrod/models/lm.py index 926e6d9..15eb5a3 100644 --- a/nimrod/models/lm.py +++ b/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