Update model.py

model.py

@@ -6,8 +6,11 @@ import json
 from safetensors.torch import save_file, load_file
 import os
 from huggingface_hub import PyTorchModelHubMixin
+from transformers import PretrainedConfig, PreTrainedModel
 
-class CosmoFormerConfig:
+class CosmoFormerConfig(PretrainedConfig):
+    model_type = "cosmoformer"
+    
     def __init__(
         self,
         d_model: int = 256,
@@ -18,7 +21,9 @@ class CosmoFormerConfig:
         num_layers: int = 6,
         vocab_size: int = 65400,
         max_len: int = 2048,
+        **kwargs
     ):
+        super().__init__(**kwargs)
         self.d_model = d_model
         self.d_ff = d_ff
         self.dropout = dropout
@@ -183,9 +188,11 @@ class DecoderLayer(nn.Module):
         return x
 
 
-class CosmoFormer(nn.Module, PyTorchModelHubMixin):
-    def __init__(self, config):
-        super().__init__()
+class CosmoFormer(PreTrainedModel):
+    config_class = CosmoFormerConfig
+
+    def __init__(self, config: CosmoFormerConfig):
+        super().__init__(config)
         self.config = config
         self.d_model = config.d_model
         self.vocab_size = config.vocab_size