1  from nltk.util import ngrams
   2  
   3  
   4  class Corpus:
   5      def __init__(self, corpus, alpha=1, start_symbol='<s>', end_symbol='</s>'):
   6          self.n = 0
   7          self.vocabulary = dict()
   8          self.unigrams = dict()
   9          self.bigrams = dict()
  10  
  11          for sentence in corpus:
  12              for word in sentence:
  13                  self.vocabulary[word] = self.vocabulary.get(word, 0) + 1
  14                  self.unigrams[word] = self.unigrams.get(word, 0) + 1
  15  
  16              self.unigrams[start_symbol] = self.unigrams.get(start_symbol, 0) + 1
  17              self.unigrams[end_symbol] = self.unigrams.get(end_symbol, 0) + 1
  18  
  19              grams = ngrams([start_symbol] + sentence + [end_symbol], 2)
  20              for gram in grams:
  21                  self.n += 1
  22                  if gram in self.bigrams:
  23                      self.bigrams[gram] += 1
  24                  else:
  25                      self.bigrams[gram] = 1
  26  
  27          # Laplace smoothing
  28          _vocabulary = list(self.vocabulary.keys()) + [start_symbol, end_symbol]
  29          for w in _vocabulary:
  30              for w2 in _vocabulary:
  31                  self.n += alpha
  32                  self.bigrams[w, w2] = self.bigrams.get((w, w2), 0) + alpha
  33                  self.unigrams[w2] += alpha
  34