FlaxAutoModelForCausalLM is the Hugging Face auto class for loading Flax causal language models and generating text. Give from_pretrained() a Model Hub name, and it infers the specialized architecture: GPT-J becomes FlaxGPTJForCausalLM, while GPT-Neo becomes FlaxGPTNeoForCausalLM. 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
Loads GPT-J-6B with Flax, JAX float16 computation, a PyTorch checkpoint conversion option, and a matching AutoTokenizer. section 11.3.1
Generation uses a tokenizer alongside the model. The tokenizer converts text into model-specific token IDs and attention masks; model.generate() then produces a continuation token by token. section 11.3.1
Keep these loading details distinct:
- revision can select an available FP16 model version. section 11.3.1
- dtype selects the computation type; it does not convert stored weights. Use to_fp16() or to_bf16() when explicit conversion is needed. section 11.3.1
