Upload model

MoEConfig.py

@@ -0,0 +1,13 @@
+from transformers import PretrainedConfig
+from typing import List
+
+
+class MoEConfig(PretrainedConfig):
+    model_type = "moewrapper"  # モデルの名前を命名？
+    torch_dtype = "float32",
+
+    def __init__(
+        self,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)

MoEModel.py

@@ -0,0 +1,33 @@
+from transformers import PreTrainedModel
+from MoEConfig import MoEConfig
+from transformers import AutoModelForCausalLM
+import torch
+
+model_name = "kanhatakeyama/01b_model_30b_token"
+
+
+class MoeModel(PreTrainedModel):
+    config_class = MoEConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.model = None
+        self.set_model()
+
+    def set_model(self):
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            device_map="auto",
+            torch_dtype=torch.float16
+        )
+
+    def generate(self, input_ids, attention_mask,
+                 **generate_kwargs):
+        if self.model is None:
+            self.set_model()
+
+        ret = self.model.generate(input_ids=input_ids,
+                                  attention_mask=attention_mask,
+                                  **generate_kwargs)
+        return ret

config.json

@@ -0,0 +1,12 @@
+{
+  "architectures": [
+    "MoeModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "MoEConfig.MoEConfig",
+    "AutoModelForCausalLM": "MoEModel.MoeModel"
+  },
+  "model_type": "moewrapper",
+  "torch_dtype": "float16",
+  "transformers_version": "4.35.0"
+}

generation_config.json

@@ -0,0 +1,4 @@
+{
+  "_from_model_config": true,
+  "transformers_version": "4.35.0"
+}

model.safetensors

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