Need to load a pretrained model without memorizing its architecture? Transformers auto classes infer the appropriate implementation from the model name passed to from_pretrained(). section 11.3.1
Choose the auto class according to the task and backend:
- Use FlaxAutoModelForCausalLM for Flax GPT-like models that generate text; EleutherAI/gpt-j-6B resolves to FlaxGPTJForCausalLM. section 11.3.1
- Use AutoTokenizer to select the tokenizer and vocabulary expected by the chosen pretrained model. section 11.3.1
- Select an auto class for the relevant backend, such as PyTorch, TensorFlow, or Flax. section 11.3.1
Listing 11.13 Loading the GPT-J-6B model from the Hugging Face Model Hub
from transformers import FlaxAutoModelForCausalLM, AutoTokenizer #1
import jax
model_name = "EleutherAI/gpt-j-6B" #2
model = FlaxAutoModelForCausalLM.from_pretrained( #3
model_name,
revision="float16", #4
dtype=jax.numpy.float16, #5
from_pt=True) #6
tokenizer = AutoTokenizer.from_pretrained(model_name) #7
model.params = model.to_fp16(model.params) #8
model
>>> <transformers.models.gptj.modeling_flax_gptj.FlaxGPTJForCausalLM at
0x7c06953b3670> #9
tokenizer
>>> GPT2TokenizerFast(name_or_path='EleutherAI/gpt-j-6B',
vocab_size=50257, model_max_length=2048, is_fast=True,
padding_side='right', truncation_side='right',
special_tokens={'bos_token': AddedToken("
<|endoftext|>", rstrip=False, lstrip=False,
single_word=False, normalized=True), 'eos_token':
AddedToken("<|endoftext|>", rstrip=False,
lstrip=False, single_word=False, normalized=True),
'unk_token': AddedToken("<|endoftext|>", rstrip=False,
lstrip=False, single_word=False, normalized=True)},
clean_up_tokenization_spaces=True) #10
#1 Imports the Hugging Face transformers library
#2 Defines the name of the model we want to use
#3 The Flax auto class for causal language modeling determines and loads a specific class for the GPT-Neo model.
#4 Loads an FP16 revision of the model (not available for every model)
#5 Defines the data type of the computation; we want to use FP16.
#6 Converts model weights from PyTorch saved model (as this particular model revision does not have separate JAX weights)
#7 Automatically determines the right tokenizer for the model
#8 Converts the model weights to FP16 (not needed here, just as an example for other models with no FP16 revision available)
#9 Autoclass FlaxAutoModelForCausalLM returns the specialized FlaxGPTNeoForCausalLM class.
#10 The auto tokenizer class returns the specialized GPT2TokenizerFast class
A compact listing shows model and tokenizer loading together, including the model name, Flax class, FP16 settings, and inferred tokenizer. section 11.3.1
