GPT-J-6B is an open-source, GPT-like autoregressive language model with 6 billion parameters. EleutherAI trained it on the Pile, an 825 GB collection of 22 open language datasets. It serves as a practical JAX example for loading a pretrained model and generating text. section 11.3.1
Keep these practical constraints in mind:
- Load the model and tokenizer through Hugging Face Transformers using Flax APIs, float16 parameters, and JAX arrays. section 11.3.1
- Use sufficient GPU or TPU memory; choose a smaller model such as Gemma 2B if GPT-J-6B is too large for the machine. section 11.3.1
- Do not deploy it for human-facing interaction without fine-tuning, supervision, and/or moderation because it may produce harmful or offensive text. section 11.3.1
Generate a SQL query
Tokenize a prompt asking for a query that selects users older than 20, then call model.generate with sampling, one beam, a 300-token limit, and temperature 0.7. The tokenizer returns JAX input_ids and an attention_mask for generation. 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
Inspect the complete FlaxAutoModelForCausalLM and AutoTokenizer loading example, including the float16 revision and parameter conversion. section 11.3.1
Listing 11.14 Using the GPT-J-6B model for text generation
prompt = """Generate SQL query from the text inside triple backticks
```select all the users from table 'users' older than 20 years```
""" #1
inputs = tokenizer(prompt, return_tensors="jax") #2
inputs
>>> {'input_ids': Array([[ 8645, 378, 16363, 12405, 422, 262, 2420,
...
>>> 812, 15506, 63, 198]], dtype=int32),
>>> 'attention_mask': Array([[1, 1, 1, … 1, 1, 1]], dtype=int32)}
generated_ids = model.generate( #3
**inputs,
do_sample=True, #4
num_beams=1, #4
max_new_tokens=300, #5
temperature=0.7, #6
pad_token_id = model.config.eos_token_id,
prng_key=jax.random.PRNGKey(4232), #7
no_repeat_ngram_size=2)
generated_text = tokenizer.decode(
generated_ids['sequences'].squeeze(0)) #8
print(generated_text)
>>> Generate SQL query from the text inside triple backticks #9
>>> ```select all the users from table 'users' older than 20 years```
>>>
>>> ```
>>> SELECT * FROM `users` WHERE age > 20
>>> ```
>>>
>>> ```
>>> SELECT * FROM `users` WHERE age > `20 years`
>>> ```
>>>
>>> ```
>>> SELECT * FROM `users` WHERE `age` > 20`
>>> ```
>>>
>>> <!--end-query-->
>>> …
Inspect the text-generation listing, including JAX tokenization, sampling settings, beam count, token limit, and temperature. section 11.3.1
