Upload model

config.json

@@ -4,7 +4,7 @@
   ],
   "attention_probs_dropout_prob": 0.1,
   "auto_map": {
-    "AutoConfig": "config.BigBrainConfig",
+    "AutoConfig": "language_config.BigBrainLanguageConfig",
     "AutoModel": "language.BigBrainLanguageModel"
   },
   "hidden_act": "gelu",

language.py

@@ -6,7 +6,7 @@ from torch.nn import functional as f
 from transformers import PreTrainedModel
 from transformers.activations import ACT2FN
 
-from config import BigBrainConfig
+from language_config import BigBrainLanguageConfig
 
 
 def _make_casual_mask(size: int) -> torch.Tensor:
@@ -26,7 +26,7 @@ class RootMeanSquareNorm(nn.Module):
 
 
 class MultiLayerPerceptron(nn.Module):
-    def __init__(self, config: BigBrainConfig):
+    def __init__(self, config: BigBrainLanguageConfig):
         super().__init__()
         self.config = config
         self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
@@ -72,7 +72,7 @@ class RotaryPositionalEmbedding(nn.Module):
 
 
 class RotaryMultiHeadAttention(nn.Module):
-    def __init__(self, config: BigBrainConfig):
+    def __init__(self, config: BigBrainLanguageConfig):
         super().__init__()
         self.config = config
         self.hidden_size = config.hidden_size
@@ -113,7 +113,7 @@ class RotaryMultiHeadAttention(nn.Module):
 
 
 class BigBrainDecoderLayer(nn.Module):
-    def __init__(self, config: BigBrainConfig):
+    def __init__(self, config: BigBrainLanguageConfig):
         super().__init__()
         self.config = config
         self.self_attn = RotaryMultiHeadAttention(config)
@@ -131,10 +131,10 @@ class BigBrainDecoderLayer(nn.Module):
 
 
 class BigBrainLanguageModel(PreTrainedModel):
-    config_class = BigBrainConfig
+    config_class = BigBrainLanguageConfig
     base_model_prefix = 'big-brain-lm'
 
-    def __init__(self, config: BigBrainConfig = BigBrainConfig()):
+    def __init__(self, config: BigBrainLanguageConfig):
         super().__init__(config)
         self.config = config
         self.tok_embed = nn.Embedding(config.vocab_size, config.hidden_size, config.pad_token_id)

language_config.py

@@ -0,0 +1,43 @@
+from transformers import PretrainedConfig
+
+
+class BigBrainLanguageConfig(PretrainedConfig):
+    model_type = 'big-brain-lm'
+
+    def __init__(
+            self,
+            vocab_size=50265,
+            hidden_size=768,
+            num_hidden_layers=12,
+            num_attention_heads=12,
+            intermediate_size=3072,
+            hidden_act='gelu',
+            hidden_dropout_probability=0.1,
+            attention_probs_dropout_prob=0.1,
+            max_position_embeddings=512,
+            initializer_range=0.02,
+            layer_norm_eps=1e-6,
+            rope_theta=10000,
+            sos_token_id=0,
+            pad_token_id=1,
+            eos_token_id=2,
+            unk_token_id=3,
+            **kwargs
+    ):
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.intermediate_size = intermediate_size
+        self.hidden_act = hidden_act
+        self.hidden_dropout_probability = hidden_dropout_probability
+        self.attention_probs_dropout_prob = attention_probs_dropout_prob
+        self.max_position_embeddings = max_position_embeddings
+        self.initializer_range = initializer_range
+        self.layer_norm_eps = layer_norm_eps
+        self.rope_theta = rope_theta
+        self.sos_token_id = sos_token_id
+        self.pad_token_id = pad_token_id
+        self.eos_token_id = eos_token_id
+        self.unk_token_id = unk_token_id
+        super().__init__(**kwargs)

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:4bc2647de306177bc7c1e1266bb7a6a98f8aee907e18e2b8d0d4457f4c1806d1
+oid sha256:c93547b3cc53ceeeaec4e5754fe86e144c1b90d9e8bbf30e82b9fcb2b53caf85
 size 774713018