forked from facebookresearch/MetaICL
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathget_pero_prompts.py
212 lines (177 loc) · 8.04 KB
/
get_pero_prompts.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
#!/usr/bin/env python
# coding: utf-8
# In[11]:
import os
import argparse
import pickle as pkl
import random
import torch
import math
import json
import string
import logging
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from tqdm import tqdm
from collections import Counter, defaultdict
from torch.utils.data import TensorDataset, DataLoader, SequentialSampler
from transformers import GPT2Tokenizer, AutoTokenizer
from metaicl.data import MetaICLData
from metaicl.model import MetaICLModel
from utils.data import load_data
class Namespace:
def __init__(self, **kwargs):
self.__dict__.update(kwargs)
def save_task_prompt_weights(task, metaicl_data, metaicl_model, train_data, checkpoint, add_newlines, args):
# predicting on train data here with no prompt
metaicl_data.tensorize([], train_data, add_newlines=add_newlines)
metaicl_data.print_tensorized_example()
if metaicl_model.is_none():
metaicl_model.load(checkpoint, gpt2=args.gpt2)
metaicl_model.cuda()
metaicl_model.eval()
if __name__ == '__main__':
args = Namespace()
args.gpt2 = 'gpt2-large'
args.checkpoint = 'checkpoints/metaicl/hr_to_lr/model.pt'
args.global_step = None
args.use_demonstrations = False
args.do_zeroshot = True
args.k = 0
args.total_data_size = 200
# args.test_batch_size = 8
args.test_batch_size = 16
args.method = 'direct'
args.task = 'custom'
args.unseen_domain_only = False
args.dataset = None
args.split = 'dev'
args.is_null = False
# args.out_dir = None
# args.out_dir = 'results/gpt2_uncertainty_sampling'
args.out_dir = 'results/gpt2finetuned_uncertainty_sampling'
args.seed = '100'
args.use_random_english_words = False
args.use_calibration = False
# args.num_prompt_samples = 32
# args.ks = [0, 1, 2, 4, 8, 16, 32]
# args.trim_dev_data = 32
args.gpt2s = 'gpt2-large'
args.checkpoints = 'checkpoints/metaicl/hr_to_lr/model.pt'
# args.gpt2s = 'gpt2-large'
# args.checkpoints = 'gpt2-large'
# args.gpt2s = 'gpt-j-6B'
# args.checkpoints = 'gpt-j-6B'
args.prompt_with_random_tasks = True
args.sampling_weights_dir = args.out_dir
args.c = 1
if not os.path.exists(args.sampling_weights_dir):
os.mkdir(args.sampling_weights_dir)
# In[4]:
config_split = "unseen_domain_test" if args.unseen_domain_only else "test"
handlers = [logging.StreamHandler()]
# if args.log_file is not None:
# handlers.append(logging.FileHandler(args.log_file))
logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s - %(message)s',
datefmt='%m/%d/%Y %H:%M:%S',
level=logging.INFO,
handlers=handlers)
logger = logging.getLogger(__name__)
logger.info(args)
# In[5]:
from tqdm import tqdm
for gpt2, chckpnt in zip(args.gpt2s.split(','), args.checkpoints.split(',')):
args.gpt2 = gpt2
args.checkpoint = chckpnt
if args.gpt2.startswith("gpt2"):
tokenizer = GPT2Tokenizer.from_pretrained(args.gpt2)
else:
tokenizer = AutoTokenizer.from_pretrained("gpt2")
add_newlines = True
### checkpoint ...
if not args.do_zeroshot:
if args.checkpoint is not None:
checkpoint = args.checkpoint
assert args.global_step is None
else:
assert args.global_step is not None
checkpoint = os.path.join(args.out_dir, "model-{}.pt".format(args.global_step))
# assert os.path.exists(checkpoint)
else:
add_newlines = not args.gpt2.startswith("gpt2")
if args.gpt2 == "gpt-j-6B":
# we are using the HF veresion where GPT-J-6B checkpoint is not officially registered
# so need to download the model checkpoint and specify checkpoint
# assert args.checkpoint is not None and os.path.exists(args.checkpoint)
args.gpt2 = args.checkpoint
checkpoint = None
metaicl_model = MetaICLModel(logger, args.out_dir)
if not os.path.exists(args.out_dir):
os.makedirs(args.out_dir)
# setup hyperparams for data
max_length_per_example = 256
max_length = 256
if args.use_demonstrations:
orig_max_length = max_length
if args.do_zeroshot:
max_length = min(max_length * args.k, 1024)
else:
max_length = min(max_length * args.k, 1024)
logger.info("batch_size=%d\tmax_length=%d\tmax_length_per_example=%d" % (
args.test_batch_size, max_length, max_length_per_example))
metaicl_data = MetaICLData(logger, tokenizer, args.method, args.use_demonstrations, args.k,
max_length, max_length_per_example)
# In[ ]:
seeds = args.seed.split(",")
for seed in seeds:
### data ...
train_data = load_data(args.task, "train", args.total_data_size, seed=seed, config_split=config_split,
datasets=None if args.dataset is None else args.dataset.split(","))
if args.use_random_english_words:
from english_words import english_words_set
english_words_set = sorted(english_words_set)
np.random.seed(int(seed))
train_counter = Counter()
dev_counter = Counter()
for dp in train_data:
train_counter[dp["task"]] += 1
for k, v in train_counter.items():
logger.info("[Train] %s\t%d" % (k, v))
logger.info("%s on %s (%d train)" % (args.method, args.task, len(train_counter)))
for test_task in tqdm(list(train_counter)):
if os.path.exists(os.path.join(args.sampling_weights_dir, test_task + '.pkl')):
continue
task_train_data = [dp for dp in train_data if dp["task"] == test_task]
# if args.sampling_weights_dir is not None:
# assert not args.prompt_with_random_tasks
# with open(os.path.join(args.sampling_weights_dir, test_task + '.pkl'), 'rb') as f:
# task_weights = pkl.load(f)
# else:
# task_weights = None
# assert len(task_dev_data) > 0
assert not args.use_demonstrations or len(task_train_data) == args.total_data_size, \
(args.use_demonstrations, len(task_train_data), args.total_data_size)
config_file = "config/tasks/{}.json".format(test_task)
assert os.path.exists(config_file), config_file
with open(config_file, "r") as f:
config = json.load(f)
is_classification = config["task_type"] == "classification"
if is_classification:
options = task_train_data[0]["options"]
assert np.all([d["options"] == options for d in task_train_data])
if args.use_random_english_words:
# create a mapping
options = task_train_data[0]["options"]
mapping = {option: np.random.choice(english_words_set) for option in options}
new_options = list(mapping.values())
for dp_idx, dp in enumerate(task_train_data):
assert dp["output"] in options, (dp, options)
task_train_data[dp_idx]["output"] = mapping[dp["output"]]
task_train_data[dp_idx]["options"] = new_options
for dp_idx, dp in enumerate(task_train_data):
assert dp["output"] in options, (dp, options)
task_train_data[dp_idx]["output"] = mapping[dp["output"]]
task_train_data[dp_idx]["options"] = new_options
save_task_prompt_weights(
test_task, metaicl_data, metaicl_model, task_train_data, checkpoint, add_newlines, args)