(self, prompts, max_n_frames)
| 220 | ) |
| 221 | |
| 222 | def __call__(self, prompts, max_n_frames): |
| 223 | batch = self.construct_input(prompts, max_n_frames) |
| 224 | with self.mesh: |
| 225 | output, self.sharded_rng = self._forward_generate( |
| 226 | self.params, self.sharded_rng, batch |
| 227 | ) |
| 228 | output = jax.device_get(output) |
| 229 | output_text = [] |
| 230 | for text in list(self.tokenizer.batch_decode(output, skip_special_tokens=True)): |
| 231 | if self.tokenizer.eos_token in text: |
| 232 | text = text.split(self.tokenizer.eos_token, maxsplit=1)[0] |
| 233 | output_text.append(text) |
| 234 | return output_text |
| 235 | |
| 236 | def main(argv): |
| 237 | assert FLAGS.prompt != '' |
nothing calls this directly
no test coverage detected