Exemplo n.º 1
0
class TestTFIDF_InfoRetrieval(unittest.TestCase):
    def setUp(self):
        stopwords = "stop".split()
        keywords = "information agency retrieval".split()
        # documents = [
        #        ("Document 1", "information retrieval information retrieval"),
        #        ("Document 2", "retrieval retrieval retrieval retrieval"),
        #        ("Document 3", "agency information retrieval agency"),
        #        ("Document 4", "retrieval agency retrieval agency"),
        #    ]
        documents = Loader.load_documents("data/documents-lab1.txt")
        self.s = TFIDF(keywords, documents, Cleaner(stopwords))

    def test_keyword_setup(self):
        actual = self.s.keywords.items()
        expected = [("agenc", 0), ("inform", 1), ("retriev", 2)]
        self.assertEqual(actual, expected)

    def test_documents_setup(self):
        actual = self.s.document_vectors
        expected = {0: [0, 2, 2], 1: [0, 0, 4], 2: [2, 1, 1], 3: [2, 0, 2]}
        self.assertEqual(actual, expected)

    def test_tf(self):
        expected_results = [(0, [0, 1, 1]), (1, [0, 0, 1]), (2, [1, 0.5, 0.5]), (3, [1, 0, 1])]
        for index, expected_vector in expected_results:
            document = self.s.document_vectors[index]
            for word, i in self.s.keywords.items():
                actual = self.s.tf(document, word)
                expected = expected_vector[i]
                self.assertEqual(actual, expected)

    def test_idf(self):
        expected_results = [("inform", math.log(2, 10)), ("retriev", 0.0), ("agenc", math.log(2, 10))]
        for term, expected in expected_results:
            actual = self.s.idf(term)
            self.assertAlmostEqual(actual, expected, places=6)

    def test_tfidf(self):
        expected_results = [
            (0, [0, math.log(2, 10), 0]),
            (1, [0, 0, 0]),
            (2, [math.log(2, 10), 0.5 * math.log(2, 10), 0]),
            (3, [math.log(2, 10), 0, 0]),
        ]
        for index, expected_vector in expected_results:
            document = self.s.document_vectors[index]
            actual_vector = self.s.tfidf(document)
            for actual, expected in zip(actual_vector, expected_vector):
                self.assertAlmostEqual(actual, expected, places=6)

    def test_similarity(self):
        expected_results = [(0, 1), (1, 0), (2, math.sqrt(0.2)), (3, 0)]
        question_vector = self.s.phrase_to_vector("information retrieval")
        question_tfidfs = self.s.tfidf(question_vector)
        for index, expected in expected_results:
            actual = self.s.doc_question_similarity(index, question_tfidfs)
            self.assertEqual(actual, expected)

    def test_search(self):
        expected = [("Document 1", 1.0, 0), ("Document 3", math.sqrt(0.2), 2)]
        actual = self.s.search("information retrieval")
        self.assertEqual(actual, expected)
Exemplo n.º 2
0
class TestTFIDF(unittest.TestCase):
    def setUp(self):
        stopwords = "stop".split()
        keywords = "aaa bbb ccc ddd eee fff".split()
        documents = [
            ("document 1 ccc", "aaa aaa aaa ccc"),
            ("document 2 stop", "stop aaa bbb ccc"),
            ("document 3 stop", "aaa"),
            ("document 4 ddd", "aaa bbb ccc ddd eee"),
        ]
        self.s = TFIDF(keywords, documents, Cleaner(stopwords))

    def test_keyword_setup(self):
        actual = self.s.keywords.items()
        expected = [("aaa", 0), ("bbb", 1), ("ccc", 2), ("ddd", 3), ("eee", 4), ("fff", 5)]
        self.assertEqual(actual, expected)

    def test_documents_setup(self):
        actual = self.s.document_vectors
        expected = {0: [3, 0, 2, 0, 0, 0], 1: [1, 1, 1, 0, 0, 0], 2: [1, 0, 0, 0, 0, 0], 3: [1, 1, 1, 2, 1, 0]}
        self.assertEqual(actual, expected)

    def test_search_with_no_results(self):
        actual = self.s.search("fff")
        expected = []
        self.assertEqual(actual, expected)

    def test_search_with_only_popular_terms(self):
        actual = self.s.search("aaa")
        expected = []  # because idf=0
        self.assertEqual(actual, expected)

    def test_tf(self):
        document = self.s.document_vectors[0]
        actual = self.s.tf(document, "ccc")
        expected = 0.6666666666
        self.assertAlmostEqual(actual, expected)

        document = self.s.document_vectors[0]
        actual = self.s.tf(document, "aaa")
        expected = 1.0
        self.assertAlmostEqual(actual, expected)

        document = self.s.document_vectors[1]
        actual = self.s.tf(document, "aaa")
        expected = 1.0
        self.assertAlmostEqual(actual, expected)

        document = self.s.document_vectors[2]
        actual = self.s.tf(document, "aaa")
        expected = 1.0
        self.assertAlmostEqual(actual, expected)

        document = self.s.document_vectors[3]
        actual = self.s.tf(document, "aaa")
        expected = 0.5
        self.assertAlmostEqual(actual, expected)

    def test_idf(self):
        expected_results = [
            ("aaa", math.log(1.0, 10)),
            ("bbb", math.log(2.0, 10)),
            ("ccc", math.log(1.3333333333333, 10)),
            ("ddd", math.log(4.0, 10)),
            ("eee", math.log(4.0, 10)),
            ("fff", 0.0),
        ]

        for term, expected in expected_results:
            actual = self.s.idf(term)
            self.assertAlmostEqual(actual, expected)
Exemplo n.º 3
0
class TestTFIDF_flies(unittest.TestCase):
    def setUp(self):
        stopwords = "stop".split()
        keywords = "bee wasp fly fruit like".split()
        documents = [
            ("D1", "Time fly like an arrow but fruit fly like a banana."),
            ("D2", "It's strange that bees and wasps don't like each other."),
            ("D3", "The fly attendant sprayed the cabin with a strange fruit " "aerosol."),
            ("D4", "Try not to carry a light, as wasps and bees may fly " "toward it."),
            ("D5", "Fruit fly fly around in swarms. When fly they flap their " "wings 220 times a second."),
        ]
        self.s = TFIDF(keywords, documents, Cleaner(stopwords))

    def test_keyword_setup(self):
        actual = self.s.keywords.items()
        expected = [("bee", 0), ("fly", 1), ("fruit", 2), ("like", 3), ("wasp", 4)]
        self.assertEqual(actual, expected)

    def test_documents_setup(self):
        actual = self.s.document_vectors
        expected = {0: [0, 2, 1, 2, 0], 1: [1, 0, 0, 1, 1], 2: [0, 1, 1, 0, 0], 3: [1, 1, 0, 0, 1], 4: [0, 3, 1, 0, 0]}
        self.assertEqual(actual, expected)

    def test_tf(self):
        expected_results = [
            (0, [0, 1, 0.5, 1, 0]),
            (1, [1, 0, 0, 1, 1]),
            (2, [0, 1, 1, 0, 0]),
            (3, [1, 1, 0, 0, 1]),
            (4, [0, 1, 0.333333333333333333, 0, 0]),
        ]
        for index, expected_vector in expected_results:
            document = self.s.document_vectors[index]
            for word, i in self.s.keywords.items():
                actual = self.s.tf(document, word)
                expected = expected_vector[i]
                self.assertEqual(actual, expected)

    def test_idf(self):
        expected_results = [
            ("bee", 0.397940009),
            ("fly", 0.096910013),
            ("fruit", 0.22184875),
            ("like", 0.397940009),
            ("wasp", 0.397940009),
        ]
        for term, expected in expected_results:
            actual = self.s.idf(term)
            self.assertAlmostEqual(actual, expected, places=6)

    def test_tfidf(self):
        expected_results = [
            (0, [0, 0.096910013, 0.110924375, 0.397940009, 0]),
            (1, [0.397940009, 0, 0, 0.397940009, 0.397940009]),
            (2, [0, 0.096910013, 0.22184875, 0, 0]),
            (3, [0.397940009, 0.096910013, 0, 0, 0.397940009]),
            (4, [0, 0.096910013, 0.073949583, 0, 0]),
        ]
        for title, expected_vector in expected_results:
            document = self.s.document_vectors[title]
            actual_vector = self.s.tfidf(document)
            for actual, expected in zip(actual_vector, expected_vector):
                self.assertAlmostEqual(actual, expected, places=6)