(
self,
requests: List[List[int]] = None,
generate: bool = False,
max_tokens: int = None,
stop: Optional[List[str]] = None,
**kwargs,
)
| 162 | return encoding |
| 163 | |
| 164 | def _model_generate( |
| 165 | self, |
| 166 | requests: List[List[int]] = None, |
| 167 | generate: bool = False, |
| 168 | max_tokens: int = None, |
| 169 | stop: Optional[List[str]] = None, |
| 170 | **kwargs, |
| 171 | ): |
| 172 | if generate: |
| 173 | kwargs = self.modify_gen_kwargs(kwargs) |
| 174 | sampling_params = SamplingParams(max_tokens=max_tokens, stop=stop, **kwargs) |
| 175 | else: |
| 176 | sampling_params = SamplingParams( |
| 177 | temperature=0, prompt_logprobs=1, max_tokens=1 |
| 178 | ) |
| 179 | if self.data_parallel_size > 1: |
| 180 | requests = [list(x) for x in divide(requests, self.data_parallel_size)] |
| 181 | inputs = [(self.model_args, sampling_params, req) for req in requests] |
| 182 | |
| 183 | with Pool(self.data_parallel_size) as pool: |
| 184 | results = pool.starmap(run_inference_one_model, inputs) |
| 185 | # Invoke ray.shutdown() to prevent hang-ups if subsequent calls required. |
| 186 | ray.shutdown() |
| 187 | # flatten results |
| 188 | return [item for sublist in results for item in sublist] |
| 189 | |
| 190 | outputs = self.model.generate( |
| 191 | prompt_token_ids=requests, |
| 192 | sampling_params=sampling_params, |
| 193 | use_tqdm=True if self.batch_size == "auto" else False, |
| 194 | ) |
| 195 | return outputs |
| 196 | |
| 197 | def _encode_pair( |
| 198 | self, context: str, continuation: str |
no test coverage detected