89 lines
3.6 KiB
Python
89 lines
3.6 KiB
Python
|
#https://github.com/huggingface/transformers
|
|||
|
|
|
|||
|
|
#https://huggingface.co/transformers/quickstart.html
|
|||
|
|
#BERT example
|
|||
|
|
|
|||
|
|
#pip install transformers
|
|||
|
|
#老的pytorch_transformers
|
|||
|
|
|
|||
|
|
import torch
|
|||
|
|
import torch.nn as nn
|
|||
|
|
from transformers import BertConfig, BertModel
|
|||
|
|
from transformers.tokenization_bert import BertTokenizer as tokenization
|
|||
|
|
import os
|
|||
|
|
|
|||
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|||
|
|
#get_bert_model
|
|||
|
|
# bert预训练模型：
|
|||
|
|
# pytorch_model.bin
|
|||
|
|
# config.json
|
|||
|
|
# vocab.txt
|
|||
|
|
bert_path = './bert'
|
|||
|
|
do_lower_case=True
|
|||
|
|
|
|||
|
|
bert_config_file = os.path.join(bert_path, f'bert_config.json')
|
|||
|
|
vocab_file = os.path.join(bert_path, f'vocab.txt')
|
|||
|
|
init_checkpoint = os.path.join(bert_path, f'pytorch_model.bin')
|
|||
|
|
|
|||
|
|
#加载配置
|
|||
|
|
bert_config = BertConfig.from_json_file(bert_config_file)
|
|||
|
|
|
|||
|
|
# 加载词典
|
|||
|
|
tokenizer = tokenization(vocab_file=vocab_file, do_lower_case=do_lower_case)
|
|||
|
|
|
|||
|
|
# 加载模型
|
|||
|
|
model_bert = BertModel.from_pretrained(bert_path)
|
|||
|
|
model_bert.to(device)
|
|||
|
|
|
|||
|
|
|
|||
|
|
# Tokenize input
|
|||
|
|
text = "乌兹别克斯坦议会立法院主席获连任"
|
|||
|
|
tokenized_text = tokenizer.tokenize(text)
|
|||
|
|
tokenized_text=['[CLS]'] + tokenized_text + ['[SEP]']
|
|||
|
|
# Convert token to vocabulary indices
|
|||
|
|
# input_ids：一个形状为[batch_size, sequence_length]的torch.LongTensor，在词汇表中包含单词的token索引
|
|||
|
|
input_ids = tokenizer.convert_tokens_to_ids(tokenized_text)
|
|||
|
|
# Define sentence A and B indices associated to 1st and 2nd sentences (see paper)
|
|||
|
|
# segment_ids ：形状[batch_size, sequence_length]的可选torch.LongTensor，在[0, 1]中选择token类型索引。类型0对应于句子A，类型1对应于句子B。
|
|||
|
|
segment_ids = [0]*len(input_ids)
|
|||
|
|
# input_mask：一个可选的torch.LongTensor，形状为[batch_size, sequence_length]，索引在[0, 1]中选择。
|
|||
|
|
input_mask = [1]*len(input_ids)
|
|||
|
|
|
|||
|
|
# Convert inputs to PyTorch tensors
|
|||
|
|
input_ids = torch.tensor([input_ids], dtype=torch.long).to(device)
|
|||
|
|
print("input_ids",input_ids.size())
|
|||
|
|
input_mask = torch.tensor([input_mask], dtype=torch.long).to(device)# attention_mask,可以不输入
|
|||
|
|
segments_tensors = torch.tensor([segment_ids], dtype=torch.long).to(device)
|
|||
|
|
#输出
|
|||
|
|
all_encoder_layer, pooled_output = model_bert(input_ids,input_mask,token_type_ids=segments_tensors)
|
|||
|
|
# all_encoder_layers：一个大小为[batch_size, sequence_length，hidden_size]的torch.FloatTensor列表，
|
|||
|
|
# 它是每个注意块末端隐藏状态的完整序列列表（即BERT-base的12个完整序列，BERT-large的24个完整序列）
|
|||
|
|
# pooled_output：一个大小为[batch_size, hidden_size] 的torch.FloatTensor，
|
|||
|
|
# 它是在与输入（CLF）的第一个字符相关联的隐藏状态之上预训练的分类器的输出，用于训练Next - Sentence任务（参见BERT的论文）
|
|||
|
|
|
|||
|
|
#如果我们要输出embeding 表示，只使用all_encoder_layer
|
|||
|
|
print('all_encoder_layer',all_encoder_layer.shape)
|
|||
|
|
print('pooled_output',pooled_output.size())
|
|||
|
|
#如果要分类，使用pooled_output
|
|||
|
|
|
|||
|
|
#padding
|
|||
|
|
max_seq_length=300
|
|||
|
|
|
|||
|
|
text = "乌兹别克斯坦议会立法院主席获连任"
|
|||
|
|
tokenized_text = tokenizer.tokenize(text)
|
|||
|
|
tokenized_text=['[CLS]'] + tokenized_text + ['[SEP]']
|
|||
|
|
input_ids = tokenizer.convert_tokens_to_ids(tokenized_text)
|
|||
|
|
input_mask = [1]*len(input_ids)
|
|||
|
|
|
|||
|
|
padding = [0] * (max_seq_length - len(input_ids))
|
|||
|
|
input_ids += padding
|
|||
|
|
input_mask += padding
|
|||
|
|
input_ids = torch.tensor([input_ids], dtype=torch.long).to(device)
|
|||
|
|
input_mask = torch.tensor([input_mask], dtype=torch.long).to(device)
|
|||
|
|
print("padding input_ids",input_ids.size())
|
|||
|
|
|
|||
|
|
model_bert.eval()
|
|||
|
|
with torch.no_grad():
|
|||
|
|
all_encoder_layer, pooled_output = model_bert(input_ids,attention_mask= input_mask)
|
|||
|
|
print('padding all_encoder_layer', all_encoder_layer.shape)
|
|||
|
|
print('padding pooled_output', pooled_output.size())
|