test: adds test case that causes etm to throw an error because a document contains a single word
@@ -14,6 +14,7 @@ from octis.models.CTM import CTM
from octis.models.NMF import NMF
from octis.models.NMF_scikit import NMF_scikit
from octis.models.ProdLDA import ProdLDA
+from octis.preprocessing.preprocessing import Preprocessing
import numpy as np
import os
@@ -82,6 +83,36 @@ def test_model_output_etm(data_dir):
assert type(output['test-topic-document-matrix']) == np.ndarray
assert output['test-topic-document-matrix'].shape == (num_topics, len(dataset.get_partitioned_corpus()[2]))
+def test_model_output_etm_with_corpus_containing_single_word_document(data_dir):
+ texts_path = data_dir+"/sample_texts/unprepr_docs.txt"
+ p = Preprocessing(vocabulary=None, max_features=None, remove_punctuation=True,
+ lemmatize=False, stopword_list='english')
+ dataset = p.preprocess_dataset(
+ documents_path=texts_path,
+ )
+ 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()[2]))
+
def test_model_output_etm_not_partitioned(data_dir):
dataset = Dataset()
