(self, model_path: str, peft_path: str = None, template: str = "default", trust_remote_code=True, tensor_parallel_size=1, gpu_memory_utilization=0.25)
| 15 | |
| 16 | class BaiChuanModel(ToolModel): |
| 17 | def __init__(self, model_path: str, peft_path: str = None, template: str = "default", trust_remote_code=True, tensor_parallel_size=1, gpu_memory_utilization=0.25): |
| 18 | self.model_path = model_path |
| 19 | self.peft_path = peft_path |
| 20 | self.template = template |
| 21 | self.trust_remote_code = trust_remote_code |
| 22 | self.tensor_parallel_size = tensor_parallel_size |
| 23 | self.gpu_memory_utilization = gpu_memory_utilization |
| 24 | self.generation_config = GenerationConfig.from_pretrained(model_path) |
| 25 | self.load_model(self.model_path, self.peft_path, self.trust_remote_code, self.tensor_parallel_size, self.gpu_memory_utilization) |
| 26 | |
| 27 | def generate( |
| 28 | self, prompts: str, |
nothing calls this directly
no test coverage detected