from datasets import load_dataset
from torch.utils.data import DataLoader
from transformers import BertTokenizer, BertModel
import torch
import torch.nn as nn
import torch.optim as optim
from tqdm import tqdm


# 加载预训练模型的分词器
tokenizer = BertTokenizer.from_pretrained("../models/bert-base-chinese")
# 加载预训练模型
bert_model = BertModel.from_pretrained("../models/bert-base-chinese")

# 1.数据预处理
def collate_fn(batch):
    # print("batch-->", batch)
    # 从这个batch中，把文本和标签(不用)拿出来
    sentences = [row['text'] for row in batch]
    # print("sentences-->", sentences)
    # 对文本进行数值化+张量化
    data = tokenizer.batch_encode_plus(sentences,
                                       padding='max_length',
                                       max_length=32,
                                       truncation=True,
                                       return_tensors='pt')
    # print("data-->", data)
    # 取出input_ids、token_type_ids、attention_mask
    input_ids = data['input_ids']
    token_type_ids = data['token_type_ids']
    attention_mask = data['attention_mask']
    # print("input_ids-->", input_ids, input_ids.shape) # [4, 32]
    # print("token_type_ids-->", token_type_ids)
    # print("attention_mask-->", attention_mask)
    # 手动构建y
    # 从数据中挖下标=16的token，当做label标签（y值）
    label = input_ids[:, 16].clone()
    # print("label-->", label)
    # 手动构建x
    # print("token1-->", tokenizer.mask_token)      # MASK
    # print("token2-->", tokenizer.mask_token_id)   # 103
    # 对原始数据进行MASK，构建x
    input_ids[:, 16] = tokenizer.mask_token_id
    # print("input_ids-->", input_ids)
    # 返回x和y
    return input_ids, token_type_ids, attention_mask, label


def load_data():
    # 加载数据
    dataset_tmp = load_dataset('csv', data_files='../data/train.csv', split='train')
    # 对数据进行过滤，过滤长度<=32的样本
    # print("len(dataset)--", len(dataset_tmp)) # 9600
    # print("(dataset[:3])--", dataset_tmp[:3])
    mydataset = dataset_tmp.filter(lambda x: len(x['text']) > 32)
    # print("len(dataset)--", len(mydataset)) # 9035

    # 把数据转换为DataLoader
    mydataloader = DataLoader(dataset=mydataset,
                              batch_size=4,
                              shuffle=True,
                              drop_last=True,
                              collate_fn=collate_fn)
    # 遍历 mydataloader
    for input_ids, token_type_ids, attention_mask, label in mydataloader:
        print("input_ids-->", input_ids)
        print("token_type_ids-->", token_type_ids)
        print("attention_mask-->", attention_mask)
        print("label-->", label)
        break


# 2.模型构建
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.linear = nn.Linear(in_features=768, out_features=tokenizer.vocab_size)

    def forward(self, input_ids, token_type_ids, attention_mask):
        # 让数据先经过预训练模型，然后经过下游线性层
        with torch.no_grad():
            out = bert_model(input_ids=input_ids,
                             token_type_ids=token_type_ids,
                             attention_mask=attention_mask)
        # [4, 32, 768]：4句话，每句话32个token，每个token使用768维度的向量表示
        # print("out.last_hidden_state-->", out.last_hidden_state.shape)
        # out.last_hidden_state[:, 16]: 4句话，16表示每句话中下标为16的token，这就是之前MASK的位置，预测结果
        # return self.linear(out.last_hidden_state[:, 16])
        return self.linear(out.last_hidden_state[:, 16, :])


def test_mymodel():
    dataset_tmp = load_dataset('csv', data_files='../data/train.csv', split='train')
    mydataset = dataset_tmp.filter(lambda x: len(x['text']) > 32)
    # 把数据转换为DataLoader
    mydataloader = DataLoader(dataset=mydataset,
                              batch_size=4,
                              shuffle=True,
                              drop_last=True,
                              collate_fn=collate_fn)
    # 遍历 mydataloader
    model = MyModel()
    for input_ids, token_type_ids, attention_mask, label in mydataloader:
        # 把数据送给模型
        result = model(input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask)
        # [4, 21128]：4句话，每句话都是21128个类别
        print("result-->", result.shape) # [4, 21128]
        break


# 3.模型训练
def train_model():
    # 准备数据
    dataset_tmp = load_dataset('csv', data_files='../data/train.csv', split='train')
    mydataset = dataset_tmp.filter(lambda x: len(x['text']) > 32)
    # 把数据转换为DataLoader
    mydataloader = DataLoader(dataset=mydataset,
                              batch_size=4,
                              drop_last=True,
                              shuffle=True,
                              collate_fn=collate_fn)
    # 准备模型
    model = MyModel()
    # 损失函数
    criterion = nn.CrossEntropyLoss()
    # 优化器
    optimizer = optim.AdamW(params=model.parameters(), lr=5e-4)
    # 冻结预训练模型参数
    for param in bert_model.parameters():
        param.requires_grad_(False)
    # 设置轮次epochs
    epochs = 1
    # 调整模型为训练模式
    model.train()
    # 外层遍历
    for epoch in range(epochs):
        # 内层遍历
        for i, (input_ids, token_type_ids, attention_mask, label) in enumerate(tqdm(mydataloader), start=1):
            # 前向传播
            out = model(input_ids=input_ids,
                        token_type_ids=token_type_ids,
                        attention_mask=attention_mask)
            # 损失计算
            loss = criterion(out, label)
            # print("loss-->", loss)
            # 梯度清零
            optimizer.zero_grad()
            # 反向传播
            loss.backward()
            # 梯度更新
            optimizer.step()
            # 打印日志
            if i % 10 == 0:
                # 得到预测结果
                y_pred = torch.argmax(out, dim=-1)
                # 累计预测对的
                accuracy = (y_pred == label).sum().item() / len(label)
                # 打印准确率
                print("epoch:", epoch + 1, "loss:", loss.item(), "accuracy:", accuracy)
        # 模型保存
        torch.save(model.state_dict(), '../save_model/model2_%d.bin' % (epoch + 1))


# 4.模型预测
def eval_model():
    # 加载数据
    dataset_tmp = load_dataset('csv', data_files='../data/test.csv', split='train')
    # 过滤数据
    mydataset = dataset_tmp.filter(lambda x: len(x['text']) > 32)
    # 转换为DataLoader
    mydataloader = DataLoader(dataset=mydataset,
                              shuffle=False,
                              drop_last=True,
                              batch_size=4,
                              collate_fn=collate_fn)
    # 准备模型
    model = MyModel()
    # 加载模型参数
    model.load_state_dict(torch.load('../save_model/model_2.pth'))
    # 设置模型为评估模式
    model.eval()
    # 定义2个变量
    correct = 0
    total = 0
    # 遍历，把数据送给模型，得到预测结果
    for i, (input_ids, token_type_ids, attention_mask, label) in enumerate(tqdm(mydataloader), start=1):
        # 把数据送给模型，得到预测结果
        out = model(input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask)
        # 累计预测正确的结果和总样本
        y_pred = torch.argmax(out, dim=-1)
        correct += (y_pred == label).sum().item()
        total += len(label)
    # 计算准确率
    print("accuracy-->", correct / total)


if __name__ == '__main__':
    # load_data()
    # test_mymodel()
    # train_model()
    eval_model()