added new ctm and nfm tests
@@ -11,6 +11,7 @@ from octis.dataset.dataset import Dataset
from octis.models.LDA import LDA
from octis.models.ETM import ETM
from octis.models.CTM import CTM
+from octis.models.NMF import NMF
import os
@@ -91,6 +92,34 @@ def test_model_output_etm(data_dir):
assert output['test-topic-document-matrix'].shape == (num_topics, len(dataset.get_partitioned_corpus()[1]))
+def test_model_output_nmf(data_dir):
+ dataset = Dataset()
+ dataset.load(data_dir + '/M10')
+ num_topics = 3
+ model = NMF(num_topics=num_topics, w_max_iter=10, h_max_iter=10, use_partitions=True)
+ 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_ctm_zeroshot(data_dir):
dataset = Dataset()
dataset.load(data_dir + '/M10')
@@ -116,3 +145,30 @@ def test_model_output_ctm_zeroshot(data_dir):
# 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_ctm_combined(data_dir):
+ dataset = Dataset()
+ dataset.load(data_dir + '/M10')
+ num_topics = 3
+ model = CTM(num_topics=num_topics, num_epochs=5, inference_type='combined')
+ 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]))
