Upload folder using huggingface_hub

modeling_havelock.py

@@ -2,11 +2,12 @@
 
 import torch
 import torch.nn as nn
-from transformers import AutoConfig, AutoModel, PreTrainedModel, PretrainedConfig
+from transformers import AutoConfig, AutoModel, PretrainedConfig, PreTrainedModel
 
 
 class HavelockTokenConfig(PretrainedConfig):
     """Config that wraps any backbone config + our custom fields."""
+
     model_type = "havelock_token_classifier"
 
     def __init__(self, num_types: int = 1, use_crf: bool = False, **kwargs):
@@ -18,7 +19,9 @@ class HavelockTokenConfig(PretrainedConfig):
 class HavelockTokenClassifier(PreTrainedModel):
     config_class = HavelockTokenConfig
 
-    def __init__(self, config: HavelockTokenConfig, backbone: PreTrainedModel | None = None):
+    def __init__(
+        self, config: HavelockTokenConfig, backbone: PreTrainedModel | None = None
+    ):
         super().__init__(config)
         self.num_types = config.num_types
         self.use_crf = config.use_crf
@@ -29,7 +32,7 @@ class HavelockTokenClassifier(PreTrainedModel):
         else:
             self.backbone = AutoModel.from_config(config)
 
-        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.dropout = nn.Dropout(getattr(config, "hidden_dropout_prob", 0.1))
         self.classifier = nn.Linear(config.hidden_size, config.num_types * 3)
 
         if self.use_crf:
@@ -75,7 +78,9 @@ class HavelockTokenClassifier(PreTrainedModel):
             mask = (
                 attention_mask.bool()
                 if attention_mask is not None
-                else torch.ones(logits.shape[:2], dtype=torch.bool, device=logits.device)
+                else torch.ones(
+                    logits.shape[:2], dtype=torch.bool, device=logits.device
+                )
             )
             return self.crf.decode(logits, mask)
-        return logits.argmax(dim=-1)
+        return logits.argmax(dim=-1)