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 policy and value losses.
  
  Args:
      old_logprobs (`torch.FloatTensor`):
          Log probabilities of the model, shape (`batch_size`, `response_length`)
      values (`torch.FloatTensor`):
          Values of the value head, shape (`batch_size`, `response_length`)
      rewards (`torch.FloatTensor`):
          Rewards from the reward model, shape (`batch_size`, `response_length`)
      logits (`torch.FloatTensor`):
          Logits of the model, shape (`batch_size`, `response_length`,
          `vocab_size`)
      v_pred (`torch.FloatTensor`):
          Values of the value head, shape (`batch_size`, `response_length`)
      logprobs (`torch.FloatTensor`):
          Log probabilities of the model, shape (`batch_size`, `response_length`)
code: def loss(
      self,
      old_logprobs: torch.FloatTensor,
      values: torch.FloatTensor,
      logits: torch.FloatTensor,
      vpreds: torch.FloatTensor,
      logprobs: torch.FloatTensor,
      mask: torch.LongTensor,
      advantages: torch.FloatTensor,
      returns: torch.FloatTensor,
  ):
      """
      Calculate policy and value losses.
  
      Args:
          old_logprobs (`torch.FloatTensor`):
              Log probabilities of the model, shape (`batch_size`, `response_length`)
          values (`torch.FloatTensor`):
              Values of the value head, shape (`batch_size`, `response_length`)
          rewards (`torch.FloatTensor`):
              Rewards from the reward model, shape (`batch_size`, `response_length`)
          logits (`torch.FloatTensor`):
              Logits of the model, shape (`batch_size`, `response_length`,
              `vocab_size`)
          v_pred (`torch.FloatTensor`):
              Values of the value head, shape (`batch_size`, `response_length`)
          logprobs (`torch.FloatTensor`):
              Log probabilities of the model, shape (`batch_size`, `response_length`)
      """
  
      vpredclipped = clip_by_value(
          vpreds,
          values - self.cfg.training.ppo_clip_value,
          values + self.cfg.training.ppo_clip_value,
      )
  
      vf_losses1 = (vpreds - returns) ** 2
      vf_losses2 = (vpredclipped - returns) ** 2
      vf_loss = 0.5 * masked_mean(
          torch.max(vf_losses1, vf_losses2), mask
      )
      vf_clipfrac = masked_mean(
          torch.gt(vf_losses2, vf_losses1).double(), mask
      )
  
      ratio = torch.exp(logprobs - old_logprobs)
      pg_losses = -advantages * ratio
      pg_losses2 = -advantages * torch.clamp(
          ratio,
          1.0 - self.cfg.training.ppo_clip_policy,
          1.0 + self.cfg.training.ppo_clip_policy,
      )
  
      pg_loss = masked_mean(torch.max(pg_losses, pg_losses2), mask)
      pg_clipfrac = masked_mean(
          torch.gt(pg_losses2, pg_losses).double(), mask
      )
  
      loss = (
          pg_loss
          + self.cfg.training.scaling_factor_value_loss * vf_loss
      )
  
      entropy = masked_mean(entropy_from_logits(logits), mask)
      approxkl = 0.5 * masked_mean((logprobs - old_logprobs) ** 2, mask)
      policykl = masked_mean(old_logprobs - logprobs, mask)
      return_mean, return_var = masked_mean(returns, mask), masked_var(
          returns, mask
      )
      value_mean, value_var = masked_mean(values, mask), masked_var(
          values, mask
      )
  
      stats = dict(
          loss=dict(
              policy=pg_loss.detach(),
              value=vf_loss.detach(),
              total=loss.detach(),
          ),
          policy=dict(
              entropy=entropy.detach(),
              approxkl=approxkl.detach(),
              policykl=policykl.detach(),
              clipfrac=pg_clipfrac.detach(),
              advantages=advantages.detach(),
              advantages_mean=masked_mean(advantages, mask).detach(),
              ratio=ratio.detach(),
          ),
          returns=dict(
              mean=return_mean.detach(), var=return_var.detach()
          ),
          val=dict(
              vpred=masked_mean(vpreds, mask).detach(),
              error=masked_mean((vpreds - returns) ** 2, mask).detach(),
              clipfrac=vf_clipfrac.detach(),
              mean=value_mean.detach(),
              var=value_var.detach(),
          ),
      )
      return (
          pg_loss,
          self.cfg.training.scaling_factor_value_loss * vf_loss,
          flatten_dict(stats),
      )
  
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 9 for calling the PPOTrainer obj's underlying member method code with example usage: obj.loss(old_logprobs, values, logits, vpreds, logprobs, mask, advantages, returns) and returns: pg_loss,  unknown,  flatten_dict