Update modeling_footballbert.py

modeling_footballbert.py

@@ -8,7 +8,7 @@ from dataclasses import dataclass
 from transformers import PreTrainedModel
 from transformers.modeling_outputs import BaseModelOutputWithPooling, MaskedLMOutput
 
-#from configuration_footballbert import FootballBERTConfig
+from .configuration_footballbert import FootballBERTConfig
 
 @dataclass
 class FootballBERTOutput(BaseModelOutputWithPooling):
@@ -72,7 +72,7 @@ class PlayerSelfAttention(nn.Module):
 class PlayerTransformerBlock(nn.Module):
     """Standard Transformer encoder layer."""
     
-    def __init__(self, config): #FootballBERTConfig):
+    def __init__(self, config: FootballBERTConfig):
         super().__init__()
         self.attention = PlayerSelfAttention(config.hidden_size, config.num_attention_heads)
 
@@ -96,7 +96,7 @@ class PlayerTransformerBlock(nn.Module):
 
 class PlayerEncoder(nn.Module):
     
-    def __init__(self, config):# FootballBERTConfig):
+    def __init__(self, config: FootballBERTConfig):
 
         super(PlayerEncoder, self).__init__()
 
@@ -133,7 +133,7 @@ class FootballBERTPreTrainedModel(PreTrainedModel):
     and loading pretrained models.
     """
     
-    #config_class = FootballBERTConfig
+    config_class = FootballBERTConfig
     base_model_prefix = "footballbert"
     supports_gradient_checkpointing = False
     
@@ -186,7 +186,7 @@ class FootballBERTModel(FootballBERTPreTrainedModel):
         ```
     """
     
-    def __init__(self, config):# FootballBERTConfig):
+    def __init__(self, config: FootballBERTConfig):
         super().__init__(config)
         self.config = config
         
@@ -259,7 +259,7 @@ class FootballBERTForMaskedPlayerPrediction(FootballBERTPreTrainedModel):
         ```
     """
     
-    def __init__(self, config):# FootballBERTConfig):
+    def __init__(self, config: FootballBERTConfig):
         super().__init__(config)
         self.config = config