GH-798: skip sentences without labels in CSV reader

flairNLP · Jun 24, 2019 · 1c5fccf · 1c5fccf
1 parent 6e96d78
commit 1c5fccf
Showing 1 changed file with 16 additions and 0 deletions.
diff --git a/flair/datasets.py b/flair/datasets.py
@@ -681,6 +681,22 @@ def __init__(
 
             for row in csv_reader:
 
+                # test if format is OK
+                wrong_format = False
+                for text_column in self.text_columns:
+                    if text_column >= len(row):
+                        wrong_format = True
+
+                # test if at least one label given
+                has_label = False
+                for column in self.column_name_map:
+                    if self.column_name_map[column].startswith("label") and row[column]:
+                        has_label = True
+                        break
+
+                if wrong_format or not has_label:
+                    continue
+
                 if self.in_memory:
 
                     text = " || ".join(