-
-
Notifications
You must be signed in to change notification settings - Fork 4.4k
/
test_segmentation.py
65 lines (53 loc) · 2 KB
/
test_segmentation.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Copyright (C) 2011 Radim Rehurek <radimrehurek@seznam.cz>
# Licensed under the GNU LGPL v2.1 - http://www.gnu.org/licenses/lgpl.html
"""
Automated tests for segmentation algorithms in the segmentation module.
"""
import logging
import unittest
import numpy as np
from gensim.topic_coherence import segmentation
from numpy import array
class TestSegmentation(unittest.TestCase):
def setUp(self):
self.topics = [
array([9, 4, 6]),
array([9, 10, 7]),
array([5, 2, 7])
]
def testSOnePre(self):
"""Test s_one_pre segmentation."""
actual = segmentation.s_one_pre(self.topics)
expected = [
[(4, 9), (6, 9), (6, 4)],
[(10, 9), (7, 9), (7, 10)],
[(2, 5), (7, 5), (7, 2)]
]
self.assertTrue(np.allclose(actual, expected))
def testSOneOne(self):
"""Test s_one_one segmentation."""
actual = segmentation.s_one_one(self.topics)
expected = [
[(9, 4), (9, 6), (4, 9), (4, 6), (6, 9), (6, 4)],
[(9, 10), (9, 7), (10, 9), (10, 7), (7, 9), (7, 10)],
[(5, 2), (5, 7), (2, 5), (2, 7), (7, 5), (7, 2)]
]
self.assertTrue(np.allclose(actual, expected))
def testSOneSet(self):
"""Test s_one_set segmentation."""
actual = segmentation.s_one_set(self.topics)
expected = [
[(9, array([9, 4, 6])), (4, array([9, 4, 6])), (6, array([9, 4, 6]))],
[(9, array([9, 10, 7])), (10, array([9, 10, 7])), (7, array([9, 10, 7]))],
[(5, array([5, 2, 7])), (2, array([5, 2, 7])), (7, array([5, 2, 7]))]
]
for s_i in range(len(actual)):
for j in range(len(actual[s_i])):
self.assertEqual(actual[s_i][j][0], expected[s_i][j][0])
self.assertTrue(np.allclose(actual[s_i][j][1], expected[s_i][j][1]))
if __name__ == '__main__':
logging.root.setLevel(logging.WARNING)
unittest.main()