def test_ngram_text_generator_generate_next_word(self):
"""
Checks that next word generates properly
"""
corpus = ('i', 'have', 'a', 'cat', '<END>', 'i', 'have', 'a', 'bruno',
'cat', '<END>')
word_storage = WordStorage()
word_storage.update(corpus)
encoded = encode_text(word_storage, corpus)
ngram = NGramTrie(3, encoded)
generator = NGramTextGenerator(word_storage, ngram)
context = (word_storage.get_id('i'), word_storage.get_id('have'))
expected = word_storage.get_id('a')
actual = generator._generate_next_word(context)
self.assertEqual(expected, actual)
def test_ngram_text_generator_generate_next_word_no_such_context(self):
"""
Checks that next word generates properly if no context found
"""
corpus = ('i', 'have', 'a', 'cat', '<END>', 'his', 'name', 'is',
'bruno', '<END>')
word_storage = WordStorage()
word_storage.update(corpus)
encoded = encode_text(word_storage, corpus)
ngram = NGramTrie(3, encoded)
generator = NGramTextGenerator(word_storage, ngram)
context = (
word_storage.get_id('i'),
word_storage.get_id('name'),
) # there is no such context in ngrams, so return most frequent option
expected_top_freq = word_storage.get_id('<END>') # as it appears twice
actual = generator._generate_next_word(context)
self.assertEqual(expected_top_freq, actual)
