class_docstr: An autoregressive model with a value head in addition to the language model head.
  This class inherits from `~trl.PreTrainedModelWrapper` and wraps a
  `transformers.PreTrainedModel` class. The wrapper class supports classic functions
  such as `from_pretrained`, `push_to_hub` and `generate`. To call a method of the wrapped
  model, simply manipulate the `pretrained_model` attribute of this class.
  
  Class attributes:
      - **transformers_parent_class** (`transformers.PreTrainedModel`) -- The parent class of the wrapped model. This
          should be set to `transformers.AutoModelForCausalLM` for this class.
      - **lm_head_namings** (`tuple`) -- A tuple of strings that are used to identify the language model head of the
          wrapped model. This is set to `("lm_head", "embed_out")` for this class but can be changed for other models
          in the future
      - **supported_args** (`tuple`) -- A tuple of strings that are used to identify the arguments that are supported
          by the `ValueHead` class. Currently, the supported args are:
          - **summary_dropout_prob** (`float`, `optional`, defaults to `None`) -- The dropout probability for the
              `ValueHead` class.
          - **v_head_initializer_range** (`float`, `optional`, defaults to `0.2`) -- The initializer range for the
              `ValueHead` if a specific initialization strategy is selected.
          - **v_head_init_strategy** (`str`, `optional`, defaults to `None`) -- The initialization strategy for the
              `ValueHead`. Currently, the supported strategies are:
              - **`None`** -- Initializes the weights of the `ValueHead` with a random distribution. This is the default
                  strategy.
              - **"normal"** -- Initializes the weights of the `ValueHead` with a normal distribution.
class_docstr_tok: ["An","autoregressive","model","with","a","value","head","in","addition","to","the","language","model","head",".","This","class","inherits","from","`","~trl.PreTrainedModelWrapper","`","and","wraps","a","`","transformers.PreTrainedModel","`","class",".","The","wrapper","class","supports","classic","functions","such","as","`","from_pretrained","`",",","`","push_to_hub","`","and","`","generate","`",".","To","call","a","method","of","the","wrapped","model",",","simply","manipulate","the","`","pretrained_model","`","attribute","of","this","class",".","Class","attributes",":","-","*","*","transformers_parent_class","*","*","(","`","transformers.PreTrainedModel","`",")","--","The","parent","class","of","the","wrapped","model",".","This","should","be","set","to","`","transformers.AutoModelForCausalLM","`","for","this","class",".","-","*","*","lm_head_namings","*","*","(","`","tuple","`",")","--","A","tuple","of","strings","that","are","used","to","identify","the","language","model","head","of","the","wrapped","model",".","This","is","set","to","`","(","``","lm_head","''",",","``","embed_out","''",")","`","for","this","class","but","can","be","changed","for","other","models","in","the","future","-","*","*","supported_args","*","*","(","`","tuple","`",")","--","A","tuple","of","strings","that","are","used","to","identify","the","arguments","that","are","supported","by","the","`","ValueHead","`","class",".","Currently",",","the","supported","args","are",":","-","*","*","summary_dropout_prob","*","*","(","`","float","`",",","`","optional","`",",","defaults","to","`","None","`",")","--","The","dropout","probability","for","the","`","ValueHead","`","class",".","-","*","*","v_head_initializer_range","*","*","(","`","float","`",",","`","optional","`",",","defaults","to","`","0.2","`",")","--","The","initializer","range","for","the","`","ValueHead","`","if","a","specific","initialization","strategy","is","selected",".","-","*","*","v_head_init_strategy","*","*","(","`","str","`",",","`","optional","`",",","defaults","to","`","None","`",")","--","The","initialization","strategy","for","the","`","ValueHead","`",".","Currently",",","the","supported","strategies","are",":","-","*","*","`","None","`","*","*","--","Initializes","the","weights","of","the","`","ValueHead","`","with","a","random","distribution",".","This","is","the","default","strategy",".","-","*","*","''","normal","''","*","*","--","Initializes","the","weights","of","the","`","ValueHead","`","with","a","normal","distribution","."]
docstr: Applies a forward pass to the wrapped model and returns the logits of the value head.
  
  Args:
      input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
          Indices of input sequence tokens in the vocabulary.
      past_key_values (`tuple(tuple(torch.FloatTensor))`, `optional`):
          Contains pre-computed hidden-states (key and values in the attention blocks) as computed by the model
          (see `past_key_values` input) to speed up sequential decoding.
      attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, `optional`):
          Mask to avoid performing attention on padding token indices. Mask values selected in ``[0, 1]``:
          - 1 for tokens that are **not masked**,
          - 0 for tokens that are **masked**.
      kwargs (`dict`, `optional`):
          Additional keyword arguments, that are passed to the wrapped model.
docstr_tok: ["Applies","a","forward","pass","to","the","wrapped","model","and","returns","the","logits","of","the","value","head",".","Args",":","input_ids","(","`","torch.LongTensor","`","of","shape","`","(","batch_size",",","sequence_length",")","`",")",":","Indices","of","input","sequence","tokens","in","the","vocabulary",".","past_key_values","(","`","tuple","(","tuple","(","torch.FloatTensor",")",")","`",",","`","optional","`",")",":","Contains","pre-computed","hidden-states","(","key","and","values","in","the","attention","blocks",")","as","computed","by","the","model","(","see","`","past_key_values","`","input",")","to","speed","up","sequential","decoding",".","attention_mask","(","`","torch.FloatTensor","`","of","shape","`","(","batch_size",",","sequence_length",")","`",",","`","optional","`",")",":","Mask","to","avoid","performing","attention","on","padding","token","indices",".","Mask","values","selected","in","``","[","0",",","1","]","``",":","-","1","for","tokens","that","are","*","*","not","masked","*","*",",","-","0","for","tokens","that","are","*","*","masked","*","*",".","kwargs","(","`","dict","`",",","`","optional","`",")",":","Additional","keyword","arguments",",","that","are","passed","to","the","wrapped","model","."]
code: def forward(
      self,
      input_ids=None,
      past_key_values=None,
      attention_mask=None,
      **kwargs,
  ):
      r"""
      Applies a forward pass to the wrapped model and returns the logits of the value head.
  
      Args:
          input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
              Indices of input sequence tokens in the vocabulary.
          past_key_values (`tuple(tuple(torch.FloatTensor))`, `optional`):
              Contains pre-computed hidden-states (key and values in the attention blocks) as computed by the model
              (see `past_key_values` input) to speed up sequential decoding.
          attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, `optional`):
              Mask to avoid performing attention on padding token indices. Mask values selected in ``[0, 1]``:
              - 1 for tokens that are **not masked**,
              - 0 for tokens that are **masked**.
          kwargs (`dict`, `optional`):
              Additional keyword arguments, that are passed to the wrapped model.
      """
      kwargs[
          "output_hidden_states"
      ] = True  # this had already been set in the LORA / PEFT examples
      kwargs["past_key_values"] = past_key_values
  
      if (
          self.is_peft_model
          and self.pretrained_model.active_peft_config.peft_type
          == "PREFIX_TUNING"
      ):
          kwargs.pop("past_key_values")
  
      base_model_output = self.pretrained_model(
          input_ids=input_ids,
          attention_mask=attention_mask,
          **kwargs,
      )
  
      last_hidden_state = base_model_output.hidden_states[-1]
      lm_logits = base_model_output.logits
      loss = base_model_output.loss
  
      if last_hidden_state.device != self.v_head.summary.weight.device:
          last_hidden_state = last_hidden_state.to(
              self.v_head.summary.weight.device
          )
  
      value = self.v_head(last_hidden_state).squeeze(-1)
  
      # force upcast in fp32 if logits are in half-precision
      if lm_logits.dtype != torch.float32:
          lm_logits = lm_logits.float()
  
      return (lm_logits, loss, value)
  
imports: import torch
  import torch.nn
  from transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM
  from .modeling_base import PreTrainedModelWrapper
label_desc: Use image node_id 3 for calling the AutoModelForCausalLMWithValueHead obj's underlying member method code with example usage: obj.forward(input_ids, past_key_values, attention_mask) and returns: lm_logits,  loss,  value