pytorch+huggingface+bert实现一个文本分类

pytorch+huggingface+bert实现一个文本分类

1，下载模型

bert模型的目前方便的有两种：一种是huggingface_hub以及/AutoModel.load, 一种torch.hub。

(1) 使用huggingface_hub或者下载模型。

# repo_id:模型名称or用户名/模型名称

from huggingface_hub import hf_hub_download
from transformers import AutoModel, AutoTokenizer
hf_hub_download(repo_id='bert-base-chinese', filename="pytorch_model.bin", local_dir='./bert-base-chinese') # 共下载五个:pytorch_model.bin, config.json,vocab.txt, tokenizer_config.json, tokenizer.json


# 直接下载/从~/.cache/huggingface/hub或者从指定文件夹
model = AutoModel.from_pretrain('bert-base-chinese')
tokenizer = AutoTokenizer.from_pretrain('bert-base-chinese')

(2) 使用torch.hub下载模型

# repo_or_dir:例如hugggingface/pytorch-transformers或者pytorch/vision， 或者本地目录
# source:github/local,默认是="github"
model = torch.hub.load('huggignface/pytorch-transformers', 'model', 'bert-base-chinese',source='github')
tokenizer = torch.hub.load('huggingface/pytorch-transformers', 'tokenizer', 'bert-base-chinese', source='github')

2， 使用模型

tokenizer(text)：返回字典，包含：input_ids, token_type_ids,attention_mask。

tokenizer.encode(text)：只返回input_dis。

两种方式都限定了text的方式，要么是单句，要么是列表。返回的结果也只有两种:seq_length, batch_size, seq_length。

注意一，如果设定return_tensors='pt'，那么只会返回一种batch_size, seq-length。预训练模型只接受batch_size, seq_length形式输入。

注意二，bert模型返回两个结果，第一个是所有的最后一层hidden, 第二个是pooler_output，也就是所有最后一层hidden的加和平均值。models默认返回dict结果，使用mdoel().keys()返回变量值。使用values()返回结果。

from transformers import AutoModel, AutoTokenizer
import torch

model = AutoModel.from_pretrained('./bert-base-chinese')
tokenizer = AutoTokenizer.from_pretrained('./bert-base-chinese')

def bert_encode(line):
    tokens = tokenizer(line, add_special_tokens=False, return_tensors='pt', padding=True, truncation=True, max_length=20)
    with torch.no_grad():
        last_hiddens, pooler_hidden = model(**tokens).values()
    return last_hiddens, pooler_hidden

3，利用last_hidden通过rnn+dense进行文本分类。

rnn模型的本质只需要注意输入的形状：input:seq_length, batch_size, feature_size ,hidden_size:num_layers, batch_size, hidden_size。batch_size,都是放在中间，符合rnn处理的特点。

其次注意一点，rnn返回两个个向量:

一个是所有token的最后一层hidden, seq_length, batch_size, hidden_size

一个是最后一个token的num_layers, batch_size, hidden_size

# 定义一个rnn模型
import torch
import torch.nn as nn
import torch.nn.functional as f

class myRNN(nn.Module):
    def __init__(self, input_size, output_size, hidden_size, num_layers):
        super().__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.rnn = nn.RNN(input_size, hidden_size, num_layers=num_layers)
        self.dense = nn.Linear(hidden_size * num_layers, output_size)

    def forward(self, input_tokens):
        """
        :param input_tokens: size()=[batch_size, seq_legnth, input_size]
        :return: [batch_size, output_size]
        """
        batch_size = input_tokens.size()[0]
        hidden0 = self.init_hidden(batch_size)
        input_tokens = input_tokens.transpose(0, 1)
        all_hiddens, last_hiddens = self.rnn(input_tokens, hidden0) # last_hiddensize.size() = [num_layers, batch_szie, hidden_size]
        last_hiddens = last_hiddens.transpose(0,1).reshape((batch_size, -1))
        last_hiddens = f.relu(last_hiddens)
        return self.dense(last_hiddens)


    def init_hidden(self, batch_size):
        return torch.zeros(size=(self.num_layers, batch_size, self.hidden_size))

4， 组装模型。

import torch.nn as nn
import torch

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.bert = bert_encode
        self.rnn = myRNN(768,2, 768, 2)

    def forward(self, lines):
        tokens = self.bert(lines)[0]
        return self.rnn(tokens)

5, 编写训练脚本。

训练脚本具有通用性，参见训练脚本介绍。

from tqdm import tqdm
import torch
import time
import pickle

def train(epoches, model, dataloader, optim, criterion, model_save_path):
    loss_mark = float('inf')
    loss_store = [] # 收集每一轮的loss
    acc_store = [] # 收集每一轮的acc率
    model_num = 0
    for epoch in tqdm(range(epoches)):
        loss_sum = 0
        acc_sum = 0
        start = time.time()
        for batch_num , (batch_line, batch_category) in tqdm(enumerate(dataloader)):
            batch_pred = model(batch_line)
            acc = (batch_pred.argmax(-1) == batch_category).sum().item()
            acc_sum += acc
            loss = criterion(batch_pred, batch_category)
            loss_sum += loss.item()
            optim.zero_grad()
            loss.backward()
            optim.step()
            if batch_num % 100 == 0:
                print(f"batch_num: {batch_num} | timesince: {time.time() - start}")
                start = time.time()
                print(f"loss: {loss} | acc: {acc/len(batch_category)}")
        print(f"epoch {epoch} | loss: {loss_sum / len(dataloader)} | acc: {acc_sum}")
        loss_store.append(loss_sum/len(dataloader))
        acc_store.append(acc_sum)
        if loss_sum  < loss_mark:
            loss_store = loss_sum
            torch.save(f'{model_save_path}/{model_num}.pth', pickle_module=pickle, pickle_protocol=2)
            model_num += 1
        return loss_store, acc_store

6， 编写测试脚本

def valid(model, dataloader, criterion):
    loss_sum = 0
    with torch.no_grad():
        for batch_line, batch_category in dataloader:
            batch_pred = model(batch_line)
            loss = criterion(batch_pred, batch_category)
            loss_sum += loss.item()
    return loss_sum / len(dataloader)

7, 读取数据，进行训练。

主要是dataset的构建。

import pandas as pd
from torch.utils.data import Dataset, DataLoader

class MyDataset(Dataset):
    def __init__(self, lines, category):
        super().__init__()
        self.lines = lines
        self.category = category

    def __getitem__(self, item):
        return self.lines[item], self.category[item]

    def __len__(self):
        return len(self.lines)

train_data = pd.read_csv('./train_data.csv', header=None, sep='\t')
category, lines = train_data[0].tolist(), train_data[1].tolist()
datasets = MyDataset(lines, category)
dataloader = DataLoader(datasets, batch_size=10, shuffle=True)

model = Model()
optim = Adam(model.parameters(), lr=1e-3)
criterion = CrossEntropyLoss()

train(epoches=1, model=model, dataloader=dataloader, optim=optim, criterion=criterion, model_save_path='./models')
	

8，读取模型进行预测。

读取模型时注意，如果保存的是模型全模型，那么在当前文件内需要有模型的定义。如果是参数，那么需要实体化一类模型类，使用load_state_dict()方法。

import torch
import pickle
from bert_encode import bert_encode
from MyRnn import myRNN
import torch.nn as nn

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.bert = bert_encode
        self.rnn = myRNN(768, 2, 768, 2)

    def forward(self, lines):
        tokens = self.bert(lines)[0]
        return self.rnn(tokens)

model = torch.load('./models/1.pth', pickle_module=pickle)

def predict(batch_lines, model):
    batch_output = model(batch_lines)
    return batch_output.argmax(-1)