class_docstr: The PPOTrainer uses Proximal Policy Optimization to optimise language models.
  Note, this trainer is heavily inspired by the original OpenAI learning to summarize
  work here: https://github.com/openai/summarize-from-feedback
  
  Attributes:
      **cfg** (`LLM Studio Config`) -- Experiment configuration object. Check the
          documentation of `LLM Studio Config` for more details.
      **model** (`PreTrainedModelWrapper`) -- Model to be optimized, Hugging Face
          transformer model with a value head. Check the documentation of
          `PreTrainedModelWrapper` for more details.
      **tokenizer** (`Union[PreTrainedTokenizer, PreTrainedTokenizerFast]`)
          Tokenizer to be used for encoding the data. Check the documentation of
          `transformers.PreTrainedTokenizer` and
          `transformers.PreTrainedTokenizerFast` for more details.
      **optimizer** (`torch.optim.Optimizer`) -- Optimizer to be used for training.
      **lr_scheduler** (`torch.optim.lr_scheduler`) -- Learning rate scheduler to be
          used for training.
      **scaler** (`torch.cuda.amp.GradScaler`) -- Gradient scaler to be used for
          training.
class_docstr_tok: ["The","PPOTrainer","uses","Proximal","Policy","Optimization","to","optimise","language","models",".","Note",",","this","trainer","is","heavily","inspired","by","the","original","OpenAI","learning","to","summarize","work","here",":","https",":","\/\/github.com\/openai\/summarize-from-feedback","Attributes",":","*","*","cfg","*","*","(","`","LLM","Studio","Config","`",")","--","Experiment","configuration","object",".","Check","the","documentation","of","`","LLM","Studio","Config","`","for","more","details",".","*","*","model","*","*","(","`","PreTrainedModelWrapper","`",")","--","Model","to","be","optimized",",","Hugging","Face","transformer","model","with","a","value","head",".","Check","the","documentation","of","`","PreTrainedModelWrapper","`","for","more","details",".","*","*","tokenizer","*","*","(","`","Union","[","PreTrainedTokenizer",",","PreTrainedTokenizerFast","]","`",")","Tokenizer","to","be","used","for","encoding","the","data",".","Check","the","documentation","of","`","transformers.PreTrainedTokenizer","`","and","`","transformers.PreTrainedTokenizerFast","`","for","more","details",".","*","*","optimizer","*","*","(","`","torch.optim.Optimizer","`",")","--","Optimizer","to","be","used","for","training",".","*","*","lr_scheduler","*","*","(","`","torch.optim.lr_scheduler","`",")","--","Learning","rate","scheduler","to","be","used","for","training",".","*","*","scaler","*","*","(","`","torch.cuda.amp.GradScaler","`",")","--","Gradient","scaler","to","be","used","for","training","."]
docstr: Calculate model outputs in multiple batches.
  
  Args:
      queries (`torch.LongTensor`):
          List of tensors containing the encoded queries, shape (`batch_size`,
          `query_length`)
      responses (`torch.LongTensor`):
          List of tensors containing the encoded responses, shape (`batch_size`,
          `response_length`)
      return_logits (`bool`, *optional*, defaults to `False`):
          Whether to return all_logits. Set to `False` if logits are not needed
          to reduce memory consumption.
      return_values (`bool`, *optional*, defaults to `True`):
          Whether to return values. Set to `False` if values are not needed to
          reduce memory consumption.
  Returns:
      (tuple):
          - all_logprobs (`torch.FloatTensor`): Log probabilities of the
              responses, shape (`batch_size`, `response_length`)
          - all_ref_logprobs (`torch.FloatTensor`): Log probabilities of the
              responses, shape (`batch_size`, `response_length`)
          - all_values (`torch.FloatTensor`): Values of the responses, shape
              (`batch_size`, `response_length`)
code: def batched_forward_pass(
      self,
      model,
      queries: torch.Tensor,
      responses: torch.Tensor,
      model_inputs: dict,
      return_logits: bool = False,
      return_values: bool = True,
  ):
      """
      Calculate model outputs in multiple batches.
  
      Args:
          queries (`torch.LongTensor`):
              List of tensors containing the encoded queries, shape (`batch_size`,
              `query_length`)
          responses (`torch.LongTensor`):
              List of tensors containing the encoded responses, shape (`batch_size`,
              `response_length`)
          return_logits (`bool`, *optional*, defaults to `False`):
              Whether to return all_logits. Set to `False` if logits are not needed
              to reduce memory consumption.
          return_values (`bool`, *optional*, defaults to `True`):
              Whether to return values. Set to `False` if values are not needed to
              reduce memory consumption.
      Returns:
          (tuple):
              - all_logprobs (`torch.FloatTensor`): Log probabilities of the
                  responses, shape (`batch_size`, `response_length`)
              - all_ref_logprobs (`torch.FloatTensor`): Log probabilities of the
                  responses, shape (`batch_size`, `response_length`)
              - all_values (`torch.FloatTensor`): Values of the responses, shape
                  (`batch_size`, `response_length`)
      """
  
      bs = len(queries)
      ppo_bs = self.cfg.training.ppo_batch_size
  
      all_logprobs = []
      all_logits = []
      all_masks = []
      all_values = []
  
      for i in range(math.ceil(bs / ppo_bs)):
          model_inputs_batch = {
              key: value[i * ppo_bs : (i + 1) * ppo_bs]
              for key, value in model_inputs.items()
          }
  
          query_batch = queries[i * ppo_bs : (i + 1) * ppo_bs]
          response_batch = responses[i * ppo_bs : (i + 1) * ppo_bs]
  
          with autocast(enabled=self.cfg.environment.mixed_precision):
              outputs = model(
                  model_inputs_batch,
                  padding=False,
              )
  
          logits = outputs["logits"]
          values = outputs["value"]
  
          input_ids = model_inputs_batch["input_ids"]
          attention_mask = model_inputs_batch["attention_mask"]
  
          logprobs = logprobs_from_logits(
              logits[:, :-1, :], input_ids[:, 1:]
          )
          masks = torch.zeros_like(attention_mask)
          masks[:, :-1] = attention_mask[:, 1:]
  
          for j in range(len(query_batch)):
              start = len(query_batch[j]) - 1
              if attention_mask[j, 0] == 0:  # offset left padding
                  start += attention_mask[j, :].nonzero()[0]
              end = start + len(response_batch[j])
  
              masks[j, :start] = 0
              masks[j, end:] = 0
  
          if return_logits:
              all_logits.append(logits)
          else:
              del logits
          if return_values:
              all_values.append(values)
          else:
              del values
          all_logprobs.append(logprobs)
          all_masks.append(masks)
  
      del outputs
  
      return (
          torch.cat(all_logprobs),
          torch.cat(all_logits)[:, :-1] if return_logits else None,
          torch.cat(all_values)[:, :-1] if return_values else None,
          torch.cat(all_masks)[:, :-1],
      )
  
imports: import math
  import time
  import warnings
  from collections.abc import Mapping
  from typing import Any, Dict, List, Optional, Union
  import numpy
  import torch
  import torch.nn.functional
  from datasets import Dataset
  from huggingface_hub import PyTorchModelHubMixin
  from torch.cuda.amp import autocast
  from torch.nn.utils.rnn import pad_sequence
  from torch.utils.data import DataLoader
  from transformers import DataCollatorForLanguageModeling, PreTrainedTokenizer, PreTrainedTokenizerFast
  from llm_studio.src.utils.modeling_utils import unwrap_model
label_desc: Use image node_id 5 for calling the PPOTrainer obj's underlying member method code with example usage: obj.batched_forward_pass(model, queries, responses, model_inputs, return_logits, return_values) and returns: torch,  unknown,  unknown,  torch