Fix little bug in SS3.fit

When one of the documents was empty, the training failed with an "IndexError: string index out of range" exception.
sergioburdisso · Feb 19, 2020 · 31eccbc · 31eccbc
1 parent 2a6af26
commit 31eccbc
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/pyss3/__init__.py b/pyss3/__init__.py
@@ -2032,7 +2032,10 @@ def fit(self, x_train, y_train, n_grams=1, prep=True, leave_pbar=True):
 
         x_train = [
             "".join([
-                x_train[i] if x_train[i][-1] == '\n' else x_train[i] + '\n'
+                x_train[i]
+                if x_train[i] and x_train[i][-1] == '\n'
+                else
+                x_train[i] + '\n'
                 for i in xrange(len(x_train))
                 if y_train[i] == cat
             ])