minor fixes

Author: vmurahari3

models/multiplexing.py

@@ -235,8 +235,6 @@ def forward(
                 .expand(modified_batch_size, modified_seq_length),
                 instance_labels,
             ]
-            retrieval_labels = torch.div(retrieval_labels, self.config.retrieval_loss_vocab_scale, rounding_mode='trunc')
-            retrieval_labels = retrieval_labels.long()
             retrieval_labels[:, :special_tokens_end_position] = -100
 
             pad_mask = retrieval_labels == 1
@@ -258,7 +256,7 @@ def forward(
             loss_fct = CrossEntropyLoss()
             task_loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
             retrieval_loss = loss_fct(
-                retrieval_predictions.view(-1, self.config.vocab_size)
+                retrieval_predictions.view(-1, self.config.vocab_size),
                 retrieval_labels.view(-1),
             )
             loss = (self.task_loss_coeff * task_loss) + (
@@ -582,8 +580,8 @@ def __init__(self, config):
         self.dense = nn.Linear(2 * config.hidden_size, config.hidden_size)
         self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
 
-        self.decoder = nn.Linear(config.hidden_size, math.ceil(config.vocab_size / config.retrieval_loss_vocab_scale))
-        self.bias = nn.Parameter(torch.zeros(math.ceil(config.vocab_size / config.retrieval_loss_vocab_scale)))
+        self.decoder = nn.Linear(config.hidden_size, config.vocab_size)
+        self.bias = nn.Parameter(torch.zeros(config.vocab_size))
         self.decoder.bias = self.bias
 
     def forward(self, features, instance_labels, **kwargs):
@@ -727,8 +725,8 @@ def __init__(self, config):
         self.layer_norm_pre_vocab = nn.LayerNorm(
             config.hidden_size, eps=config.layer_norm_eps
         )
-        self.decoder = nn.Linear(config.hidden_size, math.ceil(config.vocab_size / config.retrieval_loss_vocab_scale))
-        self.bias = nn.Parameter(torch.zeros(math.ceil(config.vocab_size / config.retrieval_loss_vocab_scale)))
+        self.decoder = nn.Linear(config.hidden_size, config.vocab_size)
+        self.bias = nn.Parameter(torch.zeros(config.vocab_size))
         self.decoder.bias = self.bias
 
     def forward(self, features, instance_labels, **kwargs):

models/trainer.py

@@ -1754,12 +1754,13 @@ def prediction_step(
                 ) = self.compute_loss(model, inputs, return_outputs=True)
                 loss = loss.mean().detach()
                 if isinstance(outputs, dict):
-                    logits = tuple(
-                        v
-                        for k, v in outputs.items()
-                        if k
-                        not in ignore_keys + ["loss", "task_loss", "retrieval_loss"]
-                    )
+                    # logits = tuple(
+                    #     v
+                    #     for k, v in outputs.items()
+                    #     if k
+                    #     not in ignore_keys + ["loss", "task_loss", "retrieval_loss"]
+                    # )
+                    logits = outputs["logits"] if "logits" in outputs else None
                 else:
                     logits = outputs[1:]
                 if "retrieval_loss" in outputs:
@@ -1778,9 +1779,10 @@ def prediction_step(
                 else:
                     outputs = model(**inputs)
                 if isinstance(outputs, dict):
-                    logits = tuple(
-                        v for k, v in outputs.items() if k not in ignore_keys
-                    )
+                    # logits = tuple(
+                    #     v for k, v in outputs.items() if k not in ignore_keys
+                    # )
+                    logits = outputs["logits"] if "logits" in outputs else None
                 else:
                     logits = outputs
                 if self.args.past_index >= 0:

run_glue.py

@@ -696,7 +696,7 @@ def compute_metrics(p: EvalPrediction):
         #     eval_datasets.append(datasets["validation_mismatched"])
 
         for eval_dataset, task in zip(eval_datasets, tasks):
-            metrics = trainer.evaluate(eval_dataset=eval_dataset) 
+            metrics = trainer.evaluate(eval_dataset=eval_dataset)
 
             max_eval_samples = (
                 data_args.max_eval_samples

run_ner.py

@@ -291,11 +291,6 @@ def main():
     last_checkpoint = None
     if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
         last_checkpoint = get_last_checkpoint(training_args.output_dir)
-        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
-            raise ValueError(
-                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
-                "Use --overwrite_output_dir to overcome."
-            )
 
     # Set seed before initializing model.
     set_seed(training_args.seed)
@@ -643,10 +638,6 @@ def compute_metrics(p):
         else:
             kwargs["dataset"] = data_args.dataset_name
 
-    if training_args.push_to_hub:
-        trainer.push_to_hub(**kwargs)
-    else:
-        trainer.create_model_card(**kwargs)
 
 
 def _mp_fn(index):

run_ner.sh

@@ -277,8 +277,10 @@ CMD="python run_ner.py \
 --demuxing_variant ${DEMUXING} \
 --should_mux ${SHOULD_MUX} \
 --gaussian_hadamard_norm ${RANDOM_ENCODING_NORM} \
---learn_muxing ${LEARN_MUXING}"
-
+--learn_muxing ${LEARN_MUXING} \
+--load_best_model_at_end 1 \
+--metric_for_best_model eval_f1 \
+--save_total_limit 1"
 if [ "$DO_TRAIN" -eq 1 ]; then
         CMD="${CMD} --do_train"
 fi