| from transformers import PreTrainedModel |
| from .MoEConfig import MoEConfig |
| from transformers import AutoModelForCausalLM |
| import torch |
| 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.set_model_id(0) |
|
|
| """ |
| 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_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 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 |
|
|
|
|
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
|
|
| def perplexity(model, tokenized_input) -> torch.Tensor: |
| with torch.inference_mode(): |
| output = model(tokenized_input.to(device), labels=tokenized_input) |
| ppl = torch.exp(output.loss) |
| return ppl.item() |
|
|
