From c599c4364cc17ce13d513c0b14e4e89734323542 Mon Sep 17 00:00:00 2001 From: Nathan Lambert Date: Wed, 7 Feb 2024 03:30:39 +0000 Subject: [PATCH] works --- analysis/README.md | 19 +++ analysis/per_token_reward.py | 237 +++++++++++++++++++++++++++++++++++ scripts/run_rm.py | 2 + 3 files changed, 258 insertions(+) create mode 100644 analysis/per_token_reward.py diff --git a/analysis/README.md b/analysis/README.md index 96d00ae8..5388ea8e 100644 --- a/analysis/README.md +++ b/analysis/README.md @@ -3,6 +3,25 @@ We're going to add visualizations for both the eval. data and results here. So far, we have the following tools: +### Per token uterrance reward +This returns the reward per-token to show how the reward evolves over a piece of text. +``` +python analysis/per_token_reward.py --model=OpenAssistant/reward-model-deberta-v3-large-v2 --chat_template=raw --text="I love to walk the dog, what do you like?" +``` +E.g. with OpenAssistant/reward-model-deberta-v3-large-v2 +Reward: -0.544 | Substring: I +Reward: -0.556 | Substring: I love +Reward: -0.566 | Substring: I love to +Reward: 0.099 | Substring: I love to walk +Reward: 0.096 | Substring: I love to walk the +Reward: 0.092 | Substring: I love to walk the dog +Reward: 0.09 | Substring: I love to walk the dog, +Reward: 0.087 | Substring: I love to walk the dog, what +Reward: 0.085 | Substring: I love to walk the dog, what do +Reward: 0.089 | Substring: I love to walk the dog, what do you +Reward: 0.09 | Substring: I love to walk the dog, what do you like +Reward: 0.093 | Substring: I love to walk the dog, what do you like? + ### Model usage within eval. dataset To run this, execute: ``` diff --git a/analysis/per_token_reward.py b/analysis/per_token_reward.py new file mode 100644 index 00000000..209ef343 --- /dev/null +++ b/analysis/per_token_reward.py @@ -0,0 +1,237 @@ +# Copyright 2023 AllenAI. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Script to output the per-token reward across a piece of text given a reward model + +import argparse +import logging +import sys + +import torch +import transformers +from accelerate import Accelerator +from accelerate.logging import get_logger +from datasets import Dataset +from tqdm import tqdm +from transformers import ( + AutoModelForSequenceClassification, + AutoTokenizer, + T5ForConditionalGeneration, + pipeline, +) + + +def get_args(): + """ + Parse arguments strings model and chat_template + """ + parser = argparse.ArgumentParser() + parser.add_argument("--model", type=str, default="natolambert/gpt2-dummy-rm", help="path to model") + parser.add_argument( + "--tokenizer", type=str, default=None, help="path to non-matching tokenizer, requires --direct_load" + ) + parser.add_argument("--chat_template", type=str, default="tulu", help="path to chat template") + parser.add_argument( + "--batch_size", type=int, default=64, help="batch size for inference (if above number of tokens)" + ) + parser.add_argument("--text", type=str, default="I love to drink coffee at work.", help="text to evaluate") + args = parser.parse_args() + + if "PairRM" in args.model or "PairRM" in args.chat_template or "SHP" in args.model or "SHP" in args.chat_template: + # Note: SHP could be used in single-output mode, but the code is not yet added + raise ValueError("PairRM and SHP require pairwise inputs, not supported") + return args + + +def main(): + args = get_args() + quantized = True # only Starling isn't quantized for now + custom_dialogue = False + # some models need custom code to be run + if "oasst" in args.model or "oasst" in args.chat_template: + from herm.models import openassistant # noqa + + model_builder = AutoModelForSequenceClassification.from_pretrained + pipeline_builder = pipeline + elif "Starling" in args.model or "Starling" in args.chat_template: + from herm.models.starling import StarlingPipeline, build_starling_rm + + model_builder = build_starling_rm + pipeline_builder = StarlingPipeline + quantized = False + elif "openbmb" in args.model or "openbmb" in args.chat_template: + from herm.models.openbmb import LlamaRewardModel, OpenBMBPipeline + + model_builder = LlamaRewardModel.from_pretrained + pipeline_builder = OpenBMBPipeline + elif "PairRM" in args.model or "PairRM" in args.chat_template: + from herm.models.pairrm import DebertaV2PairRM, PairRMPipeline + + custom_dialogue = True + model_builder = DebertaV2PairRM.from_pretrained + pipeline_builder = PairRMPipeline + elif "SHP" in args.model or "SHP" in args.chat_template: + from herm.models.shp import SHPPipeline + + custom_dialogue = True + model_builder = T5ForConditionalGeneration.from_pretrained + pipeline_builder = SHPPipeline + else: + model_builder = AutoModelForSequenceClassification.from_pretrained + pipeline_builder = pipeline + + if custom_dialogue: + raise ValueError("Custom dialogue formatting not yet supported in this script") + + ############### + # Setup logging + ############### + accelerator = Accelerator() + current_device = accelerator.process_index + + logger = get_logger(__name__) + logging.basicConfig( + format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + handlers=[logging.StreamHandler(sys.stdout)], + ) + log_level = logging.INFO + logger.setLevel(log_level) + transformers.utils.logging.set_verbosity(log_level) + transformers.utils.logging.enable_default_handler() + transformers.utils.logging.enable_explicit_format() + + logger.info(f"Running reward model on {args.model} with chat template {args.chat_template}") + + ############################ + # Load reward model pipeline + ############################ + tokenizer_path = args.tokenizer if args.tokenizer else args.model + tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) + + BATCH_SIZE = args.batch_size + logger.info("*** Load reward model ***") + reward_pipeline_kwargs = { + "batch_size": BATCH_SIZE, # eval_args.inference_batch_size, + "truncation": True, + "padding": True, + "max_length": 2048, + "function_to_apply": "none", # Compute raw logits + "return_token_type_ids": False, + } + if quantized: + model_kwargs = { + "load_in_8bit": True, + "device_map": {"": current_device}, + "torch_dtype": torch.float16 if torch.cuda.is_available() else None, + } + else: + model_kwargs = {"device_map": {"": current_device}} + # TODO remove direct load logic + # if pipeline_builder is pipeline, use built in pipeline, else custom + if not pipeline == pipeline_builder: + model = model_builder(args.model, **model_kwargs) + tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) + reward_pipe = pipeline_builder( + "text-classification", + model=model, + tokenizer=tokenizer, + ) + else: + reward_pipe = pipeline( + "text-classification", + model=args.model, + tokenizer=tokenizer, + revision="main", + model_kwargs=model_kwargs, + ) + + ############################ + # Tokenization settings & dataset preparation + ############################ + # set pad token to eos token if not set + if reward_pipe.tokenizer.pad_token_id is None: + reward_pipe.model.config.pad_token_id = reward_pipe.tokenizer.eos_token_id + reward_pipe.tokenizer.pad_token_id = reward_pipe.tokenizer.eos_token_id + + def tokenify_string(string, tokenizer): + # Tokenize the entire text + tokens = tokenizer.tokenize(string) + + cumulative_texts = [] + # Iterate over each token + for i, _ in enumerate(tokens): + # Append the current cumulative text to the list + cumulative_texts.append(tokenizer.convert_tokens_to_string(tokens[: i + 1])) + + return cumulative_texts + + substrings = tokenify_string(args.text, tokenizer) + # create dataset from list of strings substrings with huggingface + dataset = [{"text": substring} for substring in substrings] + dataset = Dataset.from_list(dataset) + + ############################ + # Run inference [1/2]" built in transformers + ############################ + # if using HF pipeline, can pass entire dataset and get results + # first, handle custom pipelines that we must batch normally + if not pipeline_builder == pipeline: + logger.info("*** Running forward pass via built in pipeline abstraction ***") + # this setup can be optimized slightly with one pipeline call + # prepare for inference + reward_pipe = accelerator.prepare(reward_pipe) + + rewards = reward_pipe(dataset["text"], **reward_pipeline_kwargs) + + ############################ + # Run inference [2/2] custom pipelines + ############################ + else: + logger.info("*** Running dataloader to collect results ***") + + dataloader = torch.utils.data.DataLoader( + dataset, + batch_size=BATCH_SIZE, + collate_fn=None, + shuffle=False, + drop_last=False, + ) + + dataloader, model = accelerator.prepare(dataloader, reward_pipe.model) + reward_pipe.model = model + + results = [] + for step, batch in enumerate(tqdm(dataloader, desc="RM batch steps")): + logger.info(f"RM inference step {step}/{len(dataloader)}") + rewards = reward_pipe(batch["text"], **reward_pipeline_kwargs) + + # for each item in batch, record 1 if chosen > rejected + # extra score from dict within batched results (e.g. logits) + # [{'label': 'LABEL_1', 'score': 0.6826171875},... ] + if isinstance(rewards[0], dict): + scores = [result["score"] for result in rewards] + # for classes that directly output scores (custom code) + else: + scores = rewards.cpu().numpy().tolist() + + results.extend(scores) + + # print the results + for i, substring in enumerate(substrings): + print(f"Reward: {round(results[i], 3)} | Substring: {substring}") + + +if __name__ == "__main__": + main() diff --git a/scripts/run_rm.py b/scripts/run_rm.py index ade9b06c..3e3a87eb 100644 --- a/scripts/run_rm.py +++ b/scripts/run_rm.py @@ -161,6 +161,8 @@ def main(): } else: model_kwargs = {"device_map": {"": current_device}} + # TODO remove direct load logic + # if pipeline_builder is pipeline, use built in pipeline, else custom if args.direct_load: model = model_builder(args.model, **model_kwargs) tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)