fix param

dnaflash.py

@@ -416,8 +416,18 @@ class FLASHTransformerForPretrained(PreTrainedModel):
             reduce_group_non_causal_attn=config.reduce_group_non_causal_attn
         )
 
-    def forward(self,inputs):
-        logits, x = self.model(inputs["input_ids"], mask=inputs["attention_mask"])
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None
+    )->Union[Tuple, MaskedLMOutput]:
+        logits, x = self.model(input_ids, mask=attention_mask)
         return MaskedLMOutput(logits=logits, hidden_states=x, loss=None, attentions=None)
 
 class FLASHTransformerForSequenceClassification(FLASHTransformerForPretrained):
@@ -438,7 +448,12 @@ class FLASHTransformerForSequenceClassification(FLASHTransformerForPretrained):
     def forward(
         self,
         input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
         labels: Optional[torch.LongTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
     ) -> Union[Tuple, SequenceClassifierOutput]:
         r"""
@@ -450,7 +465,13 @@ class FLASHTransformerForSequenceClassification(FLASHTransformerForPretrained):
 
         # 获取基模型输出
         outputs = super().forward(
-            input_ids
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
         )
         hidden_states = outputs["hidden_states"]
         input_mask_expanded = input_ids["attention_mask"].unsqueeze(-1).expand(hidden_states.size())  # 维度匹配