added lda and etm tests
@@ -9,8 +9,7 @@ from click.testing import CliRunner
from octis.evaluation_metrics.coherence_metrics import *
from octis.dataset.dataset import Dataset
from octis.models.LDA import LDA
-
-from octis import cli
+from octis.models.ETM import ETM
import os
@@ -23,6 +22,7 @@ def root_dir():
def data_dir(root_dir):
return root_dir + "/../octis/preprocessed_datasets/"
+
def test_coherence_measures(data_dir):
dataset = Dataset()
dataset.load(data_dir + '/M10')
@@ -34,3 +34,57 @@ def test_coherence_measures(data_dir):
score = metric.score(output)
assert type(score) == np.float64
+
+def test_model_output_lda(data_dir):
+ dataset = Dataset()
+ dataset.load(data_dir + '/M10')
+ num_topics = 3
+ model = LDA(num_topics=num_topics, iterations=5)
+ output = model.train_model(dataset)
+ assert 'topics' in output.keys()
+ assert 'topic-word-matrix' in output.keys()
+ assert 'test-topic-document-matrix' in output.keys()
+
+ # check topics format
+ assert type(output['topics']) == list
+ assert len(output['topics']) == num_topics
+
+ # check topic-word-matrix format
+ assert type(output['topic-word-matrix']) == np.ndarray
+ assert output['topic-word-matrix'].shape == (num_topics, len(dataset.get_vocabulary()))
+
+ # check topic-document-matrix format
+ assert type(output['topic-document-matrix']) == np.ndarray
+ assert output['topic-document-matrix'].shape == (num_topics, len(dataset.get_partitioned_corpus()[0]))
+
+ # check test-topic-document-matrix format
+ assert type(output['test-topic-document-matrix']) == np.ndarray
+ assert output['test-topic-document-matrix'].shape == (num_topics, len(dataset.get_partitioned_corpus()[1]))
+
+
+def test_model_output_etm(data_dir):
+ dataset = Dataset()
+ dataset.load(data_dir + '/M10')
+ num_topics = 3
+ model = ETM(num_topics=num_topics, num_epochs=5)
+ output = model.train_model(dataset)
+ assert 'topics' in output.keys()
+ assert 'topic-word-matrix' in output.keys()
+ assert 'test-topic-document-matrix' in output.keys()
+
+ # check topics format
+ assert type(output['topics']) == list
+ assert len(output['topics']) == num_topics
+
+ # check topic-word-matrix format
+ assert type(output['topic-word-matrix']) == np.ndarray
+ assert output['topic-word-matrix'].shape == (num_topics, len(dataset.get_vocabulary()))
+
+ # check topic-document-matrix format
+ assert type(output['topic-document-matrix']) == np.ndarray
+ assert output['topic-document-matrix'].shape == (num_topics, len(dataset.get_partitioned_corpus()[0]))
+
+ # check test-topic-document-matrix format
+ assert type(output['test-topic-document-matrix']) == np.ndarray
+ assert output['test-topic-document-matrix'].shape == (num_topics, len(dataset.get_partitioned_corpus()[1]))
+
