update: first

config.json/config.json

@@ -0,0 +1,6 @@
+{
+  "input_dim": 4,
+  "model_type": "test_model",
+  "output_dim": 16,
+  "transformers_version": "4.54.1"
+}

model.py

@@ -0,0 +1,21 @@
+from transformers import PreTrainedModel, PretrainedConfig
+from torch import nn
+
+class TestConfig(PretrainedConfig):
+    model_type = "test_model" 
+    
+    def __init__(self, input_dim=4, output_dim=16, **kwargs):
+        super().__init__(**kwargs)
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+    
+    
+class TestModel(PreTrainedModel):
+    config_class = TestConfig
+    
+    def __init__(self, config):
+        super().__init__(config)
+        self.layer = nn.Linear(config.input_dim, config.output_dim)
+        
+    def forward(self, input):
+        return self.layer(input)

pytorch_model.bin/config.json

@@ -0,0 +1,10 @@
+{
+  "architectures": [
+    "TestModel"
+  ],
+  "input_dim": 4,
+  "model_type": "test_model",
+  "output_dim": 16,
+  "torch_dtype": "float32",
+  "transformers_version": "4.54.1"
+}

pytorch_model.bin/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4d309df6d7ad88e731d193ac5329b87ea597b1175a8b0ba7b33ba05d15821fb7
+size 496

train.py

@@ -0,0 +1,18 @@
+import torch
+from transformers import (WEIGHTS_NAME, CONFIG_NAME)
+
+from model import (TestConfig, TestModel)
+
+
+
+if __name__ == "__main__":
+    config = TestConfig()
+    model = TestModel(config)
+    
+    x = torch.rand(16,4)
+    
+    pred = model(x)
+    print(pred)
+    
+    config.save_pretrained(CONFIG_NAME)
+    model.save_pretrained(WEIGHTS_NAME)