Views
No views yet
1from torch.utils.data import dataset
2from datasets import load_dataset, load_from_disk
3from tqdm import tqdm
4from datasets import load_metric
5from transformers import (
6 Seq2SeqTrainer,
7 Seq2SeqTrainingArguments,
8 AutoTokenizer,
9 AutoModelForSeq2SeqLM,
10 DataCollatorForSeq2Seq
11)
12import evaluate
13import os
14from datasets import load_dataset
15import numpy as np
16
17MAX_LENGTH_INPUT = 512+128
18MAX_LENGTH_OUTPUT = 2
19
20from datasets import load_dataset
21
22class Seq2SeqDataset(dataset.Dataset):
23
24 def __init__(self, tokenizer, type_data='train'):
25
26 # Set up the datasets
27 data_path = "CarperAI/openai_summarize_comparisons"
28 if type_data == 'train':
29 dataset = load_dataset("CarperAI/openai_summarize_comparisons", split="train")
30 else:
31 dataset = load_dataset("CarperAI/openai_summarize_comparisons", split="test").select(range(20000))
32 self.prompts = []
33 self.outputs = []
34 inputs = dataset["prompt"]
35 choosen = dataset["chosen"]
36 rejected = dataset["rejected"]
37 for i, (inp, ch, re) in enumerate(zip(inputs, choosen, rejected)):
38 choice_first = np.random.choice([ch, re])
39 res = "A" if choice_first == ch else "B"
40 choice_second = ch if choice_first == re else re
41 prompt = f"""POST: {inp}\n\nRESPONSE A: {choice_first}\n\nRESPONSE B: {choice_second}\n\nWhich response is better? RESPONSE"""
42 output = f"{res}"
43 self.prompts.append(prompt)
44 self.outputs.append(output)
45 print("Example prompt: ", self.prompts[0])
46 print("Example output: ", self.outputs[0])
47 self.tokenizer = tokenizer
48
49 def __len__(self):
50 return len(self.prompts)
51
52 def __getitem__(self, idx):
53 input_text = self.prompts[idx]
54 output_text = self.outputs[idx]
55
56 model_input = self.tokenizer(
57 input_text,
58 max_length=MAX_LENGTH_INPUT,
59 padding='max_length',
60 truncation=True
61 )
62 with self.tokenizer.as_target_tokenizer():
63 labels = self.tokenizer(
64 output_text,
65 max_length=MAX_LENGTH_OUTPUT,
66 padding='max_length',
67 truncation=True
68 )["input_ids"]
69 model_input['labels'] = labels
70 model_input['labels'] = [-100 if token == self.tokenizer.pad_token_id else token for token in model_input['labels']]
71 return model_input
72
73import wandb
74wandb.init(name="stanfordnlp/SteamSHP-flan-t5-xl", project="trlx", entity="pvduy")
75
76
77if __name__=="__main__":
78 config = {
79 "logging_steps": 100,
80 "eval_steps": 100,
81 "save_steps": 500,
82 "batch_size": 4,
83 "batch_size_val": 4,
84 "warmup_steps": 100,
85 "accum_steps": 2,
86 "num_beams": 3,
87 "output_dir": "flan-t5-rm",
88 }
89
90 accuracy_metric = evaluate.load("accuracy")
91 def compute_metrics(pred):
92 labels_ids = pred.label_ids
93 pred_ids = pred.predictions
94 pred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
95 labels_str = tokenizer.batch_decode(labels_ids, skip_special_tokens=True)
96 acc = sum(np.array(labels_str) == np.array(pred_str)) / len(labels_str)
97 return {"accuracy": acc}
98
99 training_args = Seq2SeqTrainingArguments(
100 output_dir=config["output_dir"],
101 do_train=True,
102 num_train_epochs=5,
103 do_eval=False,
104 predict_with_generate=True,
105 adam_beta1=0.9,
106 adam_beta2=0.999,
107 learning_rate=5e-5,
108 half_precision_backend=True,
109 bf16=True,
110 per_device_train_batch_size=config["batch_size"],
111 per_device_eval_batch_size=config["batch_size_val"],
112 logging_steps=config["logging_steps"],
113 evaluation_strategy="epoch",
114 warmup_steps=config["warmup_steps"],
115 eval_accumulation_steps=1,
116 lr_scheduler_type="linear",
117 save_strategy="epoch",
118 gradient_accumulation_steps=config["accum_steps"],
119 deepspeed='configs/ds_configs/ds_config_gpt_2.json',
120 )
121
122 tokenizer = AutoTokenizer.from_pretrained("stanfordnlp/SteamSHP-flan-t5-xl")
123 model = AutoModelForSeq2SeqLM.from_pretrained("stanfordnlp/SteamSHP-flan-t5-xl")
124
125 train_dataset = Seq2SeqDataset(tokenizer, type_data='train')
126 val_dataset = Seq2SeqDataset(tokenizer, type_data='val')
127 print("Train dataset size: ", len(train_dataset))
128 print("Val dataset size: ", len(val_dataset))
129
130 params = sum(p.numel() for p in model.parameters() if p.requires_grad)
131 print(f"Number of trainable parameters: {params}")
132
133 trainer = Seq2SeqTrainer(
134 model=model,
135 tokenizer=tokenizer,
136 args=training_args,
137 train_dataset=train_dataset,
138 eval_dataset=val_dataset,
139 compute_metrics=compute_metrics,
140 )
141
142 trainer.train()1from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
2from datasets import load_dataset
3import numpy as np
4import torch
5from tqdm import tqdm
6dataset = load_dataset("CarperAI/openai_summarize_comparisons", split="test")
7
8tokenizer = AutoTokenizer.from_pretrained("flan-t5-rm/checkpoint-4338/")
9model = AutoModelForSeq2SeqLM.from_pretrained("flan-t5-rm/checkpoint-4338/")
10
11device = "cuda" if torch.cuda.is_available() else "cpu"
12model.to(device)
13
14df = dataset.to_pandas()
15predictions = []
16for i, row in tqdm(df.iterrows(), total=len(df)):
17 prompt = f"""POST: {row["prompt"]}\n\nRESPONSE A: {row["chosen"]}\n\nRESPONSE B: {row["rejected"]}\n\nWhich response is better? RESPONSE"""
18 x = tokenizer([prompt], return_tensors='pt').input_ids.to(device)
19 y = model.generate(x, max_new_tokens=1)
20 predictions.append(tokenizer.batch_decode(y, skip_special_tokens=True)[0])
21
22print("Accuracy: ", sum(np.array(predictions) == 'A') / len(predictions))