forked from 920232796/bert_seq2seq
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
11 changed files
with
266 additions
and
15 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,212 @@ | ||
## THUCNews 原始数据集 | ||
import torch | ||
from tqdm import tqdm | ||
import torch.nn as nn | ||
from torch.optim import Adam | ||
import numpy as np | ||
import os | ||
import json | ||
import time | ||
import glob | ||
import bert_seq2seq | ||
from torch.utils.data import Dataset, DataLoader | ||
from bert_seq2seq.tokenizer import Tokenizer, load_chinese_base_vocab | ||
from bert_seq2seq.utils import load_bert, load_model_params, load_recent_model | ||
|
||
vocab_path = "./state_dict/roberta_wwm_vocab.txt" # roberta模型字典的位置 | ||
word2idx, keep_tokens = load_chinese_base_vocab(vocab_path, simplfied=True) | ||
model_name = "roberta" # 选择模型名字 | ||
model_path = "./state_dict/roberta_wwm_pytorch_model.bin" # 模型位置 | ||
recent_model_path = "./state_dict/bert_auto_title_model.bin" # 用于把已经训练好的模型继续训练 | ||
model_save_path = "./state_dict/bert_auto_title_model.bin" | ||
batch_size = 16 | ||
lr = 1e-5 | ||
maxlen = 256 | ||
|
||
class BertDataset(Dataset): | ||
""" | ||
针对特定数据集,定义一个相关的取数据的方式 | ||
""" | ||
def __init__(self) : | ||
## 一般init函数是加载所有数据 | ||
super(BertDataset, self).__init__() | ||
## 拿到所有文件名字 | ||
self.txts = glob.glob('./state_dict/THUCNews/*/*.txt') | ||
|
||
self.idx2word = {k: v for v, k in word2idx.items()} | ||
self.tokenizer = Tokenizer(word2idx) | ||
|
||
def __getitem__(self, i): | ||
## 得到单个数据 | ||
# print(i) | ||
text_name = self.txts[i] | ||
with open(text_name, "r", encoding="utf-8") as f: | ||
text = f.read() | ||
text = text.split('\n') | ||
if len(text) > 1: | ||
title = text[0] | ||
content = '\n'.join(text[1:]) | ||
token_ids, token_type_ids = self.tokenizer.encode( | ||
content, title, max_length=maxlen | ||
) | ||
output = { | ||
"token_ids": token_ids, | ||
"token_type_ids": token_type_ids, | ||
} | ||
return output | ||
|
||
self.__getitem__(i + 1) | ||
|
||
def __len__(self): | ||
|
||
return len(self.txts) | ||
|
||
def collate_fn(batch): | ||
""" | ||
动态padding, batch为一部分sample | ||
""" | ||
|
||
def padding(indice, max_length, pad_idx=0): | ||
""" | ||
pad 函数 | ||
""" | ||
pad_indice = [item + [pad_idx] * max(0, max_length - len(item)) for item in indice] | ||
return torch.tensor(pad_indice) | ||
|
||
token_ids = [data["token_ids"] for data in batch] | ||
max_length = max([len(t) for t in token_ids]) | ||
token_type_ids = [data["token_type_ids"] for data in batch] | ||
|
||
token_ids_padded = padding(token_ids, max_length) | ||
token_type_ids_padded = padding(token_type_ids, max_length) | ||
target_ids_padded = token_ids_padded[:, 1:].contiguous() | ||
|
||
return token_ids_padded, token_type_ids_padded, target_ids_padded | ||
|
||
class Trainer: | ||
def __init__(self): | ||
# 判断是否有可用GPU | ||
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | ||
print("device: " + str(self.device)) | ||
# 定义模型 | ||
self.bert_model = load_bert(word2idx, model_name=model_name) | ||
## 加载预训练的模型参数~ | ||
|
||
load_model_params(self.bert_model, model_path, keep_tokens=keep_tokens) | ||
# 加载已经训练好的模型,继续训练 | ||
# load_recent_model(self.bert_model, recent_model_path) | ||
|
||
# 将模型发送到计算设备(GPU或CPU) | ||
self.bert_model.to(self.device) | ||
# 声明需要优化的参数 | ||
self.optim_parameters = list(self.bert_model.parameters()) | ||
self.optimizer = torch.optim.Adam(self.optim_parameters, lr=lr, weight_decay=1e-3) | ||
# 声明自定义的数据加载器 | ||
dataset = BertDataset() | ||
self.dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn) | ||
|
||
def train(self, epoch): | ||
# 一个epoch的训练 | ||
self.bert_model.train() | ||
self.iteration(epoch, dataloader=self.dataloader, train=True) | ||
|
||
def save(self, save_path): | ||
""" | ||
保存模型 | ||
""" | ||
torch.save(self.bert_model.state_dict(), save_path) | ||
print("{} saved!".format(save_path)) | ||
|
||
def iteration(self, epoch, dataloader, train=True): | ||
total_loss = 0 | ||
start_time = time.time() ## 得到当前时间 | ||
step = 0 | ||
report_loss = 0 | ||
for token_ids, token_type_ids, target_ids in tqdm(dataloader,position=0, leave=True): | ||
step += 1 | ||
if step % 1000 == 0: | ||
self.bert_model.eval() | ||
test_data = ["夏天来临,皮肤在强烈紫外线的照射下,晒伤不可避免,因此,晒后及时修复显得尤为重要,否则可能会造成长期伤害。专家表示,选择晒后护肤品要慎重,芦荟凝胶是最安全,有效的一种选择,晒伤严重者,还请及 时 就医 。", | ||
"2007年乔布斯向人们展示iPhone并宣称它将会改变世界还有人认为他在夸大其词然而在8年后以iPhone为代表的触屏智能手机已经席卷全球各个角落未来智能手机将会成为真正的个人电脑为人类发展做出更大的贡献", | ||
"8月28日,网络爆料称,华住集团旗下连锁酒店用户数据疑似发生泄露。从卖家发布的内容看,数据包含华住旗下汉庭、禧玥、桔子、宜必思等10余个品牌酒店的住客信息。泄露的信息包括华住官网注册资料、酒店入住登记的身份信息及酒店开房记录,住客姓名、手机号、邮箱、身份证号、登录账号密码等。卖家对这个约5亿条数据打包出售。第三方安全平台威胁猎人对信息出售者提供的三万条数据进行验证,认为数据真实性非常高。当天下午 ,华 住集 团发声明称,已在内部迅速开展核查,并第一时间报警。当晚,上海警方消息称,接到华住集团报案,警方已经介入调查。"] | ||
for text in test_data: | ||
print(self.bert_model.generate(text, beam_size=3,device=self.device)) | ||
print("loss is " + str(report_loss)) | ||
report_loss = 0 | ||
# self.eval(epoch) | ||
self.bert_model.train() | ||
if step % 8000 == 0: | ||
self.save(model_save_path) | ||
|
||
token_ids = token_ids.to(self.device) | ||
token_type_ids = token_type_ids.to(self.device) | ||
target_ids = target_ids.to(self.device) | ||
# 因为传入了target标签,因此会计算loss并且返回 | ||
predictions, loss = self.bert_model(token_ids, | ||
token_type_ids, | ||
labels=target_ids, | ||
device=self.device | ||
) | ||
report_loss += loss.item() | ||
# 反向传播 | ||
if train: | ||
# 清空之前的梯度 | ||
self.optimizer.zero_grad() | ||
# 反向传播, 获取新的梯度 | ||
loss.backward() | ||
# 用获取的梯度更新模型参数 | ||
self.optimizer.step() | ||
|
||
# 为计算当前epoch的平均loss | ||
total_loss += loss.item() | ||
|
||
end_time = time.time() | ||
spend_time = end_time - start_time | ||
# 打印训练信息 | ||
print("epoch is " + str(epoch)+". loss is " + str(total_loss) + ". spend time is "+ str(spend_time)) | ||
# 保存模型 | ||
self.save(model_save_path) | ||
|
||
if __name__ == '__main__': | ||
|
||
# src, tgt = read_file("./data/train.src", "./data/train.tgt") | ||
|
||
trainer = Trainer() | ||
train_epoches = 20 | ||
|
||
for epoch in range(train_epoches): | ||
# 训练一个epoch | ||
trainer.train(epoch) | ||
|
||
# 测试一下自定义数据集 | ||
# vocab_path = "./state_dict/roberta_wwm_vocab.txt" # roberta模型字典的位置 | ||
# # sents_src, sents_tgt = read_file("./corpus/auto_title/train.src", "./corpus/auto_title/train.tgt") | ||
# sents_src= torch.load("./corpus/auto_title/train_clean.src") | ||
# sents_tgt = torch.load("./corpus/auto_title/train_clean.tgt") | ||
# import time | ||
# dataset = BertDataset(sents_src, sents_tgt, vocab_path) | ||
# word2idx = load_chinese_base_vocab(vocab_path) | ||
# tokenier = Tokenizer(word2idx) | ||
# dataloader = DataLoader(dataset, batch_size=2, shuffle=True, collate_fn=collate_fn) | ||
# for token_ids, token_type_ids, target_ids in dataloader: | ||
# # print(token_ids.shape) | ||
# print(tokenier.decode(token_ids[0].tolist())) | ||
# print(tokenier.decode(token_ids[1].tolist())) | ||
# # print(token_type_ids) | ||
# # print(target_ids.shape) | ||
# # print(tokenier.decode(target_ids[0].tolist())) | ||
# # print(tokenier.decode(target_ids[1].tolist())) | ||
# break | ||
|
||
|
||
# src, tgt = read_file("./corpus/auto_title/train.src", "./corpus/auto_title/train.tgt") | ||
# save_src, save_tgt = [], [] | ||
# for src_i, tgt_i in zip(src, tgt): | ||
# src_i = src_i.replace("“", "").replace("”", "").replace("——", "-").replace("—", "-") | ||
# tgt_i = tgt_i.replace("“", "").replace("”", "").replace("——", "-").replace("—", "-") | ||
|
||
# save_src.append(src_i) | ||
# save_tgt.append(tgt_i) | ||
|
||
# torch.save(save_src, "./corpus/auto_title/train_clean.src") | ||
# torch.save(save_tgt, "./corpus/auto_title/train_clean.tgt") |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,30 @@ | ||
import torch | ||
import torch.nn as nn | ||
import sys | ||
sys.path.append("/Users/xingzhaohu/Downloads/code/python/ml/ml_code/bert/bert_seq2seq") | ||
from torch.optim import Adam | ||
import pandas as pd | ||
import numpy as np | ||
import os | ||
import json | ||
import time | ||
import bert_seq2seq | ||
from bert_seq2seq.tokenizer import Tokenizer, load_chinese_base_vocab | ||
from bert_seq2seq.utils import load_bert, load_model_params, load_recent_model | ||
|
||
auto_title_model = "./state_dict/bert_model_poem.bin" | ||
|
||
if __name__ == "__main__": | ||
vocab_path = "./state_dict/roberta_wwm_vocab.txt" # roberta模型字典的位置 | ||
model_name = "roberta" # 选择模型名字 | ||
# model_path = "./state_dict/bert-base-chinese-pytorch_model.bin" # roberta模型位 | ||
# 加载字典 | ||
word2idx, keep_tokens = load_chinese_base_vocab(vocab_path, simplfied=True) | ||
# 定义模型 | ||
bert_model = load_bert(word2idx, model_name=model_name) | ||
load_model_params(bert_model, "./state_dict/roberta_wwm_pytorch_model.bin", keep_tokens=keep_tokens) | ||
|
||
for name, params in bert_model.named_parameters(): | ||
print(name) | ||
|
||
|