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