bugfix

WhereIsAI · Apr 27, 2024 · 09e06c5 · 09e06c5
1 parent 073e443
commit 09e06c5
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/src/billm/modeling_openelm.py b/src/billm/modeling_openelm.py
@@ -996,7 +996,7 @@ def __init__(self, config):
         super().__init__(config)
         self.num_labels = config.num_labels
         self.model = OpenELMModel(config)
-        self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
+        self.score = nn.Linear(config.model_dim, self.num_labels, bias=False)
 
         # Initialize weights and apply final processing
         self.post_init()
@@ -1110,7 +1110,7 @@ def __init__(self, config):
         else:
             classifier_dropout = 0.1
         self.dropout = nn.Dropout(classifier_dropout)
-        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+        self.classifier = nn.Linear(config.model_dim, config.num_labels)
 
         # Initialize weights and apply final processing
         self.post_init()

diff --git a/tests/test_modeling_openelm.py b/tests/test_modeling_openelm.py
@@ -8,7 +8,7 @@ def test_openelm_model():
     from billm import OpenELMModel, OpenELMConfig
 
     model = OpenELMModel(OpenELMConfig(vocab_size=128,
-                                       head_size=32,
+                                       head_dim=32,
                                        num_transformer_layers=2))
     assert model is not None
 
@@ -20,7 +20,7 @@ def test_biopenelm_model():
     from billm import OpenELMModel, OpenELMConfig
 
     model = OpenELMModel(OpenELMConfig(vocab_size=128,
-                                       head_size=32,
+                                       head_dim=32,
                                        num_transformer_layers=2))
     assert model is not None
 
@@ -33,7 +33,7 @@ def test_biopenelm_lm():
     from billm import OpenELMForCausalLM, OpenELMConfig
 
     model = OpenELMForCausalLM(OpenELMConfig(vocab_size=128,
-                                             head_size=32,
+                                             head_dim=32,
                                              num_transformer_layers=2))
     assert model is not None
     assert len(model.model.bidirectionas) > 0
@@ -46,7 +46,7 @@ def test_biopenelm_seq_clf():
     from billm import OpenELMForSequenceClassification, OpenELMConfig
 
     model = OpenELMForSequenceClassification(OpenELMConfig(vocab_size=128,
-                                                           head_size=32,
+                                                           head_dim=32,
                                                            num_transformer_layers=2))
     assert model is not None
     assert len(model.model.bidirectionas) > 0
@@ -59,7 +59,7 @@ def test_biopenelm_token_clf():
     from billm import OpenELMForTokenClassification, OpenELMConfig
 
     model = OpenELMForTokenClassification(OpenELMConfig(vocab_size=128,
-                                                        head_size=32,
+                                                        head_dim=32,
                                                         num_transformer_layers=2))
     assert model is not None
     assert len(model.model.bidirectionas) > 0