support classifier dropout

WhereIsAI · Apr 7, 2024 · 7df7de8 · 7df7de8
1 parent 72a678b
commit 7df7de8
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 deletions.
diff --git a/src/billm/modeling_llama.py b/src/billm/modeling_llama.py
@@ -232,7 +232,13 @@ def __init__(self, config):
         super().__init__(config)
         self.num_labels = config.num_labels
         self.model = LlamaModel(config)
-        self.dropout = nn.Dropout(0.1)
+        if hasattr(config, "classifier_dropout") and config.classifier_dropout is not None:
+            classifier_dropout = config.classifier_dropout
+        elif hasattr(config, "hidden_dropout") and config.hidden_dropout is not None:
+            classifier_dropout = config.hidden_dropout
+        else:
+            classifier_dropout = 0.1
+        self.dropout = nn.Dropout(classifier_dropout)
         self.classifier = nn.Linear(config.hidden_size, config.num_labels)
 
         # Initialize weights and apply final processing

diff --git a/src/billm/modeling_mistral.py b/src/billm/modeling_mistral.py
@@ -206,7 +206,13 @@ def __init__(self, config):
         super().__init__(config)
         self.num_labels = config.num_labels
         self.model = MistralModel(config)
-        self.dropout = nn.Dropout(0.1)
+        if hasattr(config, "classifier_dropout") and config.classifier_dropout is not None:
+            classifier_dropout = config.classifier_dropout
+        elif hasattr(config, "hidden_dropout") and config.hidden_dropout is not None:
+            classifier_dropout = config.hidden_dropout
+        else:
+            classifier_dropout = 0.1
+        self.dropout = nn.Dropout(classifier_dropout)
         self.classifier = nn.Linear(config.hidden_size, config.num_labels)
 
         # Initialize weights and apply final processing