(
self,
input_ids,
vision_masks,
attention_mask=None,
segment_ids=None,
position_ids=None,
deterministic: bool = True,
init_cache: bool = False,
output_attentions: bool = False,
output_hidden_states: bool = False,
return_dict: bool = True,
)
| 369 | ) |
| 370 | |
| 371 | def __call__( |
| 372 | self, |
| 373 | input_ids, |
| 374 | vision_masks, |
| 375 | attention_mask=None, |
| 376 | segment_ids=None, |
| 377 | position_ids=None, |
| 378 | deterministic: bool = True, |
| 379 | init_cache: bool = False, |
| 380 | output_attentions: bool = False, |
| 381 | output_hidden_states: bool = False, |
| 382 | return_dict: bool = True, |
| 383 | ): |
| 384 | batch_size, seq_length = input_ids.shape |
| 385 | if attention_mask is None: |
| 386 | attention_mask = jnp.ones_like(input_ids) |
| 387 | if segment_ids is None: |
| 388 | segment_ids = jnp.zeros_like(input_ids) |
| 389 | if position_ids is None: |
| 390 | position_ids = jnp.broadcast_to( |
| 391 | jnp.clip(jnp.cumsum(attention_mask, axis=-1) - 1, a_min=0), |
| 392 | (batch_size, seq_length) |
| 393 | ) |
| 394 | |
| 395 | |
| 396 | outputs = self.transformer( |
| 397 | input_ids, |
| 398 | vision_masks, |
| 399 | attention_mask, |
| 400 | segment_ids, |
| 401 | position_ids, |
| 402 | deterministic=deterministic, |
| 403 | init_cache=init_cache, |
| 404 | output_attentions=output_attentions, |
| 405 | output_hidden_states=output_hidden_states, |
| 406 | return_dict=return_dict, |
| 407 | ) |
| 408 | |
| 409 | hidden_states = outputs[0] |
| 410 | |
| 411 | if self.config.tie_vision_embeddings: |
| 412 | shared_kernel = self.transformer.variables["params"]["vte"]["embedding"].T |
| 413 | vision_logits = self.vision_head.apply({"params": {"kernel": shared_kernel}}, hidden_states) |
| 414 | else: |
| 415 | vision_logits = self.vision_head(hidden_states) |
| 416 | |
| 417 | if self.config.tie_word_embeddings: |
| 418 | shared_kernel = self.transformer.variables["params"]["wte"]["embedding"].T |
| 419 | lm_logits = self.lm_head.apply({"params": {"kernel": shared_kernel}}, hidden_states) |
| 420 | else: |
| 421 | lm_logits = self.lm_head(hidden_states) |
| 422 | |
| 423 | if self.config.sample_mode == 'all': |
| 424 | if not return_dict: |
| 425 | return (vision_logits, lm_logits,) + outputs[1:] |
| 426 | |
| 427 | return FlaxCausalLMOutput(logits=(vision_logits, lm_logits), hidden_states=outputs.hidden_states, attentions=outputs.attentions) |
| 428 | elif self.config.sample_mode == 'vision': |
nothing calls this directly
no outgoing calls
no test coverage detected