mirror of
https://github.com/wassname/SimPO.git
synced 2026-07-13 17:41:28 +08:00
Yu Meng
87 lines
3.2 KiB
Python
87 lines
3.2 KiB
Python
|
|
import torch
|
|
from transformers import AutoModelForSequenceClassification, AutoTokenizer
|
|
import json
|
|
import os
|
|
import argparse
|
|
import tqdm
|
|
import numpy as np
|
|
import datasets
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--generation_file", type=str, default="datasets/gemma2_ultrafeedback/all_outputs.json", help="Path to the output generation file")
|
|
parser.add_argument("--reward_model", type=str, default="RLHFlow/ArmoRM-Llama3-8B-v0.1", help="Path to reward model")
|
|
parser.add_argument("--output_dir", type=str, default="datasets/gemma2_ultrafeedback/", help="Path to output directory")
|
|
args = parser.parse_args()
|
|
|
|
print(args)
|
|
|
|
generation_file = args.generation_file
|
|
with open(generation_file, 'r') as f:
|
|
output_data = json.load(f)
|
|
|
|
inputs = [data["prompt"] for data in output_data]
|
|
candidates_texts = [data["all_generated_responses"] for data in output_data]
|
|
|
|
model = AutoModelForSequenceClassification.from_pretrained(args.reward_model,
|
|
device_map="cuda",
|
|
trust_remote_code=True, torch_dtype=torch.bfloat16)
|
|
tokenizer = AutoTokenizer.from_pretrained(args.reward_model, use_fast=True)
|
|
|
|
for data in tqdm.tqdm(output_data):
|
|
prompt = data["prompt"]
|
|
candidates = data["all_generated_responses"]
|
|
scores = []
|
|
for candidate in candidates:
|
|
messages = [{"role": "user", "content": prompt},
|
|
{"role": "assistant", "content": candidate}]
|
|
input_ids = tokenizer.apply_chat_template(messages, return_tensors="pt").to("cuda")
|
|
with torch.no_grad():
|
|
output = model(input_ids)
|
|
score = output.score.float().item()
|
|
scores.append(score)
|
|
data["all_rm_scores"] = scores
|
|
|
|
file_name = os.path.basename(args.generation_file).split('.json')[0] + "_rm.json"
|
|
with open(os.path.join(args.output_dir, file_name), 'w') as f:
|
|
json.dump(output_data, f, indent=4)
|
|
|
|
print(f"Annotated outputs saved to {os.path.join(args.output_dir, file_name)}")
|
|
|
|
# Binarize data: win = highest scoring reponse; lose = lowest scoring response
|
|
for data in output_data:
|
|
chosen_idx = np.argmax(data["all_rm_scores"])
|
|
rejected_idx = np.argmin(data["all_rm_scores"])
|
|
chosen = []
|
|
chosen.append({
|
|
"role": "user",
|
|
"content": data["prompt"]
|
|
})
|
|
chosen.append({
|
|
"role": "assistant",
|
|
"content": data["all_generated_responses"][chosen_idx]
|
|
})
|
|
rejected = []
|
|
rejected.append({
|
|
"role": "user",
|
|
"content": data["prompt"]
|
|
})
|
|
rejected.append({
|
|
"role": "assistant",
|
|
"content": data["all_generated_responses"][rejected_idx]
|
|
})
|
|
data.update({
|
|
"chosen": chosen,
|
|
"rejected": rejected,
|
|
})
|
|
|
|
output_file = os.path.basename(args.generation_file).split('.json')[0] + "_bin.json"
|
|
with open(os.path.join(args.output_dir, file_name), 'w') as f:
|
|
json.dump(output_data, f, indent=4)
|
|
print(f"Binarized outputs saved to {output_file}")
|
|
|
|
# Convert the data to Hugging Face datasets format
|
|
dataset = datasets.Dataset.from_list(output_data)
|
|
dataset.save_to_disk(os.path.join(args.output_dir))
|
|
print(f"Binarized dataset saved to {os.path.join(args.output_dir)}")
|