# -*- coding: utf-8 -*-

import torch
import torch.nn as nn
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 使用torch.hub加载bert中文模型的字映射器

tokenizer = torch.hub.load('huggingface/pytorch-transformers', 'tokenizer', 'bert-base-chinese')
# 使用torch.hub加载bert中文模型

model =  torch.hub.load('huggingface/pytorch-transformers', 'model', 'bert-base-chinese').to(device)

def get_bert_encode(text_1, text_2, mark=102, max_len=10):
    """
    description: 使用bert中文模型对输入的文本对进行编码
    :param text_1: 代表输入的第一句话
    :param text_2: 代表输入的第二句话
    :param mark: 分隔标记, 是预训练模型tokenizer本身的标记符号, 当输入是两个文本时,
                 得到的index_tokens会以102进行分隔
    :param max_len: 文本的允许最大长度, 也是文本的规范长度即大于该长度要被截断, 小于该长度要进行0补齐
    :return 输入文本的bert编码
    """

    # 使用tokenizer的encode方法对输入的两句文本进行字映射.

    indexed_tokens = tokenizer.encode(text_1, text_2)
    # 准备对映射后的文本进行规范长度处理即大于该长度要被截断, 小于该长度要进行0补齐

    # 所以需要先找到分隔标记的索引位置

    k = indexed_tokens.index(mark)
    # 首先对第一句话进行长度规范因此将indexed_tokens截取到[:k]判断

    if len(indexed_tokens[:k]) >= max_len:
        # 如果大于max_len, 则进行截断

        indexed_tokens_1 = indexed_tokens[:max_len]
    else:
        # 否则使用[0]进行补齐, 补齐的0的个数就是max_len-len(indexed_tokens[:k])

        indexed_tokens_1 = indexed_tokens[:k] + (max_len-len(indexed_tokens[:k]))*[0]

    # 同理下面是对第二句话进行规范长度处理, 因此截取[k:]

    if len(indexed_tokens[k:]) >= max_len:
        # 如果大于max_len, 则进行截断

        indexed_tokens_2 = indexed_tokens[k:k+max_len]
    else:
         # 否则使用[0]进行补齐, 补齐的0的个数就是max_len-len(indexed_tokens[:k])

        indexed_tokens_2 = indexed_tokens[k:] + (max_len-len(indexed_tokens[k:]))*[0]

    # 最后将处理后的indexed_tokens_1和indexed_tokens_2再进行相加

    indexed_tokens = indexed_tokens_1 + indexed_tokens_2
    # 为了让模型在编码时能够更好的区分这两句话, 我们可以使用分隔ids,

    # 它是一个与indexed_tokens等长的向量, 0元素的位置代表是第一句话

    # 1元素的位置代表是第二句话, 长度都是max_len

    segments_ids = [0]*max_len + [1]*max_len
    # 将segments_ids和indexed_tokens转换成模型需要的张量形式

    segments_tensor = torch.tensor([segments_ids])
    tokens_tensor = torch.tensor([indexed_tokens])

    # 模型不自动求解梯度

    with torch.no_grad():
        # 使用bert model进行编码, 传入参数tokens_tensor和segments_tensor得到encoded_layers

        encoded_layers, _ = model(tokens_tensor.to(device), token_type_ids=segments_tensor.to(device))
    return encoded_layers

# -*- coding: utf-8 -*-
import torch
import torch.nn as nn
import torch.nn.functional as F


class Net(nn.Module):
    """定义微调网络的类"""

    def __init__(self, char_size=20, embedding_size=768, dropout=0.2):
        """
        :param char_size: 输入句子中的字符数量, 因为规范后每条句子长度是max_len, 因此char_size为2*max_len
        :param embedding_size: 字嵌入的维度, 因为使用的bert中文模型嵌入维度是768, 因此embedding_size为768
        :param dropout: 为了防止过拟合, 网络中将引入Dropout层, dropout为置0比率, 默认是0.2
        """

        super(Net, self).__init__()
        # 将char_size和embedding_size传入其中

        self.char_size = char_size
        self.embedding_size = embedding_size
        # 实例化化必要的层和层参数:

        # 实例化Dropout层

        self.dropout = nn.Dropout(p=dropout)
        # 实例化第一个全连接层

        self.fc1 = nn.Linear(char_size*embedding_size, 256)
        self.fc2 = nn.Linear(256, 8)
        # 实例化第二个全连接层

        self.fc3 = nn.Linear(8, 2)

    def forward(self, x):
        # 对输入的张量形状进行变换, 以满足接下来层的输入要求

        x = x.view(-1, self.char_size*self.embedding_size)
        # 使用dropout层

        x = self.dropout(x)
        # 使用第一个全连接层并使用relu函数

        x = F.relu(self.fc1(x))
        # 使用dropout层

        x = self.dropout(x)
        # 使用第二个全连接层并使用relu函数

        x = F.relu(self.fc2(x))
        x = self.dropout(x)
        # 使用第三个全连接层并使用relu函数

        x = F.relu(self.fc3(x))
        return x

# -*- coding: utf-8 -*-

import pandas as pd
from sklearn.utils import shuffle
from functools import reduce
from collections import Counter
from bert_chinese_encode import get_bert_encode
import torch
import torch.nn as nn
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

max_len = 10
# 定义batch_size大小

batch_size = 32
# 数据所在路径

data_path = "./train_data.csv"
# 使用pd进行csv数据的读取

data = pd.read_csv(data_path, header=None, sep="\t")

# 打印整体数据集上的正负样本数量

print("数据集的正负样本数量:")
print(dict(Counter(data[0].values)))
# 打乱数据集的顺序

data = shuffle(data, random_state=0).reset_index(drop=True)

# 划分训练集和验证集

split = 0.1
split_point = int(len(data) * split)
valid_data = data[:split_point]
train_data = data[split_point:]

# 验证数据集中的数据总数至少能够满足一个批次

if len(valid_data) < batch_size:
    raise ("Batch size or split not match!")

# def data_loader(data_path, batch_size, ):

def data_loader(train_data, valid_data, batch_size):
    """
    description: 从持久化文件中加载数据, 并划分训练集和验证集及其批次大小
    :param data_path: 训练数据的持久化路径
    :param batch_size: 训练和验证数据集的批次大小
    :param split: 训练集与验证的划分比例
    :return: 训练数据生成器, 验证数据生成器, 训练数据数量, 验证数据数量
    """

    def _loader_generator(data):
        """
        description: 获得训练集/验证集的每个批次数据的生成器
        :param data: 训练数据或验证数据
        :return: 一个批次的训练数据或验证数据的生成器
        """

        # 以每个批次的间隔遍历数据集

        for batch in range(0, len(data), batch_size):
            # 预定于batch数据的张量列表

            batch_encoded = []
            batch_labels = []
            # 将一个bitch_size大小的数据转换成列表形式，[[label, text_1, text_2]]

            # 并进行逐条遍历

            for item in data[batch: batch+batch_size].values.tolist():
                # 每条数据中都包含两句话, 使用bert中文模型进行编码

                encoded = get_bert_encode(item[1], item[2])
                # 将编码后的每条数据装进预先定义好的列表中

                batch_encoded.append(encoded)
                # 同样将对应的该batch的标签装进labels列表中

                batch_labels.append([item[0]])
            # 使用reduce高阶函数将列表中的数据转换成模型需要的张量形式

            # encoded的形状是(batch_size, 2*max_len, embedding_size)

            encoded = reduce(lambda x, y : torch.cat((x, y), dim=0), batch_encoded)
            labels = torch.tensor(reduce(lambda x, y : x + y, batch_labels))
            # 以生成器的方式返回数据和标签

            yield (encoded, labels)

    # 对训练集和验证集分别使用_loader_generator函数, 返回对应的生成器

    # 最后还要返回训练集和验证集的样本数量

    return _loader_generator(train_data), _loader_generator(valid_data), len(train_data), len(valid_data)


# 加载微调网络

from finetuning_net import Net
import torch.optim as optim


# 定义embedding_size, char_size

embedding_size = 768
char_size = 2 * max_len
# 实例化微调网络

net = Net(embedding_size, char_size).to(device)
# 定义交叉熵损失函数

criterion = nn.CrossEntropyLoss()
# 定义SGD优化方法

optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)


def train(train_data_labels):
    """
    description: 训练函数, 在这个过程中将更新模型参数, 并收集准确率和损失
    :param train_data_labels: 训练数据和标签的生成器对象
    :return: 整个训练过程的平均损失之和以及正确标签的累加数
    """

    # 定义训练过程的初始损失和准确率累加数

    train_running_loss = 0.0
    train_running_acc = 0.0
    # 循环遍历训练数据和标签生成器, 每个批次更新一次模型参数

    for train_tensor, train_labels in train_data_labels:
        train_tensor = train_tensor.to(device)
        train_labels = train_labels.to(device)
        # 初始化该批次的优化器

        optimizer.zero_grad()
        # 使用微调网络获得输出

        train_outputs = net(train_tensor)
        # 得到该批次下的平均损失

        train_loss = criterion(train_outputs, train_labels)
        # 将该批次的平均损失加到train_running_loss中

        train_running_loss += train_loss.item()
        # 损失反向传播

        train_loss.backward()
        # 优化器更新模型参数

        optimizer.step()
        # 将该批次中正确的标签数量进行累加, 以便之后计算准确率

        train_running_acc += (train_outputs.argmax(1) == train_labels).sum().item()
    return train_running_loss, train_running_acc

def valid(valid_data_labels):
    """
    description: 验证函数, 在这个过程中将验证模型的在新数据集上的标签, 收集损失和准确率
    :param valid_data_labels: 验证数据和标签的生成器对象
    :return: 整个验证过程的平均损失之和以及正确标签的累加数
    """

    # 定义训练过程的初始损失和准确率累加数

    valid_running_loss = 0.0
    valid_running_acc = 0.0
    # 循环遍历验证数据和标签生成器

    for valid_tensor, valid_labels in valid_data_labels:
        valid_tensor = valid_tensor.to(device)
        valid_labels = valid_labels.to(device)
        # 不自动更新梯度

        with torch.no_grad():
            # 使用微调网络获得输出

            valid_outputs = net(valid_tensor)
            # 得到该批次下的平均损失

            valid_loss = criterion(valid_outputs, valid_labels)
            # 将该批次的平均损失加到valid_running_loss中

            valid_running_loss += valid_loss.item()
            # 将该批次中正确的标签数量进行累加, 以便之后计算准确率

            valid_running_acc += (valid_outputs.argmax(1) == valid_labels).sum().item()
    return valid_running_loss,  valid_running_acc

# 定义训练轮数

epochs = 5

# 定义盛装每轮次的损失和准确率列表, 用于制图

all_train_losses = []
all_valid_losses = []
all_train_acc = []
all_valid_acc = []

# 进行指定轮次的训练

for epoch in range(epochs):
    # 打印轮次

    print("Epoch:", epoch + 1)
    # 通过数据加载器获得训练数据和验证数据生成器, 以及对应的样本数量

    train_data_labels, valid_data_labels, train_data_len, valid_data_len = \
        data_loader(train_data, valid_data, batch_size)
    # 调用训练函数进行训练

    train_running_loss, train_running_acc = train(train_data_labels)
    # 调用验证函数进行验证

    valid_running_loss, valid_running_acc = valid(valid_data_labels)
    # 计算每一轮的平均损失, train_running_loss和valid_running_loss是每个批次的平均损失之和

    # 因此将它们乘以batch_size就得到了该轮的总损失, 除以样本数即该轮次的平均损失

    train_average_loss = train_running_loss * batch_size / train_data_len
    valid_average_loss = valid_running_loss * batch_size / valid_data_len

    # train_running_acc和valid_running_acc是每个批次的正确标签累加和,

    # 因此只需除以对应样本总数即是该轮次的准确率

    train_average_acc = train_running_acc /  train_data_len
    valid_average_acc = valid_running_acc / valid_data_len
    # 将该轮次的损失和准确率装进全局损失和准确率列表中, 以便制图

    all_train_losses.append(train_average_loss)
    all_valid_losses.append(valid_average_loss)
    all_train_acc.append(train_average_acc)
    all_valid_acc.append(valid_average_acc)
    # 打印该轮次下的训练损失和准确率以及验证损失和准确率

    print("Train Loss:", train_average_loss, "|", "Train Acc:", train_average_acc)
    print("Valid Loss:", valid_average_loss, "|", "Valid Acc:", valid_average_acc)


print('Finished Training')


# 导入制图工具包

import matplotlib.pyplot as plt
from matplotlib.pyplot import MultipleLocator


# 创建第一张画布

plt.figure(0)

# 绘制训练损失曲线

plt.plot(all_train_losses, label="Train Loss")
# 绘制验证损失曲线, 颜色为红色

plt.plot(all_valid_losses, color="red", label="Valid Loss")
# 定义横坐标刻度间隔对象, 间隔为1, 代表每一轮次

x_major_locator=MultipleLocator(1)
# 获得当前坐标图句柄

ax=plt.gca()
# 设置横坐标刻度间隔

ax.xaxis.set_major_locator(x_major_locator)
# 设置横坐标取值范围

plt.xlim(1,epochs)
# 曲线说明在左上方

plt.legend(loc='upper left')
# 保存图片

plt.savefig("./loss.png")



# 创建第二张画布

plt.figure(1)

# 绘制训练准确率曲线

plt.plot(all_train_acc, label="Train Acc")

# 绘制验证准确率曲线, 颜色为红色

plt.plot(all_valid_acc, color="red", label="Valid Acc")
# 定义横坐标刻度间隔对象, 间隔为1, 代表每一轮次

x_major_locator=MultipleLocator(1)
# 获得当前坐标图句柄

ax=plt.gca()
# 设置横坐标刻度间隔

ax.xaxis.set_major_locator(x_major_locator)
# 设置横坐标取值范围

plt.xlim(1,epochs)
# 曲线说明在左上方

plt.legend(loc='upper left')
# 保存图片

plt.savefig("./acc.png")

import time
# 模型保存时间

time_ = int(time.time())
# 保存路径

MODEL_PATH = './model/BERT_net_%d.pth' % time_
# 保存模型参数

torch.save(net.state_dict(), MODEL_PATH)

from flask import Flask
from flask import request
app = Flask(__name__)


import torch
# 导入中文预训练模型编码函数

from bert_chinese_encode import get_bert_encode
# 导入微调网络

from finetuning_net import Net

# 导入训练好的模型

MODEL_PATH = "./model/BERT_net.pth"
# 定义实例化模型参数

embedding_size = 768
char_size = 20
dropout = 0.2

# 初始化微调网络模型

net = Net(embedding_size, char_size, dropout)
# 加载模型参数

net.load_state_dict(torch.load(MODEL_PATH))
# 使用评估模式

net.eval()

# 定义服务请求路径和方式

@app.route('/v1/recognition/', methods=["POST"])
def recognition():
    # 接收数据

    text_1 = request.form['text1']
    text_2 = request.form['text2']
    # 对原始文本进行编码

    inputs = get_bert_encode(text_1, text_2, mark=102, max_len=10)
    # 使用微调模型进行预测

    outputs = net(inputs)
    # 获得预测结果

    _, predicted = torch.max(outputs, 1)
    # 返回字符串类型的结果

    return str(predicted.item())

import requests

url = "http://0.0.0.0:5001/v1/recognition/"
data = {"text1":"人生该如何起头", "text2": "改变要如何起手"}
res = requests.post(url, data=data)

print("预测样本:", data["text_1"], "|", data["text_2"])
print("预测结果:", res.text)