Using dict as input

modeling_flexbert.py

@@ -50,7 +50,7 @@ import os
 import sys
 import warnings
 from dataclasses import dataclass
-from typing import List, Optional, Tuple, Union
+from typing import List, Optional, Tuple, Union, Dict
 
 # 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,18 +944,36 @@ class FlexBertModel(FlexBertPreTrainedModel):
 
     def forward(
         self,
-        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,
+        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,
         **kwargs,
     ) -> Tuple[Union[List[torch.Tensor], torch.Tensor], Optional[torch.Tensor]]:
-        if attention_mask is None:
+  
+        if features["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,