ETM - Single-word document causes ETM training to fail

[lfmatosm]

• OCTIS version: 1.9.0
• Python version: 3.7.6
• Operating System: Ubuntu 20.04 LTS

Description

Probably related to #20 and #35.
If your preprocessed corpus contains any single-word document, ETM training fails. This should not happen, as the Preprocessing class has 0 as the default value for parameters min_words_docs and min_df, which define respectivelly the minimum number of words a document must have to be keep and the minimum document-frequency for words on the corpus.

What I Did

I've implemented a test case illustrating the scenario. The test fails. The test code can be found here, and the error stacktrace as shown on Github actions can be seen here.

Below, the aforementioned stacktrace (on my local machine):

Current call:  0
model: ETM(
  (t_drop): Dropout(p=0.5, inplace=False)
  (theta_act): ReLU()
  (alphas): Linear(in_features=300, out_features=16, bias=False)
  (q_theta): Sequential(
    (0): Linear(in_features=872, out_features=800, bias=True)
    (1): ReLU()
    (2): Linear(in_features=800, out_features=800, bias=True)
    (3): ReLU()
  )
  (mu_q_theta): Linear(in_features=800, out_features=16, bias=True)
  (logsigma_q_theta): Linear(in_features=800, out_features=16, bias=True)
)
Traceback (most recent call last):
  File "octis_test/unified_training.py", line 55, in <module>
    model_runs=5, plot_best_seen=True) # number of runs of the topic model
  File "/home/luizmatos/anaconda3/lib/python3.7/site-packages/octis/optimization/optimizer.py", line 160, in optimize
    results = self._optimization_loop(opt)
  File "/home/luizmatos/anaconda3/lib/python3.7/site-packages/octis/optimization/optimizer.py", line 285, in _optimization_loop
    f_val = self._objective_function(next_x)
  File "/home/luizmatos/anaconda3/lib/python3.7/site-packages/octis/optimization/optimizer.py", line 217, in _objective_function
    self.topk)
  File "/home/luizmatos/anaconda3/lib/python3.7/site-packages/octis/models/ETM.py", line 60, in train_model
    continue_training = self._train_epoch(epoch)
  File "/home/luizmatos/anaconda3/lib/python3.7/site-packages/octis/models/ETM.py", line 126, in _train_epoch
    self.hyperparameters['embedding_size'], self.device)
  File "/home/luizmatos/anaconda3/lib/python3.7/site-packages/octis/models/ETM_model/data.py", line 17, in get_batch
    doc = [doc.squeeze()]
AttributeError: 'list' object has no attribute 'squeeze'