Source code for nltk.test.unit.lm.test_vocabulary

# Natural Language Toolkit: Language Model Unit Tests
# Copyright (C) 2001-2022 NLTK Project
# Author: Ilia Kurenkov <>
# URL: <>
# For license information, see LICENSE.TXT

import unittest
from collections import Counter
from timeit import timeit

from nltk.lm import Vocabulary

[docs]class NgramModelVocabularyTests(unittest.TestCase): """tests Vocabulary Class"""
[docs] @classmethod def setUpClass(cls): cls.vocab = Vocabulary( ["z", "a", "b", "c", "f", "d", "e", "g", "a", "d", "b", "e", "w"], unk_cutoff=2, )
[docs] def test_truthiness(self): self.assertTrue(self.vocab)
[docs] def test_cutoff_value_set_correctly(self): self.assertEqual(self.vocab.cutoff, 2)
[docs] def test_unable_to_change_cutoff(self): with self.assertRaises(AttributeError): self.vocab.cutoff = 3
[docs] def test_cutoff_setter_checks_value(self): with self.assertRaises(ValueError) as exc_info: Vocabulary("abc", unk_cutoff=0) expected_error_msg = "Cutoff value cannot be less than 1. Got: 0" self.assertEqual(expected_error_msg, str(exc_info.exception))
[docs] def test_counts_set_correctly(self): self.assertEqual(self.vocab.counts["a"], 2) self.assertEqual(self.vocab.counts["b"], 2) self.assertEqual(self.vocab.counts["c"], 1)
[docs] def test_membership_check_respects_cutoff(self): # a was seen 2 times, so it should be considered part of the vocabulary self.assertTrue("a" in self.vocab) # "c" was seen once, it shouldn't be considered part of the vocab self.assertFalse("c" in self.vocab) # "z" was never seen at all, also shouldn't be considered in the vocab self.assertFalse("z" in self.vocab)
[docs] def test_vocab_len_respects_cutoff(self): # Vocab size is the number of unique tokens that occur at least as often # as the cutoff value, plus 1 to account for unknown words. self.assertEqual(5, len(self.vocab))
[docs] def test_vocab_iter_respects_cutoff(self): vocab_counts = ["a", "b", "c", "d", "e", "f", "g", "w", "z"] vocab_items = ["a", "b", "d", "e", "<UNK>"] self.assertCountEqual(vocab_counts, list(self.vocab.counts.keys())) self.assertCountEqual(vocab_items, list(self.vocab))
[docs] def test_update_empty_vocab(self): empty = Vocabulary(unk_cutoff=2) self.assertEqual(len(empty), 0) self.assertFalse(empty) self.assertIn(empty.unk_label, empty) empty.update(list("abcde")) self.assertIn(empty.unk_label, empty)
[docs] def test_lookup(self): self.assertEqual(self.vocab.lookup("a"), "a") self.assertEqual(self.vocab.lookup("c"), "<UNK>")
[docs] def test_lookup_iterables(self): self.assertEqual(self.vocab.lookup(["a", "b"]), ("a", "b")) self.assertEqual(self.vocab.lookup(("a", "b")), ("a", "b")) self.assertEqual(self.vocab.lookup(("a", "c")), ("a", "<UNK>")) self.assertEqual( self.vocab.lookup(map(str, range(3))), ("<UNK>", "<UNK>", "<UNK>") )
[docs] def test_lookup_empty_iterables(self): self.assertEqual(self.vocab.lookup(()), ()) self.assertEqual(self.vocab.lookup([]), ()) self.assertEqual(self.vocab.lookup(iter([])), ()) self.assertEqual(self.vocab.lookup(n for n in range(0, 0)), ())
[docs] def test_lookup_recursive(self): self.assertEqual( self.vocab.lookup([["a", "b"], ["a", "c"]]), (("a", "b"), ("a", "<UNK>")) ) self.assertEqual(self.vocab.lookup([["a", "b"], "c"]), (("a", "b"), "<UNK>")) self.assertEqual(self.vocab.lookup([[[[["a", "b"]]]]]), ((((("a", "b"),),),),))
[docs] def test_lookup_None(self): with self.assertRaises(TypeError): self.vocab.lookup(None) with self.assertRaises(TypeError): list(self.vocab.lookup([None, None]))
[docs] def test_lookup_int(self): with self.assertRaises(TypeError): self.vocab.lookup(1) with self.assertRaises(TypeError): list(self.vocab.lookup([1, 2]))
[docs] def test_lookup_empty_str(self): self.assertEqual(self.vocab.lookup(""), "<UNK>")
[docs] def test_eqality(self): v1 = Vocabulary(["a", "b", "c"], unk_cutoff=1) v2 = Vocabulary(["a", "b", "c"], unk_cutoff=1) v3 = Vocabulary(["a", "b", "c"], unk_cutoff=1, unk_label="blah") v4 = Vocabulary(["a", "b"], unk_cutoff=1) self.assertEqual(v1, v2) self.assertNotEqual(v1, v3) self.assertNotEqual(v1, v4)
[docs] def test_str(self): self.assertEqual( str(self.vocab), "<Vocabulary with cutoff=2 unk_label='<UNK>' and 5 items>" )
[docs] def test_creation_with_counter(self): self.assertEqual( self.vocab, Vocabulary( Counter( ["z", "a", "b", "c", "f", "d", "e", "g", "a", "d", "b", "e", "w"] ), unk_cutoff=2, ), )
[docs] @unittest.skip( reason="Test is known to be flaky as it compares (runtime) performance." ) def test_len_is_constant(self): # Given an obviously small and an obviously large vocabulary. small_vocab = Vocabulary("abcde") from nltk.corpus.europarl_raw import english large_vocab = Vocabulary(english.words()) # If we time calling `len` on them. small_vocab_len_time = timeit("len(small_vocab)", globals=locals()) large_vocab_len_time = timeit("len(large_vocab)", globals=locals()) # The timing should be the same order of magnitude. self.assertAlmostEqual(small_vocab_len_time, large_vocab_len_time, places=1)