1def generate(prompt='', num_samples=10, steps=20, do_sample=True):
2
3 # tokenize the input prompt into integer input sequence
4 vocab2id, id2vocab = get_vocab_and_token()
5 tokenizer = MyTokenizer(vocab2id, id2vocab)
6
7 x = tokenizer(prompt)
8 x = torch.tensor(x, dtype=torch.long).unsqueeze(0).to(device)
9
10
11 # we'll process all desired num_samples in a batch, so expand out the batch dim
12 x = x.expand(num_samples, -1)
13
14 # forward the model `steps` times to get samples, in a batch
15 y = model.generate(x, max_new_tokens=steps, do_sample=do_sample, top_k=40)
16
17 for i in range(num_samples):
18 pred_idxs = y[i].cpu().squeeze().tolist()
19 out = tokenizer.convert_id_2_token(pred_idxs)
20 print('-' * 80)
21 print(out)
22
23
24if __name__ == '__main__':
25
26 # set_seed(3407)
27 model_config = GPT.get_default_config()
28 model_config.model_type = 'gpt2'
29 model_config.vocab_size = 17543 # openai's model vocabulary
30 model_config.block_size = (prompt_length + info_length) - 1
31 model = GPT(model_config)
32 model.load_state_dict(torch.load('model.bin'))
33 device = 'cuda' if torch.cuda.is_available() else 'cpu'
34 model.to(device)
35 model.eval()
36
37 generate(prompt='什么是唐诗宋词?', num_samples=5, steps=200)1class Trainer:
2
3 @staticmethod
4 def get_default_config():
5 C = CN()
6 # device to train on
7 C.device = 'auto'
8 # dataloder parameters
9 C.num_workers = 4
10 # optimizer parameters
11 C.max_iters = None
12 C.batch_size = 64
13 C.learning_rate = 3e-4
14 C.betas = (0.9, 0.95)
15 C.weight_decay = 0.1 # only applied on matmul weights
16 C.grad_norm_clip = 1.0
17 return C
18
19 def __init__(self, config, model, train_dataset):
20 self.config = config
21 self.model = model
22 self.optimizer = None
23 self.train_dataset = train_dataset
24 self.callbacks = defaultdict(list)
25
26 # determine the device we'll train on
27 if config.device == 'auto':
28 self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
29 else:
30 self.device = config.device
31 self.model = self.model.to(self.device)
32 print("running on device", self.device)
33
34 # variables that will be assigned to trainer class later for logging and etc
35 self.iter_num = 0
36 self.iter_time = 0.0
37 self.iter_dt = 0.0
38
39 def add_callback(self, onevent: str, callback):
40 self.callbacks[onevent].append(callback)
41
42 def set_callback(self, onevent: str, callback):
43 self.callbacks[onevent] = [callback]
44
45 def trigger_callbacks(self, onevent: str):
46 # Trigger batch_end_callback function
47 for callback in self.callbacks.get(onevent, []):
48 callback(self)
49
50 def run(self):
51 model, config = self.model, self.config
52
53 # setup the optimizer
54 self.optimizer = model.configure_optimizers(config)
55
56 # setup the dataloader
57 train_loader = DataLoader(
58 self.train_dataset,
59 sampler=torch.utils.data.RandomSampler(self.train_dataset,
60 replacement=True,
61 num_samples=int(1e10)),
62 shuffle=False,
63 pin_memory=True,
64 batch_size=config.batch_size,
65 num_workers=config.num_workers,
66 )
67
68 model.train()
69 self.iter_num = 0
70 self.iter_time = time.time()
71
72 question_lst = ['解释一下法律法规的意思', '介绍一下原神是什么', '解释一下核心价值的意思',
73 '解释一下发达国家的意思','解释一下汽车电子的意思','解释一下技术革命的意思','什么是资源整合','啥是慈善活动','能给我讲讲青年一代吗',
74 '给我介绍下中华民族','语音识别的含义是什么','请说明一下奥林匹克运动会的含义','能给我讲讲科技成果是什么']
75 with open('data/vocab2id.json', 'r') as f:
76 vocab2id = json.load(f)
77 with open('data/id2vocab.json', 'r') as f:
78 id2vocab = json.load(f)
79 id2vocab = {int(k): v for k, v in id2vocab.items()}
80 tokenizer = MyTokenizer(vocab2id, id2vocab)
81
82 question_token = []
83 for i in range(len(question_lst)):
84 question_token.append(tokenizer(question_lst[i]))
85
86 data_iter = iter(train_loader)
87
88
89 while True:
90
91 # fetch the next batch (x, y) and re-init iterator if needed
92 try:
93 batch = next(data_iter)
94 except StopIteration:
95 data_iter = iter(train_loader)
96 batch = next(data_iter)
97 batch = [t.to(self.device) for t in batch]
98 x, y = batch
99
100 # forward the model
101 logits, self.loss = model(x, y)
102
103 # backprop and update the parameters
104 model.zero_grad(set_to_none=True)
105 self.loss.backward()
106 torch.nn.utils.clip_grad_norm_(model.parameters(),
107 config.grad_norm_clip)
108 self.optimizer.step()
109
110 self.trigger_callbacks('on_batch_end')
111 self.iter_num += 1
112 if self.iter_num % 10000 == 0:
113
114 # question = torch.tensor(question, dtype=torch.long).unsqueeze(0).to(self.device)
115 with torch.no_grad():
116 for i in range(len(question_token)):
117 my_question = torch.tensor(question_token[i], dtype=torch.long).unsqueeze(0).to(self.device)
118 y = model.generate(my_question, max_new_tokens=100, do_sample=True, top_k=40)
119 my_ids = y[0].cpu()
120 my_ids = my_ids.tolist()
121
122 # print(my_ids)
123 CRAWLERLOGGER.debug(f"the prediction: {my_ids}")
124 out = tokenizer.convert_id_2_token(ids=my_ids)
125 # print(out)
126 CRAWLERLOGGER.debug(f"the prediction: {out}")
127
128 tnow = time.time()
129 self.iter_dt = tnow - self.iter_time
130 self.iter_time = tnow
131
132 # termination conditions
133 if config.max_iters is not None and self.iter_num >= config.max_iters:
134 break1class MyTokenizer:
2 def __init__(self, vocab, token, max_len=180):
3 self.vocab = vocab
4 self.token = token
5 self.max_len = max_len
6
7 def __call__(self, str):
8 ids = []
9
10 for w in str:
11 if w in self.vocab.keys():
12 ids.append(self.vocab[w])
13
14 return ids
15
16 def convert_id_2_token(self, ids):
17 token = []
18
19 for id in ids:
20 try:
21 token.append(self.token[id])
22 except:
23 print('the error:', id)
24 exit()
25 # token = [self.token[id] for id in ids]
26 return "".join(token)
27
28 def get_vocab_size(self):
29 return len(self.vocab) + 1
30
31
32def load_resource(df, vocabsize=50257):
33 word_f = {}
34 vocab = {}
35 token = {}
36 for index, row in tqdm(df.iterrows(), desc='building the vocab library'):
37
38 line = row['prompt'] + row['info']
39
40 for ch in line:
41 try:
42 word_f[ch] += 1
43 except:
44 word_f[ch] = 1
45
46 word_f = sorted(word_f.items(), key=lambda kv: (kv[1], kv[0]), reverse=True)
47 token[0] = ""
48 token[1] = "<|endoftext|>"
49 token[2] = "<pad>"
50 id = 3
51 for it in word_f:
52 vocab[it[0]] = id
53 token[id] = it[0]
54 id += 1
55 if id >= vocabsize:
56 break
57
58
59 with open("vocab2id.json", "w") as f:
60 json.dump(vocab, f)
61
62 with open("id2vocab.json", "w") as f:
63 json.dump(token, f)
64
65
66def get_vocab_and_token():
67 with open('/AI_TEAM/colinyang/minGPT/data/vocab2id.json', 'r') as f:
68 v2id = json.load(f)
69 with open('/AI_TEAM/colinyang/minGPT/data/id2vocab.json', 'r') as f:
70 id2v = json.load(f)
71 id2v = {int(k): v for k, v in id2v.items()}
72 return v2id, id2v
73
74
75if __name__ == "__main__":
76 my_df = pd.read_pickle('baidu_wrangling.pkl')
77 load_resource(my_df)
78
79 vocab2id, id2vocab = get_vocab_and_token()
80 print(len(id2vocab))
81 print(id2vocab[2])
82 Tokenizer = MyTokenizer(vocab2id, id2vocab)
83
84 text = "科大讯飞"
85 ids = Tokenizer(text)
86 ids.append(1)
87 print(ids)
88 print(Tokenizer.convert_id_2_token(ids))