Code for main test

arguments.py

@@ -85,7 +85,7 @@ def add_training_args(parser):
                        'with larger models and sequences')
     group.add_argument('--clip-grad', type=float, default=1.0,
                        help='gradient clipping')
-    group.add_argument('--epochs', type=int, default=10,
+    group.add_argument('--epochs', type=int, default=32,
                        help='upper epoch limit')
     group.add_argument('--log-interval', type=int, default=1000000,
                        help='report interval')

bert_config.json

@@ -16,5 +16,5 @@
   "fp32_embedding": true,
   "fp32_tokentypes": false,
   "fp16": false,
-  "num_tasks": 15
+  "num_tasks": 16
 }

data_utils/datasets.py

@@ -773,7 +773,7 @@ def scale(self, vector):
         if len(vector) == 0:
             return vector
         scale = 10
-        scaling = lambda x: ((scale - 1) * (x - _min) / (_max - _min + 1e-8)) + 1
+        scaling = lambda x: ((scale) * (x - _min) / (_max - _min + 1e-8))
         _min = min(vector)
         _max = max(vector)
         return [scaling(w) for w in vector]
@@ -794,7 +794,7 @@ def get_word_labels(self, sent, doc):
             if "tf" in self.modes:
                 labels["tf"] = self.scale(tf)
             elif "tf_idf" in self.modes:
-                labels["tf_idf"] = tf_idf #self.scale(tf_idf)
+                labels["tf_idf"] = self.scale(tf_idf)
 
         return labels
 

pretrain_bert.py

@@ -97,7 +97,7 @@ def setup_model_and_optimizer(args, tokenizer):
     optimizer = get_optimizer(model, args)
     lr_scheduler = get_learning_rate_scheduler(optimizer, args)
     criterion_cls = torch.nn.CrossEntropyLoss(reduce=False, ignore_index=-1)
-    criterion_reg = torch.nn.L1Loss(reduce=False)
+    criterion_reg = torch.nn.MSELoss(reduce=False)
 
     criterion = (criterion_cls, criterion_reg)