|
"""N-gram Maximum Likelihood Probabilistic Language Model. |
|
|
|
Supports n-gram up to 5-gram. |
|
""" |
|
|
|
from __future__ import annotations |
|
|
|
import math as calc |
|
import os |
|
import pickle |
|
from collections import Counter, defaultdict |
|
from typing import List |
|
|
|
from nltk.util import ngrams |
|
|
|
|
|
class NGramMLE: |
|
def __init__(self, n: int = 1, corpus_file: str = None) -> None: |
|
"""Initializes the model. |
|
|
|
Args: |
|
n: n-gram order. Defaults to 1. |
|
corpus_file: File containing the corpus words. Defaults to None. |
|
|
|
Raises: |
|
FileNotFoundError: If corpus_file is defined but not found. |
|
""" |
|
if corpus_file and not os.path.exists(corpus_file): |
|
raise FileNotFoundError(f"Corpus file not found: {corpus_file}") |
|
|
|
self.n = n |
|
self.corpus_file = corpus_file |
|
|
|
self.ngrams = defaultdict(Counter) |
|
|
|
def _read_corpus(self) -> List[str]: |
|
"""Reads the corpus from the file. |
|
|
|
Returns: |
|
List of words in the corpus. |
|
""" |
|
with open(self.corpus_file, "r") as f: |
|
return f.read().split("\n") |
|
|
|
def create_ngrams(self): |
|
"""Creates n-grams from the corpus.""" |
|
corpus_words = self._read_corpus() |
|
self.total_words = len(corpus_words) |
|
|
|
self.ngrams[1] = Counter(corpus_words) |
|
for i in range(self.n - 1): |
|
self.ngrams[i + 2] = Counter(ngrams(corpus_words, i + 2)) |
|
|
|
def probability( |
|
self, ngram: str, higher_order_ngram: str = "", n: int = 1 |
|
) -> float: |
|
"""Computes maximum likelihood probability. |
|
|
|
Args: |
|
ngram: n-gram. |
|
higher_order_ngram: Higher order n-gram. Defaults to "". |
|
n: n-gram order. Defaults to 1. |
|
|
|
Returns: |
|
Maximum likelihood probability. |
|
""" |
|
if n == 1: |
|
return calc.log( |
|
(self.ngrams[1][ngram] + 1) |
|
/ (self.total_words + len(self.ngrams[1])) |
|
) |
|
|
|
assert n <= self.n, f"n must be less than or equal to {self.n}" |
|
return calc.log( |
|
(self.ngrams[n][higher_order_ngram] + 1) |
|
/ (self.ngrams[n - 1][ngram] + len(self.ngrams[1])) |
|
) |
|
|
|
def sentence_probability(self, sentence: str, n: int = 1) -> float: |
|
"""Computes cumulative n-gram ML probability of a sentence. |
|
|
|
Args: |
|
sentence: Sentence. |
|
n: n-gram order. Defaults to 1. |
|
|
|
Returns: |
|
Cumulative n-gram maximum likelihood probability. |
|
""" |
|
words = sentence.lower().split() |
|
cumulative_prob = 0 |
|
|
|
if n == 1: |
|
for word in words: |
|
cumulative_prob += self.probability(word) |
|
return cumulative_prob |
|
|
|
for i, word in enumerate(words): |
|
if i >= len(words) - n - 1: |
|
break |
|
cumulative_prob += self.probability( |
|
" ".join(words[i : i + n - 2]), |
|
" ".join(words[i : i + n - 1]), |
|
n, |
|
) |
|
|
|
return cumulative_prob |
|
|
|
@classmethod |
|
def load(cls, model_file: str) -> NGramMLE: |
|
"""Loads the model from a file. |
|
|
|
Args: |
|
model_file: File containing the model. |
|
|
|
Returns: |
|
Loaded model. |
|
""" |
|
return pickle.load(open(model_file, "rb")) |
|
|
|
def save(self, model_file: str) -> None: |
|
"""Saves the model to a file. |
|
|
|
Args: |
|
model_file: File to save the model. |
|
""" |
|
pickle.dump(self, open(model_file, "wb")) |
|
|
