利用BERT得到句子的表示向量（pytorch）

在文本分类和文本相似度匹配中，经常用预训练语言模型BERT来得到句子的表示向量，下面给出了pytorch环境下的操作的方法：

• 这里使用huggingface的transformers中BERT， 需要先安装该依赖包（pip install transformers）
• 具体实现如下：

import torch
from tqdm import tqdm
import joblib
import numpy as np
from torch.utils.data import DataLoader,Dataset
from sklearn.datasets import fetch_20newsgroups
from transformers import BertTokenizer,BertModel

class NewDataset(Dataset):
    def __init__(self, bert_train, mask_train=None, seg_ids_train=None):
        self.bert_train = bert_train
        self.mask_train = mask_train
        self.seg_ids_train = seg_ids_train
    def __getitem__(self, i):
        return torch.LongTensor(self.bert_train[i]), 
               torch.LongTensor(self.mask_train[i]), 
               torch.LongTensor(self.seg_ids_train[i])

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


newsgroups_train = fetch_20newsgroups(subset='train').data
newsgroups_test = fetch_20newsgroups(subset='test').data
train_label = fetch_20newsgroups(subset='train').target
test_label = fetch_20newsgroups(subset='test').target

L=512
N = len(newsgroups_train)
bert_train,mask_train,seg_ids_train = [], [],[]
all_sents = newsgroups_train+newsgroups_test
tokenizer=BertTokenizer.from_pretrained('bert-base-uncased')
for sent in tqdm(all_sents):
    tokens = tokenizer.tokenize(sent)
    tokens = ['[CLS]'] + tokens + ['[SEP]']
    padded_tokens = tokens[:L] + ['[PAD]' for _ in range(L - len(tokens))]
    attn_mask = [1 if token != '[PAD]' else 0 for token in padded_tokens]
    sent_ids = tokenizer.convert_tokens_to_ids(padded_tokens)
    seg_ids = [0 for _ in range(len(padded_tokens))]
    bert_train.append(sent_ids)
    mask_train.append(attn_mask)
    seg_ids_train.append(seg_ids)

torch.device("cuda" if torch.cuda.is_available() else "cpu")
device = "cuda:0"
data = NewDataset(bert_train,mask_train=mask_train,seg_ids_train=seg_ids_train)
bert_model = BertModel.from_pretrained('bert-base-uncased').to(device)

reps = []
batchsize = 5
for batch in tqdm(DataLoader(data, shuffle=False, batch_size=batchsize)):
    bert_train, mask_train, seg_ids_train = batch
    hidden_reps, cls_head = bert_model(bert_train.cuda(), attention_mask=mask_train.cuda(), token_type_ids=seg_ids_train.cuda())
    reps+=list(cls_head.detach().cpu().numpy())

if len(reps) != len(all_sents):
    assert "no equal size"

reps_train = reps[:N]
reps_test = reps[-N:]

newsgroups_data = {'train_vecs': reps_train, 'train_label': train_label, 'test_vecs': reps_test,'test_label': test_label}
joblib.dump(newsgroups_data,"newsgroups_data.pkl")