Update modeling_hf_nomic_bert.py

modeling_hf_nomic_bert.py

@@ -1694,7 +1694,6 @@ class NomicBertModel(NomicBertPreTrainedModel):
         return_dict=None,
         matryoshka_dim=None,
         inputs_embeds=None,
-        head_mask=None,
     ):
         if input_ids is not None and inputs_embeds is not None:
             raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
@@ -1868,7 +1867,7 @@ class NomicBertForMultipleChoice(NomicBertPreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
 
-        self.bert = NomicBertModel(config, add_pooling_layer=True)
+        self.new = NomicBertModel(config, add_pooling_layer=True)
         classifier_dropout = (
             config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
         )
@@ -1911,17 +1910,13 @@ class NomicBertForMultipleChoice(NomicBertPreTrainedModel):
             else None
         )
 
-        outputs = self.bert(
+        outputs = self.new(
             input_ids,
             attention_mask=attention_mask,
             token_type_ids=token_type_ids,
             position_ids=position_ids,
             head_mask=head_mask,
             inputs_embeds=inputs_embeds,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-            unpad_inputs=unpad_inputs,
         )
 
         pooled_output = outputs[1]
@@ -1987,9 +1982,6 @@ class NomicBertForTokenClassification(NomicBertPreTrainedModel):
             position_ids=position_ids,
             head_mask=head_mask,
             inputs_embeds=inputs_embeds,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
         )
 
         sequence_output = outputs[0]
@@ -1999,7 +1991,7 @@ class NomicBertForTokenClassification(NomicBertPreTrainedModel):
 
         loss = None
         if labels is not None:
-            loss_fct = CrossEntropyLoss()
+            loss_fct = nn.CrossEntropyLoss()
             loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
 
         if not return_dict:
@@ -2057,9 +2049,6 @@ class NomicBertForQuestionAnswering(NomicBertPreTrainedModel):
             position_ids=position_ids,
             head_mask=head_mask,
             inputs_embeds=inputs_embeds,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
         )
 
         sequence_output = outputs[0]
@@ -2081,7 +2070,7 @@ class NomicBertForQuestionAnswering(NomicBertPreTrainedModel):
             start_positions = start_positions.clamp(0, ignored_index)
             end_positions = end_positions.clamp(0, ignored_index)
 
-            loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
+            loss_fct = nn.CrossEntropyLoss(ignore_index=ignored_index)
             start_loss = loss_fct(start_logits, start_positions)
             end_loss = loss_fct(end_logits, end_positions)
             total_loss = (start_loss + end_loss) / 2