加载模型
(self, model_path, peft_path=None, trust_remote_code=True, tensor_parallel_size=1, gpu_memory_utilization=0.25)
| 76 | return params |
| 77 | |
| 78 | def load_model(self, model_path, peft_path=None, trust_remote_code=True, tensor_parallel_size=1, gpu_memory_utilization=0.25): |
| 79 | '''加载模型''' |
| 80 | self.tokenizer = AutoTokenizer.from_pretrained(self.model_path, trust_remote_code=trust_remote_code) |
| 81 | self.model = AutoModelForCausalLM.from_pretrained(self.model_path, device_map="auto", trust_remote_code=trust_remote_code).eval() |
| 82 | if peft_path: |
| 83 | self.model = PeftModel.from_pretrained(self.model, peft_path) |
| 84 | |
| 85 | # self.model = LLM(model=model_path, trust_remote_code=trust_remote_code, tensor_parallel_size=tensor_parallel_size, gpu_memory_utilization=gpu_memory_utilization) |