changed name from simple-linear-model to simple_linear_model

README.md

@@ -52,7 +52,7 @@ from simple_linear_model.model import *
 from transformers import AutoConfig, AutoModel
 
 # Load the model
-model = AutoModel.from_pretrained("sushilks/simple-linear-model")
+model = AutoModel.from_pretrained("sushilks/simple_linear_model")
 
 # Forward pass
 import torch

config.json

@@ -1,10 +1,10 @@
 {
   "architectures": [
-    "SimpleLinearPreTrainedModel"
+    "SimpleLinearModel"
   ],
   "input_dim": 768,
   "input_size": 768,
-  "model_type": "simple-linear-model",
+  "model_type": "simple_linear_model",
   "output_dim": 512,
   "output_size": 512,
   "torch_dtype": "float32",

model.py

@@ -4,16 +4,16 @@ from transformers import PreTrainedModel
 from transformers import AutoConfig, AutoModel, PretrainedConfig
 
 class SimpleLinearConfig(PretrainedConfig):
-    model_type = "simple-linear-model"
-
+    model_type = "simple_linear_model"
+    _no_split_modules = ["linear"]
     def __init__(self, input_dim=768, output_dim=512, **kwargs):
         super().__init__(**kwargs)
         self.input_dim = input_dim
         self.output_dim = output_dim
 
-class SimpleLinearPreTrainedModel(PreTrainedModel):
+class SimpleLinearModel(PreTrainedModel):
     config_class = SimpleLinearConfig
-
+    _no_split_modules = []
     def __init__(self, config: SimpleLinearConfig):
         super().__init__(config)
         self.linear = nn.Linear(config.input_dim, config.output_dim)
@@ -32,7 +32,7 @@ class SimpleLinearPreTrainedModel(PreTrainedModel):
                     nn.init.zeros_(param)
 
 # Register our config class with AutoConfig
-AutoConfig.register("simple-linear-model", SimpleLinearConfig)
+AutoConfig.register("simple_linear_model", SimpleLinearConfig)
 
 # Register our model class with AutoModel
-AutoModel.register(SimpleLinearConfig, SimpleLinearPreTrainedModel)
+AutoModel.register(SimpleLinearConfig, SimpleLinearModel)