# Load model directly
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
tokenizer = AutoTokenizer.from_pretrained("pere/test-t5-small")
model = AutoModelForSeq2SeqLM.from_pretrained("pere/test-t5-small")Test T5 small conversion
This is a test repo for the conversion of T5X to HuggingFace Flax.
The current model is first converted from MTF to T5X using the conversion script included in the T5X library:
python3 -m t5x.scripts.convert_tf_checkpoint --gin_file=t5x/examples/t5/t5_1_0/small.gin --gin.convert_checkpoint.model=%MODEL --gin.convert_checkpoint.tf_checkpoint_path=\"gs://t5-data/pretrained_models/small/model.ckpt-1000000\" --gin.convert_checkpoint.output_dir=\"/tmp/t5x_checkpoints/t5_small\" --logtostderr
After creating the T5X model, the model is converted to Huggingface Flax by a modified version of the script from @stefan-it (https://gist.githubusercontent.com/stefan-it/30e4998ef159f33696e377a46f699d9f/raw/c19da5d067dc9d31d0b8115a79e8626186e11daa/convert_t5x_checkpoint_to_flax.py). The modified version is included in this repo. The modification is basically that the wi_0 and wi_1 layers are combined into wi. This might be a difference between t5_1_0 and t5_1_1
python3 convert_t5_checkpoint_to_flax.py --t5x_checkpoint_path /tmp/t5x_checkpoints/t5_small/checkpoint_1000000/ --flax_dump_folder_path /tmp/flax_dump_folder/ --config_name t5-small
The tokenizer.json was copied from https://huggingface.co/t5-small/blob/main/tokenizer.json.
To be able to use the widgets in HuggingFace, the model was converted to pyTorch by running:
from transformers import T5ForConditionalGeneration
model =
T5ForConditionalGeneration.from_pretrained(".", from_flax=True)
model.save_pretrained(".")
- Downloads last month
- 8
# Use a pipeline as a high-level helper # Warning: Pipeline type "translation" is no longer supported in transformers v5. # You must load the model directly (see below) or downgrade to v4.x with: # 'pip install "transformers<5.0.0' from transformers import pipeline pipe = pipeline("translation", model="pere/test-t5-small")