Patch OpensciModel.forward to accept **kwargs (num_items_in_batch fix, needed for longsft/LlamaFactory training)

modeling_opensci.py

@@ -191,7 +191,7 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
 
 
 class OpensciMLP(nn.Module):
-    def __init__(self, config, **kwargs):
+    def __init__(self, config):
         super().__init__()
         self.config = config
         self.hidden_size = config.hidden_size
@@ -535,7 +535,7 @@ class OpensciModel(OpensciPreTrainedModel):
         config: OpensciConfig
     """
 
-    def __init__(self, config: OpensciConfig, **kwargs):
+    def __init__(self, config: OpensciConfig):
         super().__init__(config)
         self.padding_idx = config.pad_token_id
         self.vocab_size = config.vocab_size
@@ -573,7 +573,7 @@ class OpensciModel(OpensciPreTrainedModel):
         output_hidden_states: bool | None = None,
         return_dict: bool | None = None,
         cache_position: torch.LongTensor | None = None,
-        # **flash_attn_kwargs: Unpack[FlashAttentionKwargs],
+        **kwargs,
     ) -> tuple | BaseModelOutputWithPast:
         output_attentions = (
             output_attentions if output_attentions is not None else self.config.output_attentions
@@ -811,7 +811,7 @@ class OpensciForCausalLM(OpensciPreTrainedModel, GenerationMixin):
     _tied_weights_keys = ["lm_head.weight"]
     _tp_plan = {"lm_head": "colwise_rep"}
 
-    def __init__(self, config, **kwargs):
+    def __init__(self, config):
         super().__init__(config)
         self.model = OpensciModel(config)
         self.vocab_size = config.vocab_size
@@ -956,7 +956,7 @@ class OpensciForCausalLM(OpensciPreTrainedModel, GenerationMixin):
     Opensci_START_DOCSTRING,
 )
 class OpensciForSequenceClassification(OpensciPreTrainedModel):
-    def __init__(self, config, **kwargs):
+    def __init__(self, config):
         super().__init__(config)
         self.num_labels = config.num_labels
         self.model = OpensciModel(config)