| 382 | |
| 383 | @cached_property |
| 384 | def _forward_generate(self): |
| 385 | def fn(params, rng, batch): |
| 386 | batch = with_sharding_constraint(batch, PS(('dp', 'fsdp'), 'sp')) |
| 387 | rng_generator = JaxRNG(rng) |
| 388 | output = self.model.generate( |
| 389 | batch['input_ids'], |
| 390 | attention_mask=batch['attention_mask'], |
| 391 | params=params['params'], |
| 392 | prng_key=rng_generator(), |
| 393 | generation_config=GenerationConfig( |
| 394 | max_new_tokens=self.block_size, |
| 395 | pad_token_id=self.tokenizer.pad_token_id, |
| 396 | eos_token_id=self.tokenizer.eos_token_id, |
| 397 | temperature=0., |
| 398 | do_sample=False, |
| 399 | num_beams=1, |
| 400 | top_k=50, |
| 401 | top_p=1.0, |
| 402 | ) |
| 403 | ).sequences[:, batch['input_ids'].shape[1]:] |
| 404 | return output, rng_generator() |
| 405 | return pjit( |
| 406 | fn, |
| 407 | in_shardings=(self.model_ps, PS(), PS()), |
| 408 | out_shardings=(PS(), PS()) |
| 409 | ) |
| 410 | |
| 411 | def __call__(self, prompts, max_input_length): |
| 412 | inputs = self.prefix_tokenizer( |