fix binary labels

ahmedbesbes · Nov 1, 2019 · b8080f6 · b8080f6
1 parent 00c488d
commit b8080f6
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 3 deletions.
diff --git a/src/data_loader.py b/src/data_loader.py
@@ -52,7 +52,7 @@ def load_data(args):
             labels = [label for (text, label) in clean_data]
 
             labels = list(
-                map(lambda l: {1: 0, 2: 0, 4: 2, 5: 2}[l], labels))
+                map(lambda l: {1: 0, 2: 0, 4: 1, 5: 1}[l], labels))
 
         else:
             labels = list(
@@ -80,7 +80,7 @@ def load_data(args):
     print(
         f'data loaded successfully with {len(texts)} rows and {number_of_classes} labels')
     print('Distribution of the classes', Counter(labels))
-    
+
     sample_weights = get_sample_weights(labels)
 
     return texts, labels, number_of_classes, sample_weights

diff --git a/train.py b/train.py
@@ -368,7 +368,6 @@ def run(args, both_cases=False):
     parser.add_argument('--group_labels', type=int, default=1, choices=[0, 1])
     parser.add_argument('--ignore_center', type=int, default=1, choices=[0, 1])
     parser.add_argument('--label_ignored', type=int, default=None)
-
     parser.add_argument('--ratio', type=float, default=1)
     parser.add_argument('--balance', type=int, default=0, choices=[0, 1])
     parser.add_argument('--use_sampler', type=int,