def _convert_opt_causal_lm_to_prefix_lm(model: OPTForCausalLM) -> OPTForCausalLM:
    """Converts an OPT Causal LM to a Prefix LM.

    Supported HuggingFace model classes:
        - `OPTForCausalLM`

    See `convert_hf_causal_lm_to_prefix_lm` for more details.
    """
    if hasattr(model, '_prefix_lm_converted'):
        return model
    assert isinstance(model, OPTForCausalLM)
    assert model.config.add_cross_attention == False, 'Only supports OPT decoder-only models'
    setattr(model, '_original_forward', getattr(model, 'forward'))
    setattr(model, '_original_generate', getattr(model, 'generate'))
    model.model.decoder.bidirectional_mask = None

    def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds, past_key_values_length):
        combined_attention_mask = None
        if input_shape[-1] > 1:
            if self.bidirectional_mask == 'g':
                (bsz, src_length) = input_shape
                combined_attention_mask = torch.zeros((bsz, 1, src_length, src_length + past_key_values_length), dtype=inputs_embeds.dtype, device=inputs_embeds.device)
            else:
                combined_attention_mask = _make_causal_mask_opt(input_shape, inputs_embeds.dtype, past_key_values_length=past_key_values_length).to(inputs_embeds.device)
                if self.bidirectional_mask is not None:
                    assert attention_mask.shape == self.bidirectional_mask.shape
                    expanded_bidirectional_mask = _expand_mask_opt(self.bidirectional_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(inputs_embeds.device)
                    combined_attention_mask = torch.maximum(expanded_bidirectional_mask, combined_attention_mask)
        if attention_mask is not None:
            expanded_attn_mask = _expand_mask_opt(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(inputs_embeds.device)
            combined_attention_mask = expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
        return combined_attention_mask
    setattr(model.model.decoder, '_prepare_decoder_attention_mask', MethodType(_prepare_decoder_attention_mask, model.model.decoder))

    def forward(self: OPTForCausalLM, input_ids: Optional[torch.LongTensor]=None, attention_mask: Optional[torch.Tensor]=None, bidirectional_mask: Optional[torch.ByteTensor]=None, head_mask: Optional[torch.Tensor]=None, past_key_values: Optional[List[torch.FloatTensor]]=None, inputs_embeds: Optional[torch.FloatTensor]=None, labels: Optional[torch.LongTensor]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None):

        def call_og_forward():
            return self._original_forward(input_ids=input_ids, attention_mask=attention_mask, head_mask=head_mask, past_key_values=past_key_values, inputs_embeds=inputs_embeds, labels=labels, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
        if bidirectional_mask is None:
            return call_og_forward()
        self.model.decoder.bidirectional_mask = bidirectional_mask
        try:
            outputs = call_og_forward()
        except:
            self.model.decoder.bidirectional_mask = None
            raise
        self.model.decoder.bidirectional_mask = None
        return outputs

    def generate(self: OPTForCausalLM, *args: tuple, **kwargs: Dict[str, Any]):
        """Wraps original generate to enable PrefixLM-style attention."""
        self.model.decoder.bidirectional_mask = 'g'
        try:
            output = self._original_generate(*args, **kwargs)
        except:
            self.model.decoder.bidirectional_mask = None
            raise
        self.model.decoder.bidirectional_mask = None
        return output
    setattr(model, 'forward', MethodType(forward, model))
    setattr(model, 'generate', MethodType(generate, model))
    setattr(model, '_prefix_lm_converted', True)
    return model
Converts an OPT Causal LM to a Prefix LM.

    Supported HuggingFace model classes:
        - `OPTForCausalLM`

    See `convert_hf_causal_lm_to_prefix_lm` for more details.
    
internvl_chat_llava/llava/model/language_model/mpt/hf_prefixlm_converter.py