(self, model, device_num, llama = False, tensor_parallel_size = 4, temperature=0.7, n = 200)
| 158 | |
| 159 | class LocalModel(object): |
| 160 | def __init__(self, model, device_num, llama = False, tensor_parallel_size = 4, temperature=0.7, n = 200) -> None: |
| 161 | self.name = model |
| 162 | self.sampling_params = SamplingParams(temperature=temperature, n = n, max_tokens = 512) |
| 163 | #self.sampling_params = SamplingParams(temperature=0, n = 1, max_tokens = 512) |
| 164 | self.model = LLM(model, trust_remote_code = True, tensor_parallel_size= tensor_parallel_size) |
| 165 | self.tokenizer = self.model.get_tokenizer() |
| 166 | |
| 167 | def get_prompt_length(self, prompt): |
| 168 | return len(self.tokenizer(prompt).input_ids) |
nothing calls this directly
no outgoing calls
no test coverage detected