Initial upload of FreeChunk model with custom code

modeling_freechunker.py

@@ -16,7 +16,7 @@
 """FreeChunker model: Modified from PyTorch XLM-RoBERTa model."""
 from .utils import generate_shifted_matrix
 import math
-from typing import Optional, Tuple, Union, List
+from typing import List, Optional, Tuple, Union, List
 
 import torch
 import torch.utils.checkpoint
@@ -653,7 +653,8 @@ class FreeChunkerModel(FreeChunkerPreTrainedModel):
         input_ids: Optional[torch.Tensor] = None,
         head_mask: Optional[torch.Tensor] = None,
         encoder_hidden_states: Optional[torch.Tensor] = None,
-        encoder_attention_mask: Optional[torch.Tensor] = None
+        encoder_attention_mask: Optional[torch.Tensor] = None,
+        granularities: Optional[List[int]] = None,
     ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]:
         
         # Get input device