Upload model

MoEConfig.py

@@ -3,8 +3,11 @@ from typing import List
 
 
 class MoEConfig(PretrainedConfig):
-    model_type = "moewrapper"  # モデルの名前を命名？
-    torch_dtype = "float32",
+    model_type = "moewrapper"
+    model_list = [
+        "kanhatakeyama/01b_model_30b_token",
+        "kanhatakeyama/01b_model_30b_token",
+    ]
 
     def __init__(
         self,

MoEModel.py

@@ -1,33 +1,80 @@
 from transformers import PreTrainedModel
-from MoEConfig import MoEConfig
+from .MoEConfig import MoEConfig
 from transformers import AutoModelForCausalLM
 import torch
-
-model_name = "kanhatakeyama/01b_model_30b_token"
+import numpy as np
 
 
 class MoeModel(PreTrainedModel):
     config_class = MoEConfig
+    verbose = True
+    fix_mode = False
 
     def __init__(self, config):
         super().__init__(config)
+        self.model_list = []
+        for model_name in self.config_class.model_list:
+            self.append_model(model_name)
 
-        self.model = None
-        self.set_model()
+        self.set_model_id(0)
 
-    def set_model(self):
+    """
+    def set_model(self, model_name):
         self.model = AutoModelForCausalLM.from_pretrained(
             model_name,
             device_map="auto",
             torch_dtype=torch.float16
         )
+    """
+
+    def append_model(self, model_name):
+        print("loading ", model_name)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            device_map="auto",
+            torch_dtype=torch.float16
+        )
+        self.model_list.append(model)
+
+    # def set_tokenizer(self, tokenizer):
+    #    self.tokenizer = tokenizer
+
+    def set_model_id(self, model_id):
+        self.model = self.model_list[model_id]
+
+    def calc_perplexity(self, tokenized_input):
+        ppl_list = []
+        for model in self.model_list:
+            ppl_list.append(perplexity(model, tokenized_input))
+        return np.array(ppl_list)
+
+    def fix_model(self, model_id):
+        self.set_model_id(model_id)
+        self.fix_mode = True
+
+    def set_flexible_mode(self):
+        self.fix_mode = False
 
     def generate(self, input_ids, attention_mask,
                  **generate_kwargs):
-        if self.model is None:
-            self.set_model()
+
+        if not self.fix_mode:
+            ppl_array = self.calc_perplexity(input_ids)
+            best_model_id = np.where(ppl_array == min(ppl_array))[0][0]
+            self.set_model_id(best_model_id)
+
+            if self.verbose:
+                print(f"model {best_model_id} will be used")
+                print("ppl array: ", ppl_array)
 
         ret = self.model.generate(input_ids=input_ids,
                                   attention_mask=attention_mask,
                                   **generate_kwargs)
         return ret
+
+
+def perplexity(model, tokenized_input) -> torch.Tensor:
+    with torch.inference_mode():
+        output = model(tokenized_input, labels=tokenized_input)
+    ppl = torch.exp(output.loss)
+    return ppl.item()