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: Record training step statistics.
  
  
  Args:
      kl_coef (`float`):
          KL coefficient
      data (`dict`):
          Dictionary of training step data
  
  Returns:
      stats (`dict`):
          Dictionary of training step statistics
code: def record_step_stats(self, kl_coef: float, **data):
      """
      Record training step statistics.
  
  
      Args:
          kl_coef (`float`):
              KL coefficient
          data (`dict`):
              Dictionary of training step data
  
      Returns:
          stats (`dict`):
              Dictionary of training step statistics
      """
      mask = data.pop("masks")
  
      kl_list = ((data["logprobs"] - data["ref_logprobs"]) * mask).sum(
          axis=-1
      )
      mean_kl = kl_list.mean()
      mean_entropy = (-data["logprobs"] * mask).sum(axis=-1).mean()
  
      mean_non_score_reward = masked_mean(
          data["non_score_reward"], mask
      )  # non_score_reward is size `batch_size`, `response_length`
      mean_scores = torch.stack(
          data["scores"]
      ).mean()  # scores is size `batch_size`
      std_scores = torch.stack(data["scores"]).std()
  
      if mean_kl.item() < -1.0:
          warnings.warn(
              f"KL divergence is starting to become negative: {mean_kl.item():.2f} -"
              " this might be a precursor for failed training."
              " sometimes this happens because the generation kwargs are not"
              " correctly set. Please make sure that the generation kwargs are set"
              " correctly, or review your training hyperparameters."
          )
  
      stats = {
          "objective/kl": mean_kl,
          "objective/kl_dist": kl_list,
          "objective/logprobs": data["logprobs"],
          "objective/ref_logprobs": data["ref_logprobs"],
          "objective/kl_coef": kl_coef,
          "objective/entropy": mean_entropy,
          "ppo/mean_non_score_reward": mean_non_score_reward,
          "ppo/mean_scores": mean_scores,
          "ppo/std_scores": std_scores,
      }
  
      # Log text properties
      query_lens = torch.tensor(
          [len(query) for query in data["queries"]], dtype=torch.float
      )
      response_lens = torch.tensor(
          [len(response) for response in data["responses"]],
          dtype=torch.float,
      )
  
      stats["tokens/queries_len_mean"] = (
          torch.mean(query_lens).cpu().numpy().item()
      )
      stats["tokens/queries_len_std"] = (
          torch.std(query_lens).cpu().numpy().item()
      )
      stats["tokens/queries_dist"] = query_lens.cpu().numpy()
      stats["tokens/responses_len_mean"] = (
          torch.mean(response_lens).cpu().numpy().item()
      )
      stats["tokens/responses_len_std"] = (
          torch.std(response_lens).cpu().numpy().item()
      )
      stats["tokens/responses_dist"] = response_lens.cpu().numpy()
  
      for k, v in data["train_stats"].items():
          stats[f"ppo/{k}"] = torch.mean(v, dim=0)
      stats["ppo/val/var_explained"] = (
          1 - stats["ppo/val/error"] / stats["ppo/returns/var"]
      )
      return 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 10 for calling the PPOTrainer obj's underlying member method code with example usage: obj.record_step_stats(kl_coef) and returns: stats