fix(etm): fixes an error generated by feeding the model with a single…

…-word document Fixes #37
MIND-Lab · Oct 31, 2021 · e5c4446 · e5c4446
1 parent 47c342a
commit e5c4446
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 15 deletions.
diff --git a/octis/models/ETM.py b/octis/models/ETM.py
@@ -123,7 +123,7 @@ def _train_epoch(self, epoch):
             self.optimizer.zero_grad()
             self.model.zero_grad()
             data_batch = data.get_batch(self.train_tokens, self.train_counts, ind, len(self.vocab.keys()),
-                                        self.hyperparameters['embedding_size'], self.device)
+                                        self.device)
             sums = data_batch.sum(1).unsqueeze(1)
             if self.hyperparameters['bow_norm']:
                 normalized_data_batch = data_batch / sums
@@ -179,7 +179,7 @@ def _train_epoch(self, epoch):
                     self.model.zero_grad()
                     val_data_batch = data.get_batch(self.valid_tokens, self.valid_counts,
                                                     ind, len(self.vocab.keys()),
-                                                    self.hyperparameters['embedding_size'], self.device)
+                                                    self.device)
                     sums = val_data_batch.sum(1).unsqueeze(1)
                     if self.hyperparameters['bow_norm']:
                         val_normalized_data_batch = val_data_batch / sums
@@ -245,7 +245,7 @@ def inference(self):
         for idx, ind in enumerate(indices):
             data_batch = data.get_batch(self.test_tokens, self.test_counts,
                                         ind, len(self.vocab.keys()),
-                                        self.hyperparameters['embedding_size'], self.device)
+                                        self.device)
             sums = data_batch.sum(1).unsqueeze(1)
             if self.hyperparameters['bow_norm']:
                 normalized_data_batch = data_batch / sums

diff --git a/octis/models/ETM_model/data.py b/octis/models/ETM_model/data.py
@@ -1,24 +1,13 @@
-import os
-import random
-import pickle
 import numpy as np
 import torch
-import scipy.io
 
-def get_batch(tokens, counts, ind, vocab_size, emsize, device):
+def get_batch(tokens, counts, ind, vocab_size, device):
     """fetch input data by batch."""
     batch_size = len(ind)
     data_batch = np.zeros((batch_size, vocab_size))
     for i, doc_id in enumerate(ind):
         doc = tokens[doc_id]
         count = counts[doc_id]
-        #L = count.shape[1]
-        if len(doc) == 1:
-            doc = [doc.squeeze()]
-            count = [count.squeeze()]
-        else:
-            doc = doc#.squeeze()
-            count = count#.squeeze()
         if doc_id != -1:
             for j, word in enumerate(doc):
                 data_batch[i, word] = count[j]