Flan-T5 is an enhanced version of Google’s T5 AI model which is quite good at certain language tasks.
For example, it’s supposed to be better at a lot of zero-shot examples even than GPT-3.
Table of Contents
Install and Setup Flan-T5
First, we install the transformers module by running the below command
pip install transformers
If you run into an issue, check out the Transformers Installation Guide.
Then, we load the Flan-T5 model
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-small")
Load the tokenizer
tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small")
Note we loaded the small model. You can load any of the versions listed in a later section below.
The bigger the model, the better it is but also more resource consuming
Using Flan-T5 for language AI tasks
Next, we pass the prompt we want the AI model to generate text for.
inputs = tokenizer("A intro paragraph on a article on space travel:", return_tensors="pt")
We call the model’s generate function and get the response.
outputs = model.generate(**inputs)
You can then print the output from the AI model
print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
Flan-T5 versions
| Model | Gin File Location | Checkpoint Location |
|---|---|---|
| Flan-T5 Small | t5_1_1/small.gin | gs://t5-data/pretrained_models/t5x/flan_t5_small/checkpoint_1198000 |
| Flan-T5 Base | t5_1_1/base.gin | gs://t5-data/pretrained_models/t5x/flan_t5_base/checkpoint_1184000 |
| Flan-T5 Large | t5_1_1_large.gin | gs://t5-data/pretrained_models/t5x/flan_t5_large/checkpoint_1164000 |
| Flan-T5 XL | t5_1_1_xl.gin | gs://t5-data/pretrained_models/t5x/flan_t5_xl/checkpoint_1138000 |
| Flan-T5 XXL | t5_1_1_xxl.gin | gs://t5-data/pretrained_models/t5x/flan_t5_xxl/checkpoint_1114000 |