train.py

import argparse
from os.path import join
import math
import wandb
from tqdm.auto import tqdm, trange
import numpy as np
import torch
from torch.utils.data import DataLoader
import datasets
from transformers import AdamW, get_scheduler, GPT2Tokenizer, GPT2LMHeadModel, DataCollatorForTokenClassification
from accelerate import Accelerator
from detoxify import Detoxify

from data import load_rtp_dataset, prep_dataset_for_training, prep_dataset_for_generation, rand_split_dataset, \
                 get_ngram_overlap


parser = argparse.ArgumentParser(description="Data Feedback on Conditional Language Generation")
parser.add_argument('--wandb-log', action='store_true', help='if set, logs to wandb')
parser.add_argument('--wandb-group', type=str, default='debugging', help='group name for this experiment')
parser.add_argument('--data-dir', default='./data', type=str, help='path to data root')
parser.add_argument('--test-set-size', type=int, default=14442)
parser.add_argument('-n', '--init-train-set-size', type=int, default=20000)
parser.add_argument('-m', '--num-human-labeled-samples', type=int, default=2500)
parser.add_argument('-k', '--num-model-labeled-samples', type=int, default=2500)
parser.add_argument('-a', '--arch', choices=['gpt2', 'gpt2-medium', 'gpt2-large'], default='gpt2')
parser.add_argument("--max-generated-tokens", type=int, default=20, help="Maximum number of new tokens to generate for each sample.")
parser.add_argument("--memorization-ngrams", nargs="+", type=int, default=[5, 10], help="Which ngrams to compute memorization over.")
parser.add_argument("--sampling-type", choices=["beam_search", "nucleus_sampling"], default="nucleus_sampling")
parser.add_argument("--nucleus-top-p", type=float, default=0.9, help="Top_p parameter for nucleus sampling during generation.")
parser.add_argument("--num-beams", type=int, default=10, help="Number of beams for beam search during generation.")
parser.add_argument("--num-train-epochs", default=1, type=int)
parser.add_argument("--tox-model-batch-size", type=int, default=64, help="Batch size for toxicity model.")
parser.add_argument("--train-batch-size", type=int, help='automatically set based on arch if not provided')
parser.add_argument("--gradient-accumulation-steps", type=int, help='automatically set based on arch if not provided')
parser.add_argument("--learning-rate", type=float, help='automatically set based on arch if not provided')
args = parser.parse_args()

# Set training hparams for the model (if not given)
# by default, train_batch_size and gradient_accumulation_steps set so total_batch__size = 16 always
if not args.train_batch_size:
    if args.arch == 'gpt2': args.train_batch_size = 16
    elif args.arch == 'gpt2-medium': args.train_batch_size = 8
    elif args.arch == 'gpt2-large': args.train_batch_size = 4
if not args.gradient_accumulation_steps:
    if args.arch == 'gpt2': args.gradient_accumulation_steps = 1
    elif args.arch == 'gpt2-medium': args.gradient_accumulation_steps = 2
    elif args.arch == 'gpt2-large': args.gradient_accumulation_steps = 4
if not args.learning_rate:
    if args.arch == 'gpt2': args.learning_rate = 5e-5
    elif args.arch == 'gpt2-medium': args.learning_rate = 1e-5
    elif args.arch == 'gpt2-large': args.learning_rate = 1e-5
args.total_batch_size = args.train_batch_size * args.gradient_accumulation_steps

# Set evaluation batch sizes
args.eval_batch_size = args.train_batch_size * 5
if args.sampling_type == 'beam_search':
    args.gen_batch_size = args.train_batch_size * 25 // args.num_beams
elif args.sampling_type == 'nucleus_sampling':
    args.gen_batch_size = args.train_batch_size * 25


@torch.no_grad()
def evaluate_ppl(model, data_loader, accelerator):
    model.eval()
    losses = []
    for batch in tqdm(data_loader, disable=not accelerator.is_main_process, desc='Evaluating'):
        outputs = model(**batch)
        loss = outputs.loss
        losses.append(accelerator.gather(loss.repeat(data_loader.batch_sampler.batch_size)))
    losses = torch.cat(losses)[:len(data_loader.dataset)]
    perplexity = math.exp(torch.mean(losses))
    return perplexity


@torch.no_grad()
def generate_samples(model, data_loader, accelerator, tokenizer, max_generated_tokens, sampling_type, nucleus_top_p, num_beams):
    tokenizer.padding_side = 'left'
    model.eval()
    sample_ids, samples = [], []
    
    for batch in tqdm(data_loader, disable=not accelerator.is_main_process, desc='Generating'):
        if sampling_type == 'nucleus_sampling':
            outputs = model.generate(batch['input_ids'], attention_mask=batch['attention_mask'], max_new_tokens=max_generated_tokens, 
                                     do_sample=True, top_p=nucleus_top_p, top_k=0)
        elif sampling_type == 'beam_search':
            outputs = model.generate(batch['input_ids'], attention_mask=batch['attention_mask'], max_new_tokens=max_generated_tokens, 
                                     do_sample=False, num_beams=num_beams)
        # grab only the continuation generated by the model
        generations = outputs[:, batch['input_ids'].size(1):]
        generations = accelerator.gather(generations)

        # convert generations to text, and also store generated IDs but remove eos tokens
        samples.extend(tokenizer.batch_decode(generations, skip_special_tokens=True))
        for generation in generations:
            sample_ids.append(generation[generation != tokenizer.eos_token_id].tolist())

    samples = samples[:len(data_loader.dataset)]
    sample_ids = sample_ids[:len(data_loader.dataset)]
    tokenizer.padding_side = 'right'
    return samples, sample_ids


@torch.no_grad()
def evaluate_samples_tox(samples, accelerator, tox_model, batch_size, stats_prefix):
    tox_scores, quote_counts = [], []
    for i in trange(0, len(samples), batch_size, disable=not accelerator.is_main_process, desc='Tox-scoring'):
        # track quotes
        quote_counts.extend([sample.count('”') for sample in samples[i:i+batch_size]])
        # score toxicity using pretrained model
        pred_tox_scores = tox_model.predict(samples[i:i+batch_size])['toxicity']
        pred_tox_scores = accelerator.gather(torch.tensor(pred_tox_scores))
        tox_scores.append(pred_tox_scores)

    tox_score_results = {}
    tox_scores = torch.cat(tox_scores)
    tox_score_results[stats_prefix + '/avg_toxicity'] = tox_scores.mean()  
    tox_score_results[stats_prefix + '/avg_threshold_toxicity'] = (tox_scores > 0.5).float().mean()     
    tox_score_results[stats_prefix + '/avg_num_quotes'] = np.mean(quote_counts)
    tox_score_results[stats_prefix + '/avg_has_a_quote'] = (np.array(quote_counts) > 0).mean()
    return tox_score_results


def train(model, train_loader, optimizer, lr_scheduler, num_train_epochs, gradient_accumulation_steps, num_train_steps, accelerator):
    model.train()
    progress_bar = tqdm(range(num_train_steps), disable=not accelerator.is_main_process, desc='Training')

    for _ in range(num_train_epochs):
        for step, batch in enumerate(train_loader):
            outputs = model(**batch)
            loss = outputs.loss
            loss = loss / gradient_accumulation_steps
            accelerator.backward(loss)

            if step % gradient_accumulation_steps == 0 or step == len(train_loader) - 1:
                wandb.log({'train_loss': loss.item(), 'lr': lr_scheduler.get_last_lr()[0]})
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()
                progress_bar.update(1)


if __name__ == "__main__":
    wandb.init(project='data-feedback', group=args.wandb_group, mode='online' if args.wandb_log else 'disabled')
    wandb.config.update(vars(args))

    accelerator = Accelerator()

    # Load GPT2 tokenizer
    tokenizer = GPT2Tokenizer.from_pretrained(args.arch)
    tokenizer.pad_token = tokenizer.eos_token

    # Handles dynamic padding of examples
    data_collator = DataCollatorForTokenClassification(tokenizer=tokenizer)
    
    # Get language model
    model = GPT2LMHeadModel.from_pretrained(args.arch, pad_token_id=tokenizer.eos_token_id)

    # Load and split dataset into train, test, and reserve sets
    dataset = load_rtp_dataset(join(args.data_dir, 'realtoxicityprompts-data/prompts.jsonl'), tokenizer)
    test_set, reserve_set = rand_split_dataset(dataset, args.test_set_size)
    train_set, reserve_set = rand_split_dataset(reserve_set, args.init_train_set_size)

    # Set up test set for evaluation and generation
    test_set_gen = prep_dataset_for_generation(test_set)
    test_set_eval = prep_dataset_for_training(test_set)
    
    test_gen_loader = DataLoader(test_set_gen, shuffle=False, collate_fn=data_collator, batch_size=args.gen_batch_size)
    test_eval_loader = DataLoader(test_set_eval, shuffle=False, collate_fn=data_collator, batch_size=args.eval_batch_size)

    # Get toxicity model
    tox_model = Detoxify('original', device=accelerator.device)

    # Move objects to GPU
    model, test_gen_loader, test_eval_loader = accelerator.prepare(model, test_gen_loader, test_eval_loader)

    # Do initial pretrained model evaluation.
    round = -1
    print(f'*** ROUND {round} (pretrained model evaluation) ***')

    # Evaluate initial (pretrained) model ppl and toxicity
    test_ppl = evaluate_ppl(model, test_eval_loader, accelerator)
    test_set_model_samples, _ = generate_samples(model, test_gen_loader, accelerator, tokenizer, args.max_generated_tokens,
                                                    args.sampling_type, args.nucleus_top_p, args.num_beams)
    model_tox_scores = evaluate_samples_tox(test_set_model_samples, accelerator, tox_model, args.tox_model_batch_size, 'generated')

    # Evaluate test set toxicity
    test_set_tox_scores = evaluate_samples_tox(test_set['continuation'], accelerator, tox_model, args.tox_model_batch_size, 'test_set')

    # Log initial statistics
    stats = {'round': round, 'test_ppl': test_ppl} | model_tox_scores | test_set_tox_scores
    wandb.log(stats)


    while True:
        round += 1
        print(f'*** ROUND {round} ***')

        # Create model, optimizer, and dataloader. Move to GPU with accelerator
        model = GPT2LMHeadModel.from_pretrained(args.arch, pad_token_id=tokenizer.eos_token_id)
        optimizer = AdamW(model.parameters(), lr=args.learning_rate, weight_decay=0)
        train_set_prepped = prep_dataset_for_training(train_set)
        train_loader = DataLoader(train_set_prepped, shuffle=True, collate_fn=data_collator, batch_size=args.train_batch_size)
        model, optimizer, train_loader = accelerator.prepare(model, optimizer, train_loader)
        
        # Scheduler and math around the number of training steps.
        num_update_steps_per_epoch = math.ceil(len(train_loader) / args.gradient_accumulation_steps)
        num_train_steps = args.num_train_epochs * num_update_steps_per_epoch
        lr_scheduler = get_scheduler(name='linear', optimizer=optimizer, num_training_steps=num_train_steps, num_warmup_steps=0)

        # Free models and optimizers from previous loop
        accelerator.free_memory()
        
        # Train model
        train(model, train_loader, optimizer, lr_scheduler, args.num_train_epochs, 
              args.gradient_accumulation_steps, num_train_steps, accelerator)

        # Free training objects from memory
        del lr_scheduler
        del optimizer
        accelerator.free_memory()

        # Evaluate model ppl and toxicity
        test_ppl = evaluate_ppl(model, test_eval_loader, accelerator)
        test_set_model_samples, test_set_model_sample_ids = generate_samples(model, test_gen_loader, accelerator, tokenizer, 
                        args.max_generated_tokens, args.sampling_type, args.nucleus_top_p, args.num_beams)
        model_tox_scores = evaluate_samples_tox(test_set_model_samples, accelerator, tox_model, args.tox_model_batch_size, 'generated')

        # Calculate bias for the models tox scores
        model_bias_scores = {}
        for tox_name, tox_score in model_tox_scores.items():
            tox_name = tox_name.split('/')[1]
            model_bias_scores[f'bias/{tox_name}'] = tox_score - wandb.run.summary[f'test_set/{tox_name}']
        model_tox_scores |= model_bias_scores

        # Evaluate model memorization
        for n in args.memorization_ngrams:
            model_tox_scores |= {f'memorization/{n}_gram_overlap': get_ngram_overlap(train_set['continuation_ids'], 
                                                                                     test_set_model_sample_ids, n)}

        # Evaluate train set toxicity
        train_set_tox_scores = evaluate_samples_tox(train_set['continuation'], accelerator, tox_model, args.tox_model_batch_size, 'train_set')

        # Log stats from the round
        stats |= model_tox_scores | train_set_tox_scores | {'round': round, 'test_ppl': test_ppl, 'size_train_set': len(train_set), 
                                                            'size_reserve_set': len(reserve_set)}
        wandb.log(stats)

        if len(reserve_set) < args.num_human_labeled_samples + args.num_model_labeled_samples:
            print(f'Ending retraining - not enough remaining samples in reserve set ({len(reserve_set)} left).')
            break
    
        if args.num_human_labeled_samples > 0:
            # Take some samples from reserve_set, preprocess them for training, and add them to train set
            reserve_set_selected, reserve_set = rand_split_dataset(reserve_set, args.num_human_labeled_samples)
            train_set = datasets.concatenate_datasets([train_set, reserve_set_selected])

        if args.num_model_labeled_samples > 0:
            # Take some samples from reserve_set set and preprocess them for generation (model will complete the prompts)
            reserve_set_selected, reserve_set = rand_split_dataset(reserve_set, args.num_model_labeled_samples)
            reserve_set_gen = prep_dataset_for_generation(reserve_set_selected)

            # Generate completions with model
            reserve_set_loader = DataLoader(reserve_set_gen, shuffle=False, collate_fn=data_collator, batch_size=args.gen_batch_size)
            reserve_set_loader = accelerator.prepare(reserve_set_loader)
            reserve_samples, reserve_sample_ids = generate_samples(model, reserve_set_loader, accelerator, tokenizer, args.max_generated_tokens,
                                                                   args.sampling_type, args.nucleus_top_p, args.num_beams)

            # Use model generations as ground truth & add to train set
            reserve_set_remapped = reserve_set_selected.remove_columns(['continuation', 'continuation_ids'])
            reserve_set_remapped = reserve_set_remapped.add_column('continuation', reserve_samples)
            reserve_set_remapped = reserve_set_remapped.add_column('continuation_ids', reserve_sample_ids)
            train_set = datasets.concatenate_datasets([train_set, reserve_set_remapped])