0.3 prob to return aux task samples

Manu87DS · Apr 5, 2023 · 67f864a · 67f864a
1 parent 0de5b05
commit 67f864a
Showing 1 changed file with 9 additions and 4 deletions.
diff --git a/src/uie_dataset.py b/src/uie_dataset.py
@@ -27,6 +27,7 @@
 INSTRUCTION_STRATEGIES = ['single', 'multiple']
 ANSWER_PREFIX = "Answer:"
 SINGLE_QUOTES_SUBSTITUTE = "#$%#"
+AUX_PROB = 0.3
 
 
 def gen_cache_path(cache_dir, data_args):
@@ -327,7 +328,8 @@ def load_ES_dataset(self, dataset_path, labels_path, dataset_name, sampling_stra
                 "instruction": instruction
             }
 
-            yield example
+            if random.random() < AUX_PROB:
+                yield example
 
     def load_ET_dataset(self, dataset_path, labels_path, dataset_name, sampling_strategy, max_num_instances, subset):
         # ET = Entity Type
@@ -367,7 +369,8 @@ def load_ET_dataset(self, dataset_path, labels_path, dataset_name, sampling_stra
                 "instruction": instruction
             }
 
-            yield example
+            if random.random() < AUX_PROB:
+                yield example
 
     def load_EP_dataset(self, dataset_path, labels_path, dataset_name, sampling_strategy, max_num_instances, subset):
         # EP = Entity Pair
@@ -409,7 +412,8 @@ def load_EP_dataset(self, dataset_path, labels_path, dataset_name, sampling_stra
                 "instruction": instruction
             }
 
-            yield example
+            if random.random() < AUX_PROB:
+                yield example
 
     def load_EPR_dataset(self, dataset_path, labels_path, dataset_name, sampling_strategy, max_num_instances, subset):
         # EPR = Entity Pair Relationship
@@ -458,7 +462,8 @@ def load_EPR_dataset(self, dataset_path, labels_path, dataset_name, sampling_str
                 "instruction": instruction
             }
 
-            yield example
+            if random.random() < AUX_PROB:
+                yield example
 
     def load_RE_dataset(self, dataset_path, labels_path, dataset_name, sampling_strategy, max_num_instances, subset):
         instances, labels = self._load_dataset(dataset_path, labels_path)