| 195 | |
| 196 | @cached_property |
| 197 | def _forward_generate(self): |
| 198 | def fn(params, rng, batch): |
| 199 | batch = with_sharding_constraint(batch, PS(('dp', 'fsdp'), 'sp')) |
| 200 | rng_generator = JaxRNG(rng) |
| 201 | output = self.model.generate( |
| 202 | batch['input_ids'], |
| 203 | vision_masks=batch['vision_masks'], |
| 204 | attention_mask=batch['attention_mask'], |
| 205 | params=params['params'], |
| 206 | prng_key=rng_generator(), |
| 207 | generation_config=GenerationConfig( |
| 208 | max_new_tokens=self.block_size, |
| 209 | pad_token_id=self.tokenizer.pad_token_id, |
| 210 | eos_token_id=self.tokenizer.eos_token_id, |
| 211 | temperature=FLAGS.temperature, |
| 212 | do_sample=True, |
| 213 | ) |
| 214 | ).sequences[:, batch['input_ids'].shape[1]:] |
| 215 | return output, rng_generator() |
| 216 | return pjit( |
| 217 | fn, |
| 218 | in_shardings=(self.model_ps, PS(), PS()), |
| 219 | out_shardings=(PS(), PS()) |
| 220 | ) |
| 221 | |
| 222 | def __call__(self, prompts, max_n_frames): |
| 223 | batch = self.construct_input(prompts, max_n_frames) |