<repo_name>xwwwwww/text-classification<file_sep>/README.md
### 1. 文件结构

![image-20200531092115476](../../pictures/typora_pic/image-20200531092115476.png)

```
-- bert/　bert模型实现

　-- chinese_wwm_ext_pytorch是bert的预训练数据

    -- main.py utils.py 是head+tail做法的Bert实现

    -- myBert目录下是分段输入Bert的实现，如报告中所说，我因显存限制未能测试模型效果，因此报告中没有汇报
　　这一方法的性能，也没有提供模型参数，无需复现。

-- train.py model.py solve.py提供了RNN,MLP,CNN,RNN+Attention的实现

-- newdata目录下是数据，已处理完毕
-- model_save目录下是保存的模型数据，提交时会删除掉，需下载此目录，注意保持目录结构，不要新增嵌套。
	清华云盘下载链接: https://cloud.tsinghua.edu.cn/d/f885293c5c444630a084/
-- word2vec目录下是词向量，提交时会删除掉，需下载至此目录,注意保持目录结构，不要新增嵌套。
	清华云盘下载链接:https://cloud.tsinghua.edu.cn/d/491f0e4168fb433ba960/
```



### 2. 复现方式

-e开启evaluate mode, --cuda指定GPU

#### MLP

```
python train.py --model_type MLP --lr 0.0001 -b 128 --wd 0.001 --name mixed_vector --vector mixed --fix_length 1560 -e --cuda 3
```



#### RNN

```
python train.py --model_type RNN --lr 1e-3 -b 16 --wd 1e-5 --name 1 --vector sogou --min_freq 1 -e --cuda 3
```



#### RNN+Attention

##### biGRU+Attention

```
 python train.py --model_type AttRNN --lr 1e-3 -b 32 --wd 1e-4 --name gru_1e-3_att --vector sogou --min_freq 1 --max_epoch 100 -e --cuda 3
```

##### biLSTM+Attention

```
  python train.py --model_type AttRNN --lr 1e-4 -b 32 --wd 1e-4 --name lstm --vector sogou --min_freq 1 --cell LSTM   --max_epoch 100 -e --cuda 3
```



#### BERT

```
cd bert/

python main.py --lr 1e-6 -b 6 --wd 0.0005 --max_epoch 20 --name fix_np --cuda 1 -e
```

<file_sep>/bert/myBert/utils_myBert.py
import torch
import torch.autograd as autograd
import torch.nn.functional as F
from torchtext import data
from torchtext.vocab import Vectors, Vocab
from collections import Counter
import numpy as np
import sys
from scipy.stats import pearsonr
from sklearn.metrics import f1_score
import numpy as np 

CLS = '[CLS]'
UNK = '[UNK]'
SEP = '[SEP]'
PAD = '[PAD]'


def field_load_vocab(field, vocab_path, novector=False):
    if not novector:
        vector=Vectors(name=vocab_path)
        counter=Counter()
        with open(vocab_path,'r',encoding='utf-8') as file:
            # file head
            file.readline()
            q=file.readline()
            while q:
                q=q.strip()
                counter[q.split()[0]]=1
                q=file.readline()
        field.vocab=Vocab(counter,vectors=vector)
    else:
        counter=Counter()
        with open(vocab_path,'r',encoding='utf-8') as file:
            lines=file.readlines()
            sum=1000000
            for index,line in enumerate(lines):
                line=line.strip()
                # 构造Vocab时，词频大的词排在词表的前面
                counter[line]=sum-index
        field.vocab=Vocab(counter,specials=[])


def train_one_epoch(model, train_iter, criterion, optimizer, epoch, args):
    model.train()

    train_loss = 0
    train_acc = 0
    times = 0

    print_loss = 0
    print_acc = 0

    for batch in train_iter:
        text, ori_label = batch.text, batch.label
        label = torch.argmax(ori_label,dim=1)
        model.zero_grad()
        output = model(text)
        loss = criterion(output,label)
        _, pred = torch.max(output.data, dim=1)
        acc = (pred == label).sum().item() / label.shape[0]

        train_loss += loss.item()
        train_acc += acc
        print_loss += loss.item()
        print_acc += acc
        times += 1

        if times % args.print_frequency == 0:
            print('Train Epoch: {}, Iteration:{}, loss: {:.6f}, acc: {:.6f}'.format(
                epoch, times, print_loss / args.print_frequency, print_acc / args.print_frequency))
            sys.stdout.flush()

            print_loss = 0
            print_acc = 0
        
        loss.backward()
        optimizer.step()

    train_loss /= times
    train_acc /= times
    return train_loss, train_acc


def evaluate(model, eval_iter,criterion):
    model.eval()
    test_loss = 0
    test_acc = 0
    times = 0

    f1 = 0
    corr = 0
    for batch in eval_iter:
        text, ori_label = batch.text, batch.label
        label = torch.argmax(ori_label,dim=1)
        output = model(text)
        loss = criterion(output,label)
        _, pred = torch.max(output.data, dim=1)
        acc = (pred == label).sum().item() / label.shape[0]  # 平均值

        temp_corr = 0
        for i in range(output.shape[0]):
            temp_corr += pearsonr(output[i].cpu().detach().numpy(), ori_label[i].cpu().detach().numpy())[0]
        temp_corr /= output.shape[0]
        corr += temp_corr

        f1 += f1_score(label.cpu().detach().numpy(), pred.cpu().detach().numpy(), np.arange(8),average='micro')

        test_loss += loss.item()
        test_acc += acc
        times += 1

    test_loss /= times
    test_acc /= times
    corr /= times
    f1 /= times

    print('\nValidation set: Average loss: {:.4f}, Accuracy: {:.4f} ({:.2f}%, Correlation:{:4f}), F-score:{:4f})\n'.format(
        test_loss, test_acc, 100. * test_acc, corr, f1))
    return test_loss, test_acc,corr,f1


if __name__=='__main__':
    test()
<file_sep>/model.py
import torch
import torch.nn as nn
import torch.nn.functional as F


class RNN(nn.Module):
    def __init__(self, vocab, input_size, hidden_size, fc_size, num_layers, cell,dropout):
        super(RNN, self).__init__()
        self.embedding = nn.Embedding(len(vocab), input_size)
        self.embedding.weight.requires_grad = False
        self.dropout = nn.Dropout(dropout)
        
        if cell == 'GRU':
            self.rnn = nn.GRU(
                input_size=input_size,
                hidden_size=hidden_size,
                num_layers=num_layers,
                #dropout=dropout,
                bias=True,
                #batch_first = True,
                bidirectional=True,
            )
        elif cell == 'LSTM':
            self.rnn = nn.LSTM(
                input_size =input_size,
                hidden_size = hidden_size,
                num_layers = num_layers,
                #dropout = dropout,
                bias=True,
                bidirectional =True
            )
        
        self.fc1 = nn.Linear(
            in_features=hidden_size * 2, 
            out_features=8,
            bias=True
        )

    def forward(self, x):
        x = self.embedding(x)
        x = self.dropout(x)
        output,h = self.rnn(x)
        x = output[-1]
        x = self.fc1(x)
        return x


class AttentionRNN(nn.Module):
    def __init__(self, vocab, input_size, hidden_size, fc_size, num_layers,cell, dropout):
        super(AttentionRNN, self).__init__()
        self.embedding = nn.Embedding(len(vocab), input_size)
        self.embedding.weight.data.copy_(vocab.vectors)
        self.embedding.weight.requires_grad = False
        self.dropout = nn.Dropout(dropout)
        self.w = nn.Parameter(torch.randn((input_size*2,1)))
        
        if cell =='GRU':
            self.rnn = nn.GRU(
                input_size=input_size,
                hidden_size=hidden_size,
                num_layers=num_layers,
                #dropout=dropout,
                bias=True,
                batch_first = True,
                bidirectional=True,
            )
        
        elif cell =='LSTM':
            self.rnn = nn.LSTM(
                input_size =input_size,
                hidden_size = hidden_size,
                num_layers = num_layers,
                #dropout = dropout,
                bias=True,
                bidirectional =True
            )
    
        self.fc = nn.Linear(
            in_features=hidden_size*2,
            out_features=8,  #8种情感
            bias=True
        )

    def forward(self, x):
        x = self.embedding(x)
        x = self.dropout(x)

        output,h = self.rnn(x)
        output = torch.transpose(output, 1,2)

        m = torch.tanh(output)
        a = F.softmax(torch.matmul(self.w.t(),m), dim=2)   #w: []
        r = torch.matmul(output, torch.transpose(a,1,2))
        r = r.squeeze(2)
        r = torch.tanh(r)
        x = self.fc(r)
        return x


class CNN(nn.Module):
    def __init__(self, vocab, input_size,fix_length, dropout):
            super(CNN, self).__init__()
            self.fix_length = fix_length
            self.embedding = nn.Embedding(len(vocab), input_size)
            self.embedding.weight.data.copy_(vocab.vectors)
            self.embedding.weight.requires_grad = True  #默认是false　不参与训练
            
            self.filter_sizes=[3,4,5]
            self.features = 100#300
            self.convs = nn.ModuleList([nn.Conv2d(1,self.features,(k, input_size) )for k in self.filter_sizes])

            self.dropout = nn.Dropout(dropout)
            
            self.fc = nn.Linear(self.features * len(self.filter_sizes), 8)

    def conv_and_pool(self,x,conv):
        x = conv(x)
        x = F.relu(x).squeeze(3)
        x = F.max_pool1d(x,x.size(2)).squeeze(2)
        return x

    def forward(self,x):
        
            x = self.embedding(x)
            x = x.unsqueeze(1)
            x = torch.cat([self.conv_and_pool(x,conv) for conv in self.convs],1)
            x = self.dropout(x)
            x = self.fc(x)
            return x


class MLP(nn.Module):
    def __init__(self, vocab, input_size,fix_length, dropout):
            super(MLP, self).__init__()
            self.fix_length = fix_length
            self.input_size = input_size
            self.embedding = nn.Embedding(len(vocab), input_size)
            self.embedding.weight.data.copy_(vocab.vectors)
            self.embedding.weight.requires_grad = False  #默认是false　不参与训练
            self.fc1 = nn.Linear(fix_length*input_size,512)
            self.fc2 = nn.Linear(512,32)
            self.dropout = nn.Dropout(dropout)
            self.fc3 = nn .Linear(32,8)
    
    def forward(self,x):
            x = self.embedding(x)
            x = x.view(-1,self.fix_length*self.input_size)
            x = self.fc1(x)
            x = F.relu(x)
            x = self.dropout(x)
            x = self.fc2(x)
            x = F.relu(x)
            x = self.fc3(x)
            return x 
<file_sep>/bert/myBert/bert.py
import torch
import torch.nn as nn
import torch.nn.functional as F


class BERT(nn.Module):
    def __init__(self, baseBert, batch_size,device):
        super(BERT,self).__init__()
        self.baseBert = baseBert #加载的预训练BERT
        print(device)
        self.device = device
        self.batch_size =batch_size
        self.fc = nn.Linear(192,8)
        #self.fc = nn.Linear(768,8)
    def forward(self, x):
        text_len = x.shape[1]
        output_list = []
        step = 256-32 #
        length = 256
        start = -1*step
        end = -1
        flag = False
        while not flag and (start < text_len):  #128区间　64步长滑动
            if (start+step) >= text_len:
               start = text_len
            else:
               start = start+step
       
       	    if (start+length) >= text_len:
               end = text_len
            else:
               end = start+length

            if end == text_len: #标记本次为最后一次滑动
                flag = True
            
            temp_text = x[:,start:end] 
            self.cls_column = torch.LongTensor([[101] for i in range(x.shape[0])]).to(self.device)
            self.sep_column = torch.LongTensor([[102] for i in range(x.shape[0])]).to(self.device)
            temp_text = torch.cat([self.cls_column, temp_text, self.sep_column], dim = 1)

            #在一整个batch加入clk和sep的id
            mask = temp_text>0
            #print('temp_text.shape =',temp_text.shape)
            output = self.baseBert(
                temp_text,
                attention_mask=mask,
            )
            output = output[1] #output[1]是clk的输出 　维度是[batch_size, hidden_size]
            output_list.append(output)
        #print('list len = ', len(output_list))

        kernel_size = len(output_list)*4
        concat_output = torch.cat(output_list,1)
        concat_output = concat_output.unsqueeze(1)
        final_output = F.max_pool1d(concat_output, kernel_size=kernel_size) #max_pooling
        final_output = final_output.squeeze(1)
        x = self.fc(final_output)
        return x
<file_sep>/train.py
import torch
import torch.nn as nn
import argparse
import time
import numpy as np
from torchtext import data
from torchtext.vocab import Vectors, GloVe
import jieba
import os
from model import RNN,CNN,MLP,AttentionRNN
from solve import train_one_epoch, evaluate


parser = argparse.ArgumentParser()
# learning
parser.add_argument('--dropout',default=0.5,type=float)
parser.add_argument('--fix_length',default='512',type=int)
parser.add_argument('--model_type',default='RNN',type=str,help="model used, model list = ['RNN','CNN','MLP','AttRNN']")
parser.add_argument('--max_epoch', default=20, type=int, metavar='N', help='number of total epochs to run')
parser.add_argument('-b', '--batch-size', default=64, type=int, metavar='N', help='mini-batch size(default: 1')
parser.add_argument('--lr', '--learning-rate', default=0.001, type=float, metavar='LR', help='initial learning rate',
                    dest='lr')
parser.add_argument('--momentum', default=0.9, type=float, metavar='M', help='momentum')
parser.add_argument('--wd', '--weight-decay', default=0.0005, type=float, metavar='W',
                    help='weight decay(default: 1e-4)', dest='weight_decay')
parser.add_argument('-print-frequency', type=int, default=10,
                    help='how many steps to wait before logging training status [default: 50]')
parser.add_argument('-e', '--evaluate', dest='evaluate', action='store_true', help='evaluate model on validation set')
parser.add_argument('--cuda',default=None)
parser.add_argument('--name', default='name',help = 'name of the experiment')
parser.add_argument('--vector', type=str, default='sogou')
parser.add_argument('--min_freq',type=int,default=1)
parser.add_argument('--cell',default='GRU')

# device
# parser.add_argument('--device', type=int, default=0, help='device to use for iterate data, -1 mean cpu [default: 0]')
parser.add_argument('--no-cuda', action='store_true', default=False, help='disable the gpu')

args = parser.parse_args()
#args.cuda = (not args.no_cuda) and torch.cuda.is_available()


def tokenizer(text):
    return text.split(' ')


def main():
    # load data
    fix_length = args.fix_length

    if args.model_type != 'RNN' : #batch first && fix_length
        LABEL = data.Field(sequential=False,use_vocab=False,batch_first=True, dtype=torch.float32)
        TEXT = data.Field(sequential=True,fix_length = fix_length, tokenize=tokenizer, batch_first=True) 

    else:
        LABEL = data.Field(sequential=False, use_vocab=False,dtype=torch.float32)
        TEXT = data.Field(sequential=True,tokenize=tokenizer)

    fields = {'label': ('label', LABEL), 'text': ('text', TEXT)}
    print("ready to load data")
    train, val, test = data.TabularDataset.splits(path='newdata', train='train.json', validation='val.json',
                                                  test='test.json', format='json', fields=fields)
    train_length=len(train.examples)
    val_length=len(val.examples)
    test_length=len(test.examples)
    print('train_length:{}\nval_length:{}\ntest_length:{}'.format(train_length,val_length,test_length))
    
    print("ready to build vocab")
    # build vocab
    print('vector = ', args.vector)
    if args.vector == 'sogou':
        cache = 'word2vec/sogou_vector_cache'
        vector_path = 'word2vec/sgns.sogou.word'
    elif args.vector == 'baidu':
        cache = 'word2vec/baidu_vector_cache'
        vector_path = 'word2vec/sgns.target.word-word.dynwin5.thr10.neg5.dim300.iter5'
    elif args.vector == 'mixed':
        cache= 'word2vec/mixed_vector_cache'
        vector_path = 'word2vec/sgns.merge.word'
    #cache = '.vector_cache'
    if not os.path.exists(cache):
        print('cache not found')
        os.mkdir(cache)
    else:
        print('found cache')

    
    TEXT.build_vocab(train, min_freq=args.min_freq,vectors=Vectors(name = vector_path,cache=cache))
    #TEXT.build_vocab(train, val, test, min_freq=5,vectors=Vectors(name = '../sgns.target.word-word.dynwin5.thr10.neg5.dim300.iter5',cache=cache))
    #TEXT.build_vocab(train, val, test, min_freq=1,vectors=Vectors(name = '../sgns.merge.word',cache=cache))
    print("finish build vocab")
    vocab = TEXT.vocab
    print('vocab len = ', len(vocab))


    # 迭代器　返回batch
    device = torch.device('cuda:{}'.format(args.cuda) if args.cuda != None else 'cpu')
    train_iter, val_iter, test_iter = data.BucketIterator.splits(
    (train, val, test), batch_size=args.batch_size, device=device, shuffle=True,sort_key=lambda x: len(x.text), sort_within_batch=False)
    # 迭代器　返回batch

    # define model
    dropout = args.dropout
    model = 0
    print('model = ', args.model_type)
    if args.model_type == 'CNN':
        model = CNN(
            vocab = TEXT.vocab,
            input_size = 300,
            fix_length = fix_length,
            dropout = dropout
        )
    elif args.model_type == 'MLP':
        model = MLP(
            vocab = TEXT.vocab,
            input_size = 300,
            fix_length = fix_length,
            dropout = dropout
        )
    elif args.model_type == 'RNN':
        model = RNN(
        vocab=TEXT.vocab,
        input_size=300,
        hidden_size=300,
        num_layers=1,
        fc_size=64,
        cell = args.cell,
        dropout=dropout
        )
    elif args.model_type == 'AttRNN':
        model = AttentionRNN(
            vocab=TEXT.vocab,
            input_size=300,
            hidden_size=300,
            num_layers=1,
            fc_size=64,
            cell = args.cell,
            dropout=dropout
        )
    print(type(model))
    

    #weight = torch.Tensor([0.177625, 0.052946,0.061913,0.420154, 0.15670,0.07686, 0.042271,0.011528])

    if args.cuda != None:
        print('to cuda! device = ', device)
        model.to(device)
        #weight = torch.Tensor([0.177625, 0.052946,0.061913,0.420154, 0.15670, 0.07686, 0.042271, 0.011528]).to(device)
    print('device = ', device)
        #model.to(device)
    # define loss function and  optimizer
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)

    # train & val  & test
    train_loss_list = []
    train_acc_list = []
    val_loss_list = []
    val_acc_list = []
    if not args.evaluate:
            best_val_acc = 0.0
            test_loss, test_acc,test_corr,test_f1= evaluate(model, test_iter, criterion)
            for epoch in range(args.max_epoch):
                train_loss, train_acc = train_one_epoch(model, train_iter, criterion, optimizer, epoch, args)
                val_loss, val_acc,val_corr,val_f1= evaluate(model, val_iter, criterion)
                
                train_loss_list.append(train_loss)
                train_acc_list.append(train_acc)
                val_loss_list.append(val_loss)
                val_acc_list.append(val_acc)           
                
                if val_acc > best_val_acc:
                    best_val_acc = val_acc
                    print('ready to save data')
                    save_path = 'model_save/{}'.format(args.model_type)
                    if not os.path.exists(save_path):
                        os.mkdir(save_path)
                        print('mkdir ',save_path)
                    else:
                        print('{} has existed'.format(save_path))
                    torch.save(model.state_dict(), '{}/parameter_{}.pkl'.format(save_path, args.name))
                

                print('Epoch {} result is: train_acc={}, train_loss={}, val_acc={}, val_loss={}'.format(epoch, train_acc,
                                                                                                        train_loss, val_acc,
                                                                                                        val_loss))

            test_loss, test_acc,test_corr,testf1= evaluate(model, test_iter, criterion)
            '''
            np.save('{}_{}_train_loss.npy'.format(args.model_type, args.name), train_loss_list)
            np.save('{}_{}_train_acc.npy'.format(args.model_type, args.name), train_acc_list)
            np.save('{}_{}_val_loss.npy'.format(args.model_type, args.name), val_loss_list)
            np.save('{}_{}_val_acc.npy'.format(args.model_type, args.name), val_acc_list)
            '''
    else:
        print('evaluate mode!')
        #model.load_state_dict(torch.load('net_parameters.pkl'))
        if args.cuda == None:
            model.load_state_dict(torch.load('model_save/{}/parameter_{}.pkl'.format(args.model_type, args.name), map_location=device))  
        else:
            model.load_state_dict(torch.load('model_save/{}/parameter_{}.pkl'.format(args.model_type, args.name)))
        test_loss, test_acc,test_corr, test_f1 = evaluate(model, test_iter, criterion)
        

if __name__ == '__main__':
    main()
<file_sep>/solve.py
import torch
import torch.autograd as autograd
import torch.nn.functional as F
from scipy.stats import pearsonr
from sklearn.metrics import f1_score
import numpy as np 

def train_one_epoch(model, train_iter, criterion, optimizer, epoch, args):
    model.train()

    train_loss = 0
    train_acc = 0
    times = 0
   
    print_loss = 0
    print_acc = 0

    for batch in train_iter:
        optimizer.zero_grad()
        text, ori_label = batch.text, batch.label
        label = torch.argmax(ori_label, dim = 1)  #分类任务计算
        output = model(text)
        loss = criterion(output, label) #label用regression的形式和classification(one hot)的形式试一下吧
        #乘一个权重
        _, pred = torch.max(output.data, dim=1)
        acc = (pred == label).sum().item() / label.shape[0]

        train_loss += loss.item()
        train_acc += acc
        print_loss += loss.item()
        print_acc += acc
        times += 1

        if times % args.print_frequency == 0:
            print('Train Epoch: {}, Iteration:{}, loss: {:.6f}, acc: {:.6f}'.format(
                epoch, times, print_loss / args.print_frequency, print_acc / args.print_frequency))
            print_loss = 0
            print_acc = 0

        loss.backward()
        optimizer.step()

    train_loss /= times
    train_acc /= times
    return train_loss, train_acc


def evaluate(model, eval_iter, criterion):
    model.eval()
    test_loss = 0
    test_acc = 0
    times = 0

    f1 = 0
    corr = 0

    for batch in eval_iter:
        text, ori_label = batch.text, batch.label #归一化的label用于计算cor
        label = torch.argmax(ori_label,dim=1)
        output = model(text)
        loss = criterion(output, label)
        
        _, pred = torch.max(output.data, dim=1)
        acc = (pred == label).sum().item() / label.shape[0]  # 平均值

        temp_corr = 0
        for i in range(output.shape[0]):
            temp_corr += pearsonr(output[i].cpu().detach().numpy(), ori_label[i].cpu().detach().numpy())[0]
        temp_corr /= output.shape[0]
        corr += temp_corr

        f1 += f1_score(label.cpu().detach().numpy(), pred.cpu().detach().numpy(), np.arange(8),average='micro')

        test_loss += loss.item()
        test_acc += acc
        times += 1

    test_loss /= times
    test_acc /= times
    corr /= times
    f1 /= times
    

    print('\nValidation set: Average loss: {:.4f}, Accuracy: {:.4f} ({:.2f}%), Correlation:{:4f}, F-score:{:4f})\n'.format(
        test_loss, test_acc, 100. * test_acc, corr, f1))
    return test_loss, test_acc,corr,f1
<file_sep>/bert/myBert/myBert.py
import torch
import torch.nn as nn
import argparse
from torchtext import data
import os
import transformers
from utils_myBert import train_one_epoch, evaluate, field_load_vocab, CLS, UNK, PAD, SEP
from bert import BERT

parser = argparse.ArgumentParser()
# learning
parser.add_argument('--max_epoch', default=20, type=int, metavar='N', help='number of total epochs to run')
parser.add_argument('-b', '--batch-size', default=16, type=int, metavar='N', help='mini-batch size(default: 1')
parser.add_argument('--lr', '--learning-rate', default=0.0001, type=float, metavar='LR', help='initial learning rate',
                    dest='lr')
parser.add_argument('--momentum', default=0.9, type=float, metavar='M', help='momentum')
parser.add_argument('--wd', '--weight-decay', default=0.0005, type=float, metavar='W',
                    help='weight decay(default: 1e-4)', dest='weight_decay')
parser.add_argument('-print-frequency', type=int, default=10,
                    help='how many steps to wait before logging training status [default: 50]')
parser.add_argument('-e', '--evaluate', dest='evaluate', action='store_true', help='evaluate model on validation set')
parser.add_argument('--model', default='rnn_cache', help='model cache')
parser.add_argument('--name',default=1)
parser.add_argument('--no-cuda', action='store_true', default=False, help='disable the gpu')


args = parser.parse_args()

#从预训练词表加载tokenizer
tokenizer = transformers.BertTokenizer.from_pretrained('../chinese_wwm_ext_pytorch/vocab.txt')

def myTokenizer(text):
    t = ''.join(text.split(' '))
    t = tokenizer.tokenize(t)[:1024]#max
    #t.insert(0, CLS) #加入clk和sep
    #t.append(SEP)
    return t


def main():
    # load data
    LABEL = data.Field(sequential=False, use_vocab=False, batch_first=True, dtype=torch.float32)
    TEXT = data.Field(sequential=True, tokenize=myTokenizer, lower=False, batch_first=True, unk_token=UNK, pad_token=PAD)
    fields = {'label': ('label', LABEL), 'text': ('text', TEXT)}
    train, val, test = data.TabularDataset.splits(path='../../newdata', train='train.json', validation='val.json',
                                                  test='test.json',format='json', fields=fields)
    train_length = len(train.examples)
    val_length = len(val.examples)
    test_length = len(test.examples)
    print('train_length:{}\nval_length:{}\ntest_length:{}'.format(train_length, val_length, test_length))


    #加载与训练数据的词表
    vocab_path = '../chinese_wwm_ext_pytorch/vocab.txt'
    with open(vocab_path, 'r', encoding='utf-8') as file:
        original_vocab_size = len(file.readlines())
    field_load_vocab(TEXT, vocab_path, novector=True)
    vocab_size = len(TEXT.vocab.stoi)
    print('original vocab size:{}'.format(original_vocab_size))

    print('vocab size:{}'.format(vocab_size))
    assert vocab_size == original_vocab_size

    cls_id = TEXT.vocab.stoi[CLS]
    sep_id = TEXT.vocab.stoi[SEP]
    print('cls id = {}, sep_id = {}'.format(cls_id, sep_id)) 


    device = torch.device('cuda:{}'.format(args.cuda) if args.cuda else 'cpu')
    train_iter, val_iter, test_iter = data.BucketIterator.splits(
    (train, val, test), batch_size=args.batch_size, device=device, shuffle=True,sort_key=lambda x: len(x.text), sort_within_batch=False)
    
    # define model
    
    config = transformers.BertConfig.from_pretrained('../chinese_wwm_ext_pytorch/bert_config.json',output_attentions=False,output_hidden_states=False)
    baseBert = transformers.BertModel.from_pretrained(
        '../chinese_wwm_ext_pytorch/pytorch_model.bin',
        config=config
    )
    model = BERT(baseBert,args.batch_size, device)
    if args.cuda:
        model.to(device)

    # define loss function and  optimizer
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)

    # train & val  & test
    best_val_acc = 0.0
    print('begin to train!')
    val_loss, val_acc,val_corr,val_f1 = evaluate(model, val_iter , criterion)
    for epoch in range(args.max_epoch):
        train_loss, train_acc = train_one_epoch(model, train_iter, criterion,optimizer, epoch, args)
        val_loss, val_acc,val_corr,val_f1= evaluate(model, val_iter,criterion )
        '''
        if val_acc > best_val_acc:
            best_val_acc = val_acc
            if not os.path.exists(sample_method + '/'):
                os.makedirs(sample_method + '/')
            torch.save(model.state_dict(), '{}/net_parameters.pkl'.format(sample_method))
        '''
        print('Epoch {} result is: train_acc={}, train_loss={}, val_acc={}, val_loss={}'.format(epoch, train_acc,
                                                                                                train_loss, val_acc,
                                                                                                val_loss))

    test_loss, test_acc,test_corr, test_f1 = evaluate(model, test_iter, criterion)
    print('\nTest set: Average loss: {:.4f}, Accuracy: {:.4f} ({:.2f}%)\n'.format(test_loss, test_acc, test_acc * 100))


if __name__ == '__main__':
    main()
