pytorch中的包的用法_Python pytorch-transformers包_程序模块 - PyPI - Python中文网

importtorchfrompytorch_transformersimport*# PyTorch-Transformers has a unified API# for 6 transformer architectures and 27 pretrained weights.# Model | Tokenizer | Pretrained weights shortcutMODELS=[(BertModel,BertTokenizer,'bert-base-uncased'),(OpenAIGPTModel,OpenAIGPTTokenizer,'openai-gpt'),(GPT2Model,GPT2Tokenizer,'gpt2'),(TransfoXLModel,TransfoXLTokenizer,'transfo-xl-wt103'),(XLNetModel,XLNetTokenizer,'xlnet-base-cased'),(XLMModel,XLMTokenizer,'xlm-mlm-enfr-1024')]# Let's encode some text in a sequence of hidden-states using each model:formodel_class,tokenizer_class,pretrained_weightsinMODELS:# Load pretrained model/tokenizertokenizer=tokenizer_class.from_pretrained(pretrained_weights)model=model_class.from_pretrained(pretrained_weights)# Encode textinput_ids=torch.tensor([tokenizer.encode("Here is some text to encode")])withtorch.no_grad():last_hidden_states=model(input_ids)[0]# Models outputs are now tuples# Each architecture is provided with several class for fine-tuning on down-stream tasks, e.g.BERT_MODEL_CLASSES=[BertModel,BertForPreTraining,BertForMaskedLM,BertForNextSentencePrediction,BertForSequenceClassification,BertForMultipleChoice,BertForTokenClassification,BertForQuestionAnswering]# All the classes for an architecture can be initiated from pretrained weights for this architecture# Note that additional weights added for fine-tuning are only initialized# and need to be trained on the down-stream tasktokenizer=BertTokenizer.from_pretrained('bert-base-uncased')formodel_classinBERT_MODEL_CLASSES:# Load pretrained model/tokenizermodel=model_class.from_pretrained('bert-base-uncased')# Models can return full list of hidden-states & attentions weights at each layermodel=model_class.from_pretrained(pretrained_weights,output_hidden_states=True,output_attentions=True)input_ids=torch.tensor([tokenizer.encode("Let's see all hidden-states and attentions on this text")])all_hidden_states,all_attentions=model(input_ids)[-2:]# Models are compatible with Torchscriptmodel=model_class.from_pretrained(pretrained_weights,torchscript=True)traced_model=torch.jit.trace(model,(input_ids,))# Simple serialization for models and tokenizersmodel.save_pretrained('./directory/to/save/')# savemodel=model_class.from_pretrained('./directory/to/save/')# re-loadtokenizer.save_pretrained('./directory/to/save/')# savetokenizer=tokenizer_class.from_pretrained('./directory/to/save/')# re-load# SOTA examples for GLUE, SQUAD, text generation...