Revert dict

modeling_flexbert.py

@@ -50,7 +50,7 @@ import os
 import sys
 import warnings
 from dataclasses import dataclass
-from typing import List, Optional, Tuple, Union, Dict
+from typing import List, Optional, Tuple, Union
 
 # Add folder root to path to allow us to use relative imports regardless of what directory the script is run from
 sys.path.append(os.path.dirname(os.path.realpath(__file__)))
@@ -944,36 +944,18 @@ class FlexBertModel(FlexBertPreTrainedModel):
 
     def forward(
         self,
-        features: Dict[str, torch.Tensor],
-        # input_ids: torch.Tensor,
-        # attention_mask: Optional[torch.Tensor] = None,
-        # position_ids: Optional[torch.Tensor] = None,
-        # indices: Optional[torch.Tensor] = None,
-        # cu_seqlens: Optional[torch.Tensor] = None,
-        # max_seqlen: Optional[int] = None,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        indices: Optional[torch.Tensor] = None,
+        cu_seqlens: Optional[torch.Tensor] = None,
+        max_seqlen: Optional[int] = None,
         **kwargs,
     ) -> Tuple[Union[List[torch.Tensor], torch.Tensor], Optional[torch.Tensor]]:
-  
-        if features["attention_mask"] is None:
+        if attention_mask is None:
             attention_mask = torch.ones_like(input_ids)
-        else:
-            attention_mask = features["attention_mask"]
-        input_ids = features["input_ids"]
-        if "position_ids" not in features:
-            position_ids = None
-        else:
-            position_ids = features["position_ids"]
+
         embedding_output = self.embeddings(input_ids, position_ids)
-        if "indices" not in features:
-            indices = None
-        else:
-            indices = features["indices"]
-        if "cu_seqlens" not in features:
-            cu_seqlens = None
-        else:
-            cu_seqlens = features["cu_seqlens"]
-        if "max_seqlen" not in features:
-            max_seqlen = None
 
         encoder_outputs = self.encoder(
             hidden_states=embedding_output,