#首先导入需要用到的数据包

from transformers import BertModel, BertTokenizer
import torch.nn as nn
import torch
from torch.utils.data import Dataset, DataLoader
from torch import optim
import os

class BertClassifier(nn.Module):
    def __init__(self, bert_model, output_size):
        super(BertClassifier, self).__init__()
        self.bert = bert_model
        self.classifier = nn.Linear(bert_model.config.hidden_size, output_size)

    def forward(self, input_ids, attention_mask):
        # 获取BERT模型的CLS输出
        text_output = self.bert(input_ids=input_ids, attention_mask=attention_mask)  
        #得到线性层的结果
        logits=self.classifier(text_output.pooler_output)
        return logits


#读取数据
class data_load(Dataset):
    def __init__(self,path):
        self.data=list()
        file=open(path,'r',encoding='utf-8')
        for line in file:
            text,label=line.strip().split('\t')
            self.data.append((text,int(label)))
        file.close()
    def __len__(self):
        return len(self.data)
    def __getitem__(self, index):
        return self.data[index]


#用于dataloader,对于每个小批量的数据，进行分词和填充
def collate_fn(batch,tokenizer):
    texts=[text[0] for text in batch]
    labels=[text[1] for text in batch]
    labels=torch.tensor(labels,dtype=torch.long)
    tokens=tokenizer(
                texts,
                add_special_tokens=True,
                max_length=512,
                padding=True,
                truncation=True,
                return_tensors='pt',
                )
    return tokens['input_ids'],tokens['attention_mask'],labels

if __name__=="__main__":
    dataset=data_load('./train.txt')

    print(len(dataset))
    
    output :180000
            
 
    #加载模型,生成分词器
    tokenizer=BertTokenizer.from_pretrained('bert-base-chinese')
    bert_model = BertModel.from_pretrained('bert-base-chinese')

    #dataset：要加载的数据集对象，必须是实现了len()和getitem()方法的对象
    data_loader=DataLoader(dataset,
                           batch_size=128,
                           shuffle=True,
                           collate_fn=lambda x:collate_fn(x,tokenizer))
    
    # 指定机器
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    #打印分词器支持的最大长度，输入的中文数据不能超过512
    #如果进行长文本分类，需要进行文本截断或分块处理
    # print(tokenizer.model_max_length)
    

    #定义bertclassifier模型为10分类
    model=BertClassifier(bert_model,output_size=10).to(device)
    model.train()
    #优化器
    optimizer=optim.AdamW(model.parameters(),lr=5e-5)
    #交叉熵损失误差
    criterion=nn.CrossEntropyLoss()
    #存放模型
    os.makedirs('output_models',exist_ok=True)
    epoch_n=10
    for epoch in range(1,epoch_n+1):
        for batch_index,data in enumerate(data_loader):
            input_ids=data[0].to(device)
            attention_mask=data[1].to(device)
            label=data[2].to(device)
            #清空梯度
            optimizer.zero_grad()
            #前向传播
            output=model(input_ids,attention_mask)
            loss=criterion(output,label)
            loss.backward() #计算梯度
            optimizer.step()  #更新模型参数
            
            #计算正确率，用于观察模型结果
            predict=torch.argmax(output,dim=1)
            correct=(predict==label).sum().item()
            acc=correct/output.size(0)
            print(f"Epoch {epoch}/{epoch_n}") #迭代轮数
            print(f"Batch {batch_index+1}/{len(data_loader)}") 
            print(f"Loss: {loss.item():.4f}") #损失
            print((f"Acc {correct}/{output.size(0)}=={acc:.3f}")) #正确率
            #每一次迭代都保存一次模型结果
            model_name=f'./output_models/chinese_news_classify{epoch}.pth'
            print("saved model: %s" % (model_name))
            torch.save(model.state_dict(),model_name)