在文本分类和文本相似度匹配中,经常用预训练语言模型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")
本站文章如无特殊说明,均为本站原创,如若转载,请注明出处:利用BERT得到句子的表示向量(pytorch) - Python技术站