| 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 | |