# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/02-ppo.ipynb (unless otherwise specified).

__all__ = ['AdaptiveKLController', 'FixedKLController', 'PPOTrainer']

# Cell

import pandas as pd

import numpy as np

import torch.nn.functional as F

from torch.optim import Adam

from transformers import pipeline

from transformers import AutoModelForSequenceClassification, AutoTokenizer

import torch

import time

import random

import os

import math

from tqdm import tqdm

from load_finetuned_model import load_model_and_tokenizer

from modeling_value_head import AutoModelForCausalLMWithValueHead

from core import (logprobs_from_logits,

whiten,

clip_by_value,

entropy_from_logits,

flatten_dict,

average_torch_dicts,

stats_to_np,

stack_dicts,

add_suffix,

build_bert_batch_from_txt)

# Cell

class AdaptiveKLController:

"""

Adaptive KL controller described in the paper:

https://arxiv.org/pdf/1909.08593.pdf

"""

def __init__(self, init_kl_coef, target, horizon):

self.value = init_kl_coef

self.target = target

self.horizon = horizon

def update(self, current, n_steps):

target = self.target

proportional_error = np.clip(current / target - 1, -0.2, 0.2)

mult = 1 + proportional_error * n_steps / self.horizon

self.value *= mult

# Cell

class PPOTrainer:

"""

The PPO_trainer uses Proximal Policy Optimization to tune a language model.

"""

default_params = {

"lr": 1.41e-5,

"adap_kl_ctrl": True,

"init_kl_coef":0.2,

"target": 6,

"horizon":10000,

"gamma":1,

"lam":0.95,

"cliprange": .2,

"cliprange_value":.2,

"vf_coef":.1,

"batch_size": 64,

"forward_batch_size": 8,

"ppo_epochs": 4,

}

def __init__(self, model, ref_model, **ppo_params):

"""

Initialize PPOTrainer.

Args:

model (torch.model): Huggingface GPT2 model

ref_model (torch.model): Huggingface GPT2 refrence model used for KL penalty

ppo_params (dict or None): PPO parameters for training. Can include following keys:

'lr' (float): Adam learning rate, default: 1.41e-5

'batch_size' (int): Number of samples per optimisation step, default: 256

'forward_batch_size' (int): Number of samples forward passed through model at a time, default: 16

'ppo_epochs' (int): Number of optimisation epochs per batch of samples, default: 4

'gamma' (float)): Gamma parameter for advantage calculation, default: 1.

'lam' (float): Lambda parameter for advantage calcualation, default: 0.95

'cliprange_value' (float): Range for clipping values in loss calculation, default: 0.2

'cliprange' (float): Range for clipping in PPO policy gradient loss, default: 0.2

'vf_coef' (float): Scaling factor for value loss, default: 0.1

'adap_kl_ctrl' (bool): Use adaptive KL control, otherwise linear, default: True

'init_kl_coef' (float): Initial KL penalty coefficient (used for adaptive and linear control), default: 0.2

'target' (float): Target KL value for adaptive KL control, default: 6.0

'horizon' (float): Horizon for adaptive KL control, default: 10000

"""

self.ppo_params = self.default_params

self.ppo_params.update(ppo_params)

self.ref_model = ref_model

self.model = model

self.optimizer = Adam(model.parameters(), lr=self.ppo_params['lr'])

# learning rate scheduler

self.epochs = self.ppo_params["epochs"] / self.ppo_params["batch_size"]

self.max_steps = int(self.epochs * self.ppo_params["ppo_epochs"] * self.ppo_params["batch_size"])

self.warmup_steps = int(0.1 * self.max_steps)

self.max_lr = self.ppo_params['lr']

self.min_lr = 0.1 * self.ppo_params['lr']

self.current_step = 0

self.kl_ctl = AdaptiveKLController(self.ppo_params['init_kl_coef'],

self.ppo_params['target'],

self.ppo_params['horizon'])

def get_lr(self, it):

# 1) linear warmup for warmup_iters steps

if it < self.warmup_steps:

return self.max_lr * (it+1) / self.warmup_steps

# 2) if it > lr_decay_iters, return min learning rate

if it > self.max_steps:

return self.min_lr

# 3) in between, use cosine decay down to min learning rate

decay_ratio = (it - self.warmup_steps) / (self.max_steps - self.warmup_steps)

assert 0 <= decay_ratio <= 1

coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio)) # coeff starts at 1 and goes to 0

return self.min_lr + coeff * (self.max_lr - self.min_lr)

def step(self, query, response, scores):

"""

Run a PPO optimisation step.

args:

query (torch.tensor): tensor containing the encoded queries, shape [batch_size, query_length]

response (torch.tensor): tensor containing the encoded responses, shape [batch_size, response_length]

scores (torch.tensor): tensor containing the scores, shape [batch_size]

returns:

train_stats (dict): a summary of the training statistics

"""

bs = self.ppo_params['batch_size']

timing = dict()

t0 = time.time()

gen_len = response.shape[1]

model_input = torch.cat((query, response), axis=1)

t = time.time()

logprobs, ref_logprobs, values = self.batched_forward_pass(model_input, gen_len)

timing['time/ppo/forward_pass'] = time.time()-t

t = time.time()

rewards, non_score_reward, kl_coef = self.compute_rewards(scores, logprobs, ref_logprobs)

timing['time/ppo/compute_rewards'] = time.time()-t

t = time.time()

all_stats = []

idxs = list(range(bs))

# off policy learning

for _ in range(self.ppo_params['ppo_epochs']):

random.shuffle(idxs)

for i in range(bs):

idx = idxs[i]

train_stats = self.train_minibatch(logprobs[idx:idx+1], values[idx:idx+1],

rewards[idx:idx+1], query[idx:idx+1],

response[idx:idx+1], model_input[idx:idx+1], step=self.current_step)

all_stats.append(train_stats)

self.current_step += 1

timing['time/ppo/optimize_step'] = time.time()-t

t = time.time()

train_stats = stack_dicts(all_stats)

# reshape advantages/ratios such that they are not averaged.

train_stats['policy/advantages'] = torch.flatten(train_stats['policy/advantages']).unsqueeze(0)

train_stats['policy/ratio'] = torch.flatten(train_stats['policy/ratio']).unsqueeze(0)

stats = self.record_step_stats(scores=scores, logprobs=logprobs, ref_logprobs=ref_logprobs,

non_score_reward=non_score_reward, train_stats=train_stats,

kl_coef=kl_coef)

stats = stats_to_np(stats)

timing['time/ppo/calc_stats'] = time.time()-t

self.kl_ctl.update(stats['objective/kl'], self.ppo_params['batch_size'])

timing['time/ppo/total'] = time.time()-t0

stats.update(timing)

return stats

def batched_forward_pass(self, model_input, gen_len):

"""Calculate model outputs in multiple batches.

In PPO, we update the policy π_θ. The logprobs output represents log(π_θ(a_t|s_t)), which is crucial for computing the probability ratio:

r_t(θ) = π_θ(a_t|s_t) / π_θ_old(a_t|s_t)

In log space, this becomes:

log(r_t(θ)) = log(π_θ(a_t|s_t)) - log(π_θ_old(a_t|s_t))

Value Function (V_θ):

The values output represents V_θ(s_t), which is used in computing advantages:

A_t = δ_t + (γλ)δ_(t+1) + ... + (γλ)^(T-t+1)δ_(T-1)

where δ_t = r_t + γV(s_(t+1)) - V(s_t)

The ref_logprobs are used to compute the KL divergence between the current policy and the reference policy:

KL(π_ref || π_θ) = E_s[Σ_a π_ref(a|s) * (log π_ref(a|s) - log π_θ(a|s))]

This is used as a soft constraint to prevent the policy from deviating too much from the reference policy.

"""

bs = self.ppo_params['batch_size']

fbs = self.ppo_params['forward_batch_size']

logprobs = []

ref_logprobs = []

values = []

for i in range(int(self.ppo_params['batch_size']/fbs)):

m_input = model_input[i*fbs:(i+1)*fbs]

logits, _, v = self.model(m_input)

ref_logits, _, _ = self.ref_model(m_input)

values.append(v[:, -gen_len-1:-1].detach()) # By selecting -gen_len-1:-1, we ensure that we’re only considering the value predictions for the generated tokens and not for future tokens that haven’t been generated yet.

logprobs.append(logprobs_from_logits(logits[:,:-1,:], m_input[:,1:])[:, -gen_len:].detach()) # we are interested in the log probability of the actual tokens that were generated or observed, not the predictions for tokens that haven’t been generated yet.

ref_logprobs.append(logprobs_from_logits(ref_logits[:,:-1,:], m_input[:,1:])[:, -gen_len:].detach()) # m_input[:, 1:] is used to shift the input tokens to align with the target tokens that the model should predict.

return torch.cat(logprobs), torch.cat(ref_logprobs), torch.cat(values)

def train_minibatch(self, logprobs, values, rewards, query, response, model_input, step):

"""Train one PPO minibatch"""

loss_p, loss_v, train_stats = self.loss(logprobs, values, rewards, query, response, model_input)

loss = loss_p + loss_v

self.optimizer.zero_grad()

loss.backward()

# lr = self.get_lr(step)

# print(f"learning rate {lr} at step {step}")

# for param_group in self.optimizer.param_groups:

# param_group['lr'] = lr

self.optimizer.step()

return train_stats

def compute_rewards(self, scores, logprobs, ref_logprobs):

"""Compute per token rewards from scores and KL-penalty.

Per-token Rewards:

The non_score_reward provides a per-token penalty based on KL divergence.

This is a form of dense reward, giving feedback for each action (token) in the sequence.

"""

kl = logprobs - ref_logprobs

non_score_reward = -self.kl_ctl.value * kl # tune this into a penalty

rewards = non_score_reward.clone().detach()

rewards[:, -1] += scores # add the score to the last token because we compute advatage starting from last token going backwards

return rewards, non_score_reward, self.kl_ctl.value

def loss(self, old_logprobs, values, rewards, query, response, model_input):

"""Calculate policy and value losses."""

lastgaelam = 0

advantages_reversed = []

gen_len = response.shape[1]

for t in reversed(range(gen_len)): # boostrapping

# estimating the advantage term

nextvalues = values[:, t + 1] if t < gen_len - 1 else 0.0

delta = rewards[:, t] + self.ppo_params['gamma'] * nextvalues - values[:, t] # Measure the difference between what we expected and what actually happened.

# total expected return from taking the chosen action

# values[:, t] represents the average expected return from being in the current state, regardless of which action is taken. This is essentially what we'd expect if we chose actions randomly.

# By subtracting values[:, t] from the total expected return, we're calculating how much better (or worse) the chosen action is compared to this average expectation.

lastgaelam = delta + self.ppo_params['gamma'] * self.ppo_params['lam'] * lastgaelam

# It allows advantages from future time steps to influence the advantage estimate of earlier time steps.

# We add to this a discounted (gamma) version of our previous estimate (lastgaelam).

# The 'lam' (lambda) parameter lets us control how much we trust future estimates vs. current feedback.

advantages_reversed.append(lastgaelam)

advantages = torch.stack(advantages_reversed[::-1]).transpose(0, 1)

returns = advantages + values

advantages = whiten(advantages)

advantages = advantages.detach()

logits, _, vpred = self.model(model_input)

logprob = logprobs_from_logits(logits[:,:-1,:], model_input[:, 1:])

#only the generation part of the values/logprobs is needed

logprob, vpred = logprob[:, -gen_len:], vpred[:,-gen_len-1:-1]

vpredclipped = clip_by_value(vpred,

values - self.ppo_params["cliprange_value"],

values + self.ppo_params["cliprange_value"])

vf_losses1 = (vpred - returns)**2

vf_losses2 = (vpredclipped - returns)**2

vf_loss = .5 * torch.mean(torch.max(vf_losses1, vf_losses2))

vf_clipfrac = torch.mean(torch.gt(vf_losses2, vf_losses1).double())

ratio = torch.exp(logprob - old_logprobs)

pg_losses = -advantages * ratio

pg_losses2 = -advantages * torch.clamp(ratio,

1.0 - self.ppo_params['cliprange'],

1.0 + self.ppo_params['cliprange'])

pg_loss = torch.mean(torch.max(pg_losses, pg_losses2))

pg_clipfrac = torch.mean(torch.gt(pg_losses2, pg_losses).double())

loss = pg_loss + self.ppo_params['vf_coef'] * vf_loss

entropy = torch.mean(entropy_from_logits(logits))

approxkl = .5 * torch.mean((logprob - old_logprobs)**2)

policykl = torch.mean(logprob - old_logprobs)

return_mean, return_var = torch.mean(returns), torch.var(returns)

value_mean, value_var = torch.mean(values), torch.var(values)

stats = dict(

loss=dict(policy=pg_loss, value=vf_loss, total=loss),

policy=dict(entropy=entropy, approxkl=approxkl,policykl=policykl, clipfrac=pg_clipfrac,

advantages=advantages, advantages_mean=torch.mean(advantages), ratio=ratio),

returns=dict(mean=return_mean, var=return_var),

val=dict(vpred=torch.mean(vpred), error=torch.mean((vpred - returns) ** 2),

clipfrac=vf_clipfrac, mean=value_mean, var=value_var),

logp=dict(logp=logprob, logp_old=old_logprobs, logp_diff=logprob-old_logprobs)

)

print(f"pg_loss {pg_loss} | vf_loss {self.ppo_params['vf_coef'] * vf_loss} | returns: {return_mean} | value {value_mean} | policykl: {policykl} | advantages: {torch.mean(advantages)}")

return pg_loss, self.ppo_params['vf_coef'] * vf_loss, flatten_dict(stats)

def record_step_stats(self, kl_coef, **data):

"""Record training step statistics."""

kl = data['logprobs'] - data['ref_logprobs']

mean_kl = torch.mean(torch.sum(kl, axis=-1))

mean_entropy = torch.mean(torch.sum(-data['logprobs'], axis=1))

mean_non_score_reward =torch.mean(torch.sum(data['non_score_reward'], axis=1))

stats = {

'objective/kl': mean_kl,

'objective/kl_dist': kl,

'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,

}

for k, v in data['train_stats'].items():

stats[f'ppo/{k}'] = torch.mean(v, axis=0)

stats['ppo/val/var_explained'] = 1 - stats['ppo/val/error'] / stats['ppo/returns/var']

return stats

config = {

"tb_path": "runs",

"hidden_size": 768,

"cls_model_name": "lvwerra/bert-imdb",

"tk_name": "gpt2",

"txt_in_len": 5,

"txt_out_len": 15,

"lr": 1.41e-5 * 0.7,

"init_kl_coef":0.2,

"target": 6,

"horizon":1000,

"gamma":1,

"lam":0.95,

"cliprange": .2,

"cliprange_value":.2,

"vf_coef":.1,

"epochs": 25600,

"batch_size": 64,

"forward_batch_size": 4,

"ppo_epochs": 4,

}

df = pd.read_csv('data/imdb-dataset.csv')

df = df.loc[df['review'].str.len() > 500]

df['review'] = df['review'].apply(lambda x: x[:1000])

model_path = "checkpoints/checkpoint_57.pt"

config_path = "checkpoints/config.json"

tokenizer, model, device, _ = load_model_and_tokenizer(model_path, config_path)

_, ref_model, _, model_config = load_model_and_tokenizer(model_path, config_path)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

df['tokens'] = df['review'].apply(lambda x: tokenizer.encode(x, return_tensors="pt").to(device)[0, :config['txt_in_len']])

df['query'] = df['tokens'].apply(lambda x: tokenizer.decode(x))

model = AutoModelForCausalLMWithValueHead(model, model_config, device).to(device)

ref_model = AutoModelForCausalLMWithValueHead(ref_model, model_config, device).to(device)

ppo_trainer = PPOTrainer(model, ref_model, device=device, **config)

def respond_to_batch(model, input_ids, text_length, pad_token_id, device="cuda"):

input_ids = input_ids.to(device)

batch_size = input_ids.size(0)

input_length = input_ids.shape[1]

while input_ids.size(1) < text_length:

with torch.no_grad():

logits, loss, last_hidden_state = model(input_ids)

logits = logits[:, -1, :]

probs = F.softmax(logits, dim=-1)

topk_probs, topk_indices = torch.topk(probs, 50, dim=-1)

ix = torch.multinomial(topk_probs, 1)

next_token = torch.gather(topk_indices, -1, ix)

input_ids = torch.cat((input_ids, next_token), dim=1)

if input_ids.size(1) < text_length:

padding = torch.full((batch_size, text_length - input_ids.size(1)),

pad_token_id,

dtype=torch.long,

device=device)

input_ids = torch.cat((input_ids, padding), dim=1)

generated_sequence = input_ids[:, input_length:]

return generated_sequence

#TODO -> train distil bert on imdb dataset for sentiment analysis

sentiment_model = AutoModelForSequenceClassification.from_pretrained(config["cls_model_name"])

sentiment_tokenizer = AutoTokenizer.from_pretrained(config["cls_model_name"])

text = 'this movie was really good!!'

output = sentiment_model.forward(sentiment_tokenizer.encode(text, return_tensors="pt"))

print("predicted sentiment :", output)

fbs = config['forward_batch_size']

model.to(device)

sentiment_model.to(device)

ref_model.to(device)

all_stats = []

for epoch in tqdm(range(int(np.ceil(config["epochs"]/config['batch_size'])))):

torch.cuda.empty_cache()

logs = dict()

game_data = dict()

timing = dict()

t0 = time.time()

#### get a batch from the dataset

df_batch = df.sample(config['batch_size'])

game_data['query'] = df_batch['query'].tolist()

query_tensors = torch.stack(df_batch['tokens'].tolist())

#### get response from gpt2

t = time.time()

total_length = config['txt_in_len']+config['txt_out_len']

response_tensors = []

for i in range(int(config['batch_size']/fbs)):

response = respond_to_batch(model, query_tensors[i*fbs:(i+1)*fbs],

text_length=total_length, pad_token_id=tokenizer.eos_token_id, device=device)

response_tensors.append(response)

response_tensors = torch.cat(response_tensors)

game_data['response'] = [tokenizer.decode(response_tensors[i, :]) for i in range(config['batch_size'])]

timing['time/get_response'] = time.time()-t

#### tokenize text for sentiment analysis

t = time.time()

texts = [q + r for q,r in zip(game_data['query'], game_data['response'])]

sentiment_inputs, attention_masks = build_bert_batch_from_txt(texts, sentiment_tokenizer, device)

timing['time/build_input_sentiment'] = time.time()-t

#### get sentiment score

t = time.time()

rewards = []

for i in range(int(config['batch_size']/fbs)):

res = sentiment_model.forward(sentiment_inputs[i*fbs:(i+1)*fbs],

attention_masks[i*fbs:(i+1)*fbs])[0][:, 1].detach()

rewards.append(res)

rewards = torch.cat(rewards)

timing['time/get_sentiment_preds'] = time.time()-t

#### Run PPO training

t = time.time()

stats = ppo_trainer.step(query_tensors, response_tensors, rewards)

timing['time/optimization'] = time.time()-t

#### Log everything

logs.update(timing)

logs.update(stats)

logs['env/reward_mean'] = torch.mean(rewards).cpu().numpy()

logs['env/reward_std'] = torch.std(rewards).cpu().numpy()

logs['env/reward_dist'] = rewards.cpu().numpy()

all_stats.append(logs)

checkpoint_dir = "ppo_checkpoints"

model.save_pretrained(os.path.join(checkpoint_dir, f"policy_model.pt"), config=config, safe_serialization=True)

ref_model.save_pretrained(os.path.join(checkpoint_dir, f"ref_model.pt"), config=config, safe_serialization=True)

stats_df = pd.DataFrame(all_stats)

stats_df.to_csv("ppo_training_stats.csv", index=False)

#### get a batch from the dataset

bs = 16

game_data = dict()

df_batch = df.sample(bs)

game_data['query'] = df_batch['query'].tolist()

query_tensors = torch.stack(df_batch['tokens'].tolist())

#### get response from gpt2 and gpt2_ref

total_length = config['txt_in_len']+config['txt_out_len']

response_tensors_ref = respond_to_batch(ref_model, query_tensors, text_length=config['txt_out_len'], pad_token_id=tokenizer.eos_token_id, device=device)

game_data['response (before)'] = [tokenizer.decode(response_tensors_ref[i, :]) for i in range(bs)]

response_tensors = respond_to_batch(model, query_tensors, text_length=config['txt_out_len'], pad_token_id=tokenizer.eos_token_id, device=device)

game_data['response (after)'] = [tokenizer.decode(response_tensors[i, :]) for i in range(bs)]

#### sentiment analysis of query/response pairs before/after

texts = [q + r for q,r in zip(game_data['query'], game_data['response (before)'])]

sentiment_inputs, attention_masks = build_bert_batch_from_txt(texts, sentiment_tokenizer, device)

rewards = sentiment_model.forward(sentiment_inputs, attention_masks)[0][:, 1].detach()

game_data['rewards (before)'] = rewards.cpu().numpy()

texts = [q + r for q,r in zip(game_data['query'], game_data['response (after)'])]

sentiment_inputs, attention_masks = build_bert_batch_from_txt(texts, sentiment_tokenizer, device)

rewards = sentiment_model.forward(sentiment_inputs, attention_masks)[0][:, 1].detach()

game_data['rewards (after)'] = rewards.cpu().numpy()

# store results in a dataframe

df_results = pd.DataFrame(game_data)

df_results.to_csv("ref_model_response.csv", index=False)