Upload model

config.json

@@ -1,14 +1,14 @@
-{
-  "architectures": [
-    "TestModel"
-  ],
-  "auto_map": {
-    "AutoConfig": "configuration_test.TestConfig",
-    "AutoModel": "modeling_test.TestModel"
-  },
-  "input_dim": 10,
-  "model_type": "my_test_model",
-  "output_dim": 5,
-  "torch_dtype": "float32",
-  "transformers_version": "4.39.1"
-}
+{
+  "architectures": [
+    "TestModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_test.TestConfig",
+    "AutoModel": "modeling_test.TestModel"
+  },
+  "input_dim": 10,
+  "model_type": "my_test_model",
+  "output_dim": 5,
+  "torch_dtype": "float32",
+  "transformers_version": "4.37.2"
+}

configuration_test.py

@@ -1,16 +1,16 @@
-from transformers import PretrainedConfig
-from typing import List
-
-
-class TestConfig(PretrainedConfig):
-    model_type = "my_test_model"
-
-    def __init__(
-        self,
-        input_dim: int = 20,
-        output_dim: int = 10,
-        **kwargs,
-    ):
-        self.input_dim = input_dim
-        self.output_dim = output_dim
-        super().__init__(**kwargs)
+from transformers import PretrainedConfig
+from typing import List
+
+
+class TestConfig(PretrainedConfig):
+    model_type = "my_test_model"
+
+    def __init__(
+        self,
+        input_dim: int = 20,
+        output_dim: int = 10,
+        **kwargs,
+    ):
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+        super().__init__(**kwargs)

modeling_test.py

@@ -1,20 +1,20 @@
-from transformers import PreTrainedModel
-from .configuration_test import TestConfig
-import torch.nn as nn
-from transformers import AutoModelForMaskedLM, AutoConfig
-from transformers import AutoModelForSequenceClassification
-
-
-class TestModel(PreTrainedModel):
-    config_class = TestConfig
-
-    def __init__(self, config: TestConfig):
-        super().__init__(config)
-        self.input_dim = config.input_dim
-        self.model1 = nn.Linear(config.input_dim, config.output_dim)
-        self.model2 = AutoModelForMaskedLM.from_config(
-            AutoConfig.from_pretrained("albert/albert-base-v2")
-        )
-
-    def forward(self, tensor):
-        return self.model1(tensor)
+from transformers import PreTrainedModel
+from .configuration_test import TestConfig
+import torch.nn as nn
+from transformers import AutoModelForMaskedLM, AutoConfig
+from transformers import AutoModelForSequenceClassification
+
+
+class TestModel(PreTrainedModel):
+    config_class = TestConfig
+
+    def __init__(self, config: TestConfig):
+        super().__init__(config)
+        self.input_dim = config.input_dim
+        self.model1 = nn.Linear(config.input_dim, config.output_dim)
+        self.model2 = AutoModelForMaskedLM.from_config(
+            AutoConfig.from_pretrained("albert/albert-base-v2")
+        )
+
+    def forward(self, tensor):
+        return self.model1(tensor)

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:0f88faa6c688ec9ab7b59910120c1e3bf97817b8d41a68104f608db5b2830260
-size 44897099
+oid sha256:c636fd32fccd184fcfaf1c852070f2144848824682ae2048a02e0658ce545a24
+size 44897544