(self, prompts, max_input_length)
| 400 | ) |
| 401 | |
| 402 | def __call__(self, prompts, max_input_length): |
| 403 | inputs = self.prefix_tokenizer( |
| 404 | prompts, |
| 405 | padding='max_length', |
| 406 | truncation=True, |
| 407 | max_length=max_input_length, |
| 408 | return_tensors='np' |
| 409 | ) |
| 410 | batch = dict( |
| 411 | input_ids=inputs.input_ids, |
| 412 | attention_mask=inputs.attention_mask |
| 413 | ) |
| 414 | with self.mesh: |
| 415 | output, self.sharded_rng = self._forward_generate( |
| 416 | self.params, self.sharded_rng, batch |
| 417 | ) |
| 418 | output = jax.device_get(output) |
| 419 | output_text = [] |
| 420 | for text in list(self.tokenizer.batch_decode(output, skip_special_tokens=True)): |
| 421 | if self.tokenizer.eos_token in text: |
| 422 | text = text.split(self.tokenizer.eos_token, maxsplit=1)[0] |
| 423 | output_text.append(text) |
| 424 | return output_text |
| 425 | |
| 426 | |
| 427 | def main(argv): |
nothing calls this directly
no test coverage detected