Make bigram interpolation work for training and unseen transitions

.github/workflows/tests.yml

@@ -0,0 +1,17 @@
+name: N-gram correctness
+on: [push, pull_request]
+permissions:
+  contents: read
+jobs:
+  tests:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python: ['3.12', '3.13']
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python }}
+      - run: python -m unittest -v
+      - run: python demo.py --text 'the agent tests the system'

README.md

@@ -21,7 +21,7 @@
 - **Trigram (3-gram)**: Extended context modeling with two preceding words.  
 
 ### **Smoothing Techniques**  
-- **Add-k Smoothing**: Mitigates zero-probability issues by adding constant `k` to counts.  
+- **Add-k Smoothing**: Not implemented in the extracted offline module.
 - **Linear Interpolation**: Blends probabilities across n-gram orders using weights (`λ₁`, `λ₂`, `λ₃`).  
 
 ### **Preprocessing**  
@@ -39,9 +39,9 @@
 ### **Class Hierarchy**  
 | Class       | Description                                                                 |  
 |-------------|-----------------------------------------------------------------------------|  
-| `Unigram`   | Base frequency model with add-k smoothing support.                          |  
-| `Bigram`    | Inherits from `Unigram`; implements bigram MLE and interpolation.           |  
-| `Trigram`   | Inherits from `Bigram`; calculates trigram probabilities with backoff.      |  
+| `Unigram`   | Base frequency model; no lower-order interpolation or add-k implementation. |
+| `Bigram`    | Holds a trained `Unigram`; implements bigram MLE and interpolation.          |
+| `Trigram`   | Holds a trained `Bigram`; calculates trigram probabilities with backoff.     |
 
 ### **Core Methods**  
 - `train_regular()`: Trains vanilla n-gram models using Maximum Likelihood Estimation (MLE).  
@@ -64,8 +64,11 @@
 ### **Training**  
 ```python  
 # Conceptual Example  
-train_data = load_data("corpus.txt")  
-model = train_interpolation_trigram(train_data, l1=0.3, l2=0.3, l3=0.4)  
+from ngram import load_data, tokenize, get_pretrained_bigram, train_interpolation_trigram
+train_data = tokenize(load_data("corpus.txt"))
+test_data = tokenize(load_data("test.txt"))
+pretrained = get_pretrained_bigram(train_data)
+model = train_interpolation_trigram(train_data, pretrained, l1=0.3, l2=0.3, l3=0.4)
 print(f"Test Perplexity: {model.perplexity(test_data):.2f}")  
 ```  
 
@@ -109,3 +112,19 @@ python -m unittest -v
 ```
 
 Supply `--train path/to/sentences.txt` for your own corpus. The built-in corpus is an explicitly labeled toy fixture, not a trained production model or a reproduction of the original benchmark. The notebook still requires the missing `NLP/1b_benchmark.*.tokens` files for its historical experiment. Zero probabilities use the original notebook floor of 1e-10; reported perplexities are therefore floored scores.
+
+### Bigram interpolation example
+
+```python
+from ngram import Bigram, tokenize, train_regular
+corpus = tokenize(["a b", "a c"])
+model = Bigram()
+model.set_unigram(train_regular(corpus, 1))
+model.train_bigram(corpus, smoothing=True, l1=0.3, l2=0.7)
+print(model.perplexity(tokenize(["a a"])))  # 5.0 (up to floating-point rounding)
+```
+
+The weights must be finite, nonnegative, and sum to one. Unseen transitions
+receive their unigram component when smoothing is enabled; unknown words with
+no trained `<UNK>` mass still use the existing numerical floor. These toy checks
+validate the formulas and runnable API, not the historical benchmark results.

ngram.py

@@ -118,7 +118,9 @@ def perplexity(self, corpus):
     entropy = 0
     for sent in processed_corpus:
       # if not exist(unseen bigram), just give 0, and we will handle it in the log_joint_probs function
-      probs = [self.vocab_prob.get(bigram, 0) for bigram in self.get_bigrams(sent)]
+      probs = [self.vocab_prob[bigram] if bigram in self.vocab_prob else
+               self.MLE(bigram, self.smoothing, self.l1, self.l2)
+               for bigram in self.get_bigrams(sent)]
       entropy += self.log_joint_probs(probs)
     l = entropy / m
     return pow(2, -l)
@@ -130,15 +132,21 @@ def log_joint_probs(self, probs):
   def handle_zero_prob(self):
     return pow(10, -10)
   def train_bigram(self, train_corpus, smoothing=None, l1=0, l2=1):
+      if self.unigram is None:
+        raise ValueError("Set a trained unigram before training a bigram")
+      if smoothing and (not all(math.isfinite(w) and w >= 0 for w in (l1, l2))
+                        or not math.isclose(l1 + l2, 1)):
+        raise ValueError("Interpolation weights must be finite, nonnegative, and sum to 1")
       self.smoothing = smoothing
       self.l1 = l1
       self.l2 = l2
       self.vocab_freq = get_vocab_freq([self.get_bigrams(sent) for sent in train_corpus])
+      self.vocab_prob = {}
       for bigram in self.vocab_freq.keys():
         self.vocab_prob[bigram] = self.MLE(bigram, smoothing, l1, l2)
   def MLE(self, bigram, smoothing=None, l1=0, l2=1):
     if smoothing : 
-      return self.bigram_prob_with_interpolation(bigram, l1, l2)
+      return self.bigram_prob_interpolation(bigram, l1, l2)
     else:
       return self.bigram_prob(bigram)
   def get_bigrams(self, sent):
@@ -148,9 +156,10 @@ def vocab_size(self):
     # excluding the <START> token
     return len(self.unigram.vocab_freq.keys()) - 1
   def bigram_prob(self, bigram):
-    return (self.vocab_freq.get(bigram, 0) / self.unigram.vocab_freq[bigram[0]])
+    count = self.unigram.vocab_freq.get(bigram[0], 0)
+    return self.vocab_freq.get(bigram, 0) / count if count else 0
   def bigram_prob_interpolation(self, bigram, l1=0, l2=1):
-    return l2 * self.bigram_prob(bigram) + l1 * self.unigram.word_prob(bigram[1])
+    return l2 * self.bigram_prob(bigram) + l1 * self.unigram.vocab_prob.get(bigram[1], 0)
 
 
 class Trigram:

test_interpolation.py

@@ -0,0 +1,74 @@
+import math
+import unittest
+
+from ngram import Bigram, START_TOKEN, tokenize, train_regular, train_interpolation_trigram
+
+
+class BigramInterpolationTests(unittest.TestCase):
+    def setUp(self):
+        self.corpus = tokenize(["a b", "a c"])
+        self.unigram = train_regular(self.corpus, 1)
+
+    def smoothed(self, l1=.3, l2=.7):
+        model = Bigram()
+        model.set_unigram(self.unigram)
+        model.train_bigram(self.corpus, smoothing=True, l1=l1, l2=l2)
+        return model
+
+    def test_training_and_seen_probabilities(self):
+        model = self.smoothed()
+        self.assertAlmostEqual(model.vocab_prob[(START_TOKEN, "a")], .8)
+        self.assertAlmostEqual(model.vocab_prob[("a", "b")], .4)
+
+    def test_unseen_transition_uses_unigram_component(self):
+        model = self.smoothed()
+        self.assertAlmostEqual(model.MLE(("a", "a"), True, .3, .7), .1)
+        # P(a|START)=.8, P(a|a)=.1, P(STOP|a)=.1 => perplexity=5.
+        self.assertAlmostEqual(model.perplexity(tokenize(["a a"])), 5)
+
+    def test_observed_context_probability_mass(self):
+        model = self.smoothed()
+        for context in (START_TOKEN, "a", "b", "c"):
+            mass = sum(model.MLE((context, word), True, .3, .7)
+                       for word in self.unigram.vocab_prob)
+            self.assertAlmostEqual(mass, 1)
+
+    def test_unsmoothed_scores_preserved(self):
+        model = train_regular(self.corpus, 2)
+        self.assertAlmostEqual(model.perplexity(tokenize(["a b"])), 2 ** (1 / 3))
+        interpolated = self.smoothed(0, 1)
+        for sentence in ("a b", "a a", "unseen a"):
+            self.assertAlmostEqual(model.perplexity(tokenize([sentence])),
+                                   interpolated.perplexity(tokenize([sentence])))
+
+    def test_unknown_context_and_word_do_not_raise(self):
+        model = self.smoothed()
+        self.assertTrue(math.isfinite(model.perplexity(tokenize(["unseen a"]))))
+        self.assertAlmostEqual(model.MLE(("not-a-context", "a"), True, .3, .7), .1)
+
+    def test_missing_unigram_rejected_clearly(self):
+        with self.assertRaisesRegex(ValueError, "trained unigram"):
+            Bigram().train_bigram(self.corpus, smoothing=True, l1=.3, l2=.7)
+
+    def test_invalid_mixture_weights_rejected(self):
+        for l1, l2 in ((-.1, 1.1), (.3, .3), (float("nan"), .7), (.3, float("inf"))):
+            with self.subTest(l1=l1, l2=l2):
+                with self.assertRaises(ValueError):
+                    self.smoothed(l1, l2)
+
+    def test_retraining_discards_old_probability_cache(self):
+        model = self.smoothed()
+        corpus = tokenize(["x y"])
+        model.set_unigram(train_regular(corpus, 1))
+        model.train_bigram(corpus, smoothing=True, l1=.3, l2=.7)
+        self.assertNotIn(("a", "b"), model.vocab_prob)
+
+    def test_trigram_shared_interpolation_path_preserved(self):
+        bigram = train_regular(self.corpus, 2)
+        model = train_interpolation_trigram(self.corpus, bigram, .3, .3, .4)
+        self.assertAlmostEqual(model.perplexity(tokenize(["a b"])),
+                               (.8 * .4 * .8) ** (-1 / 3))
+
+
+if __name__ == "__main__":
+    unittest.main()