forked from piskvorky/gensim
-
Notifications
You must be signed in to change notification settings - Fork 6
/
test_dtm.py
49 lines (37 loc) · 1.65 KB
/
test_dtm.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
#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Automated tests for DTM/DIM model
"""
import logging
import gensim
import os
import sys
import unittest
from gensim import corpora
# needed because sample data files are located in the same folder
module_path = os.path.dirname(__file__)
datapath = lambda fname: os.path.join(module_path, 'test_data', fname)
class TestDtmModel(unittest.TestCase):
def setUp(self):
self.time_slices = [3, 7]
self.corpus = corpora.mmcorpus.MmCorpus(datapath('dtm_test.mm'))
self.id2word = corpora.Dictionary.load(datapath('dtm_test.dict'))
# first you need to setup the environment variable $DTM_PATH for the dtm executable file
self.dtm_path = os.environ.get('DTM_PATH', None)
if self.dtm_path is None:
if sys.version_info >= (2, 7, 0):
self.skipTest("$DTM_PATH is not properly set up.")
else:
logging.warning("$DTM_PATH is not properly set up.")
def testDtm(self):
if self.dtm_path is not None:
model = gensim.models.wrappers.DtmModel(self.dtm_path, self.corpus, self.time_slices, num_topics=2, id2word=self.id2word, model='dtm', initialize_lda=True)
topics = model.show_topics(topics=2, times=2, topn=10)
def testDim(self):
if self.dtm_path is not None:
model = gensim.models.wrappers.DtmModel(self.dtm_path, self.corpus, self.time_slices, num_topics=2, id2word=self.id2word, model='fixed', initialize_lda=True)
topics = model.show_topics(topics=2, times=2, topn=10)
if __name__ == '__main__':
logging.basicConfig(level=logging.DEBUG)
unittest.main()