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