Duplicate from parler-tts/parler-tts-mini-expresso

Co-authored-by: Sanchit Gandhi <sanchit-gandhi@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,363 @@
+---
+library_name: transformers
+tags:
+- text-to-speech
+- annotation
+license: apache-2.0
+language:
+- en
+pipeline_tag: text-to-speech
+inference: false
+datasets:
+  - ylacombe/expresso
+  - reach-vb/jenny_tts_dataset
+  - blabble-io/libritts_r
+---
+
+<img src="https://huggingface.co/datasets/parler-tts/images/resolve/main/thumbnail.png" alt="Parler Logo" width="800" style="margin-left:'auto' margin-right:'auto' display:'block'"/>
+
+
+# Parler-TTS Mini: Expresso
+
+<a target="_blank" href="https://huggingface.co/spaces/parler-tts/parler-tts-expresso">
+  <img src="https://huggingface.co/datasets/huggingface/badges/raw/main/open-in-hf-spaces-sm.svg" alt="Open in HuggingFace"/>
+</a>
+
+**Parler-TTS Mini: Expresso** is a fine-tuned version of [Parler-TTS Mini v0.1](https://huggingface.co/parler-tts/parler_tts_mini_v0.1)
+on the [Expresso](https://huggingface.co/datasets/ylacombe/expresso) dataset. It is a lightweight text-to-speech (TTS) 
+model that can generate high-quality, natural sounding speech. Compared to the original model, Parler-TTS Expresso provides
+superior control over **emotions** (happy, confused, laughing, sad) and **consistent voices** (Jerry, Thomas, Elisabeth, Talia).
+
+It is part of the first release from the [Parler-TTS](https://github.com/huggingface/parler-tts) project, which aims to 
+provide the community with TTS training resources and dataset pre-processing code. Details for reproducing this entire 
+training run are provided in the section [Training Procedure](#training-procedure).
+
+## Usage
+
+Using Expresso is as simple as "bonjour". Simply install the library from source:
+
+```sh
+pip install git+https://github.com/huggingface/parler-tts.git
+```
+
+You can then use the model with the following inference snippet:
+
+```py
+import torch
+from parler_tts import ParlerTTSForConditionalGeneration
+from transformers import AutoTokenizer, set_seed
+import soundfile as sf
+
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
+
+model = ParlerTTSForConditionalGeneration.from_pretrained("parler-tts/parler-tts-mini-expresso").to(device)
+tokenizer = AutoTokenizer.from_pretrained("parler-tts/parler-tts-mini-expresso")
+
+prompt = "Why do you make me do these examples? They're *so* generic."
+description = "Thomas speaks moderately slowly in a sad tone with emphasis and high quality audio."
+
+input_ids = tokenizer(description, return_tensors="pt").input_ids.to(device)
+prompt_input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
+
+set_seed(42)
+generation = model.generate(input_ids=input_ids, prompt_input_ids=prompt_input_ids)
+audio_arr = generation.cpu().numpy().squeeze()
+sf.write("parler_tts_out.wav", audio_arr, model.config.sampling_rate)
+```
+
+**Tips**:
+* Specify the name of a male speaker (Jerry, Thomas) or female speaker (Talia, Elisabeth) for consistent voices
+* The model can generate in a range of emotions, including: "happy", "confused", "default" (meaning no particular emotion conveyed), "laughing", "sad", "whisper", "emphasis"
+* Include the term "high quality audio" to generate the highest quality audio, and "very noisy audio" for high levels of background noise
+* Punctuation can be used to control the prosody of the generations, e.g. use commas to add small breaks in speech
+* To emphasise particular words, wrap them in asterisk (e.g. `*you*` in the example above) and include "emphasis" in the prompt
+
+## Training Procedure
+
+Expresso is a high-quality, expressive speech dataset that includes samples from four speakers (two male, two female).
+By fine-tuning Parler-TTS Mini v0.1 on this dataset, we can train the model to follow emotion and speaker prompts.
+
+To reproduce this fine-tuning run, we need to perform two steps:
+1. Create text descriptions from the audio samples in the Expresso dataset
+2. Train the model on the (text, audio) pairs
+
+Step 1 is performed using the [DataSpeech](https://github.com/huggingface/dataspeech) library, and step 2 using 
+[Parler-TTS](https://github.com/huggingface/parler-tts). Should you wish to use the pre-annotated dataset from our 
+experiments, you can jump straight to [step 2](#step-2--fine-tune-the-model). For both, you can follow step 0 for 
+getting set-up.
+
+### Step 0: Set-Up
+
+We'll start by creating a fresh Python environment:
+
+```sh
+python3 -m venv parler-env
+source parler-env/bin/activate
+```
+
+Next, install PyTorch according to the [official instructions](https://pytorch.org/get-started/locally/). We can then 
+install DataSpeech and Parler-TTS sequentially:
+
+```sh
+git clone git@github.com:huggingface/dataspeech.git && cd dataspeech && pip install -r requirements.txt
+cd ..
+git clone https://github.com/huggingface/parler-tts.git && cd parler-tts && pip install -e ."[train]"
+cd ..
+```
+
+You can link your Hugging Face account so that you can push model repositories on the Hub. This will allow you to save 
+your trained models on the Hub so that you can share them with the community. Simply run the command:
+
+```sh
+git config --global credential.helper store
+huggingface-cli login
+```
+
+And then enter an authentication token from https://huggingface.co/settings/tokens. Create a new token if you do not 
+have one already. You should make sure that this token has "write" privileges.
+
+You also have the option to configure Accelerate by running the following command. Note that you should set the number 
+of GPUs you wish to use for training/inference, and also the data type (dtype) based on your device (e.g. bfloat16 on 
+A100 GPUs, float16 on V100 GPUs, etc.):
+
+```sh
+accelerate config
+```
+
+Optionally, you can also login to Weights and Biases for automatic logging:
+
+```sh
+wandb login
+```
+
+### Step 1: Create Text Descriptions
+
+Creating text descriptions for the dataset comprises three sub-stages from DataSpeech, which we'll cover below.
+
+#### 1.A. Annotate the Expresso dataset
+
+We'll use the [`main.py`](https://github.com/huggingface/dataspeech/blob/main/main.py) file from DataSpeech to label 
+the following continuous variables:
+- Speaking rate
+- Signal-to-noise ratio (SNR)
+- Reverberation
+- Speech monotony
+
+This can be done with the following command:
+```sh
+python ./dataspeech/main.py "ylacombe/expresso" \
+  --configuration "default" \
+  --text_column_name "text" \
+  --audio_column_name "audio" \
+  --cpu_num_workers 8 \
+  --rename_column \
+  --repo_id "expresso-tags"
+```
+
+Note that the script will be faster if you have GPUs at your disposal. It will automatically scale up to every GPU available in your environment. To control which GPUs to run the script on consider indicating via `CUDA_VISIBLE_DEVICES` environment variable.
+
+The resulting dataset will be pushed to the Hugging Face Hub under your Hugging Face handle. Mine was pushed to [reach-vb/expresso-tags](https://huggingface.co/datasets/reach-vb/expresso-tags).
+We can see that the dataset is annotated with continuous features like "speaking_rate" and "snr".
+
+#### 1.B. Map annotations to text bins
+
+The next step involves mapping the continuous variables to discrete ones. This is achieved by binning the continuous variables
+into buckets, and assigning each one a text label.
+
+Since the ultimate goal here is to fine-tune the [Parler-TTS v0.1 checkpoint](https://huggingface.co/parler-tts/parler_tts_mini_v0.1) 
+on the Expresso dataset, we want to stay consistent with the text bins of the dataset on which the original model was trained.
+
+To do this, we'll pass [`v01_bin_edges.json`](https://github.com/huggingface/dataspeech/blob/main/examples/tags_to_annotations/v01_bin_edges.json)
+as an input argument to our script, which holds the bin edges from the original dataset:
+
+```sh
+python ./dataspeech/scripts/metadata_to_text.py \
+    "reach-vb/expresso-tags" \
+    --repo_id "expresso-tags" \
+    --configuration "default" \
+    --cpu_num_workers "8" \
+    --path_to_bin_edges "./examples/tags_to_annotations/v01_bin_edges.json" \
+    --avoid_pitch_computation
+```
+
+Since we leverage the bins from the original dataset, the above script only takes a few seconds. The resulting dataset 
+will be pushed to the Hugging Face Hub under your Hugging Face handle. Mine was pushed to [reach-vb/expresso-tags](https://huggingface.co/datasets/reach-vb/expresso-tags).
+
+You can notice that text bins such as "slightly noisy", "quite monotone" have been added to the samples.
+
+#### 1.C. Create natural language descriptions from those text bins
+
+Now that we have text bins associated to the Expresso dataset, the next step is to create natural language descriptions. 
+This involves passing the text bins to a large-language model (LLM), and have it generate corresponding descriptions.
+
+There is a template [prompt creation script](https://github.com/huggingface/dataspeech/blob/main/scripts/run_prompt_creation.py) 
+in Parler-TTS which can be used to generate descriptions from the features tagged in [step 1.A](#1a-annotate-the-expresso-dataset) (reverberation, noise, speaking rate, etc).
+
+However, not all of these features are relevant for the Expresso dataset. For instance, Expresso was recorded in a 
+professional recording studio, so all the samples are high quality. Thus, we chose to create prompts with the following subset of features:
+1. Name: we mapped the speaker ids (ex1, ex2, ex3, ex4) to unique speaker names (Jerry, Elisabeth, Thomas, Talia). This encourages the model to learn specific speakers from the training data
+2. Emotion: we include the emotion provided in the Expresso dataset
+3. Speaking rate: we use the pre-computed text bins from the previous step
+
+4. In addition, we also hard-coded the quality of the audio to be "very high-quality", given the studio recording conditions.
+
+As an example, if we passed:
+1. Speaker: Jerry
+2. Emotion: confused
+3. Speaking rate: moderate speed
+
+We would expect to generate a sample along the lines of: "Jerry speaks with a confused tone and at a moderate speed with high quality audio."
+
+The modified prompt creation script can be found in this repository. You can download this script with the following Python command:
+
+```python
+from huggingface_hub import hf_hub_download
+
+hf_hub_download(repo_id="parler-tts/parler-tts-mini-expresso", filename="run_prompt_creation.py", local_dir="./run_prompt_creation_expresso.py")
+```
+
+You can then launch prompt creation using the [Mistral Instruct 7B](https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2) 
+model with the following command:  
+
+```sh
+accelerate launch ./dataspeech/run_prompt_creation_expresso.py \
+  --dataset_name "reach-vb/expresso-tags" \
+  --dataset_config_name "default" \
+  --model_name_or_path "mistralai/Mistral-7B-Instruct-v0.2" \
+  --per_device_eval_batch_size 32 \
+  --attn_implementation "sdpa" \
+  --dataloader_num_workers 8 \
+  --output_dir "./tmp_expresso" \
+  --load_in_4bit \
+  --push_to_hub \
+  --hub_dataset_id "expresso-tagged-w-speech-mistral" \
+  --preprocessing_num_workers 16
+```
+
+Note that the Mistral model is gated, so you should ensure you have accepted the terms-of-use from the [model card](https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2).
+You can find the annotated dataset under TODO [reach-vb/expresso-tagged-w-speech-mistral](https://huggingface.co/datasets/reach-vb/expresso-tagged-w-speech-mistral),
+where you'll find sensible descriptions from the features that we passed.
+
+This step generally demands more resources and times and should use one or many GPUs. Scaling to multiple GPUs using [distributed data parallelism (DDP)](https://pytorch.org/tutorials/beginner/ddp_series_theory.html)
+is trivial: simply run `accelerate config` and select the multi-GPU option, specifying the IDs of the GPUs you wish to use. The 
+above script can then be run using DDP with no code changes.
+
+If you are resource constrained and need to use a smaller model, [Gemma 2B](https://huggingface.co/google/gemma-2b-it) 
+is an excellent choice.
+
+### Step 2: Fine-Tune the Model
+
+Fine-tuning is performed using the Parler-TTS training script [run_parler_tts_training.py](https://github.com/huggingface/parler-tts/blob/main/training/run_parler_tts_training.py).
+It is the same script used to pre-train the model, and can be used for fine-tuning without any code-changes.
+
+To preserve the model's ability to generate speech with generic voice descriptions, such as in the style of 
+[Parler-TTS Mini v0.1](https://huggingface.co/parler-tts/parler_tts_mini_v0.1), we fine-tuned the model 
+on a combination of three datasets, including the test split of LibriTTS-R:
+1. [Expresso](https://huggingface.co/datasets/ylacombe/expresso)
+2. [Jenny](https://huggingface.co/datasets/reach-vb/jenny_tts_dataset)
+3. [LibriTTS-R](https://huggingface.co/datasets/blabble-io/libritts_r)
+
+This was achieved through the following command:
+
+```sh
+accelerate launch ./training/run_parler_tts_training.py \
+    --model_name_or_path "parler-tts/parler_tts_mini_v0.1" \
+    --feature_extractor_name "parler-tts/dac_44khZ_8kbps" \
+    --description_tokenizer_name "parler-tts/parler_tts_mini_v0.1" \
+    --prompt_tokenizer_name "parler-tts/parler_tts_mini_v0.1" \
+    --report_to "wandb" \
+    --overwrite_output_dir true \
+    --train_dataset_name "ylacombe/expresso+reach-vb/jenny_tts_dataset+blabble-io/libritts_r+blabble-io/libritts_r" \
+    --train_metadata_dataset_name "reach-vb/expresso-tagged-w-speech-mistral-v3+ylacombe/jenny-tts-10k-tagged+parler-tts/libritts_r_tags_tagged_10k_generated+parler-tts/libritts_r_tags_tagged_10k_generated" \
+    --train_dataset_config_name "read+default+clean+other" \
+    --train_split_name "train+train[:20%]+test.clean+test.other" \
+    --eval_dataset_name "ylacombe/expresso+reach-vb/jenny_tts_dataset+blabble-io/libritts_r+blabble-io/libritts_r" \
+    --eval_metadata_dataset_name "reach-vb/expresso-tagged-w-speech-mistral-v3+ylacombe/jenny-tts-10k-tagged+parler-tts/libritts_r_tags_tagged_10k_generated+parler-tts/libritts_r_tags_tagged_10k_generated" \
+    --eval_dataset_config_name "read+default+clean+other" \
+    --eval_split_name "train+train[:20%]+test.clean+test.other" \
+    --max_eval_samples 8 \
+    --per_device_eval_batch_size 16 \
+    --target_audio_column_name "audio" \
+    --description_column_name "text_description" \
+    --prompt_column_name "text" \
+    --max_duration_in_seconds 30.0 \
+    --min_duration_in_seconds 2.0 \
+    --max_text_length 400 \
+    --preprocessing_num_workers 2 \
+    --do_train true \
+    --num_train_epochs 8 \
+    --gradient_accumulation_steps 8 \
+    --gradient_checkpointing true \
+    --per_device_train_batch_size 16 \
+    --learning_rate 0.00008 \
+    --adam_beta1 0.9 \
+    --adam_beta2 0.99 \
+    --weight_decay 0.01 \
+    --lr_scheduler_type "cosine" \
+    --warmup_steps 250 \
+    --logging_steps 2 \
+    --freeze_text_encoder true \
+    --audio_encoder_per_device_batch_size 4 \
+    --dtype "bfloat16" \
+    --seed 456 \
+    --output_dir "./parler-tts-mini-expresso" \
+    --temporary_save_to_disk "./audio_code_tmp" \
+    --save_to_disk "./tmp_dataset_audio" \
+    --dataloader_num_workers 4 \
+    --do_eval \
+    --predict_with_generate \
+    --include_inputs_for_metrics \
+    --group_by_length true
+```
+
+On a single 80GB A100 GPU, training took approximately 1.5 hours and returned a final evaluation loss of 4.0. Again, the 
+script can be configured for multiple GPUs by running `accelerate config` from the command line; no further 
+code-changes are required.
+
+Training performance is quite sensitive to learning rate and number of epochs: you should tune these according to your task 
+and the size of your dataset. In our experiments, we found the best performance to occur after 8 epochs of training
+with a learning rate of 8e-5.
+
+If you followed to the end of these steps: congratulations! You should now have a fine-tuned model you can use for your
+downstream applications using the [inference code-example](#usage) above. You can try substituting your own dataset, or
+run training using a single-speaker dataset, like the [Jenny example](https://colab.research.google.com/github/ylacombe/scripts_and_notebooks/blob/main/Finetuning_Parler_TTS_on_a_single_speaker_dataset.ipynb).
+
+## Motivation
+
+Parler-TTS is a reproduction of work from the paper [Natural language guidance of high-fidelity text-to-speech with synthetic annotations](https://www.text-description-to-speech.com) by Dan Lyth and Simon King, from Stability AI and Edinburgh University respectively. 
+
+Contrarily to other TTS models, Parler-TTS is a **fully open-source** release. All datasets, pre-processing, training code and weights are released publicly under permissive license, enabling the community to build on our work and develop their own powerful TTS models.
+Parler-TTS was released alongside:
+* [The Parler-TTS repository](https://github.com/huggingface/parler-tts) - you can train and fine-tuned your own version of the model.
+* [The Data-Speech repository](https://github.com/huggingface/dataspeech) - a suite of utility scripts designed to annotate speech datasets.
+* [The Parler-TTS organization](https://huggingface.co/parler-tts) - where you can find the annotated datasets as well as the future checkpoints.
+
+## Citation
+
+If you found this repository useful, please consider citing this work and also the original Stability AI paper:
+
+```
+@misc{lacombe-etal-2024-parler-tts,
+  author = {Yoach Lacombe and Vaibhav Srivastav and Sanchit Gandhi},
+  title = {Parler-TTS},
+  year = {2024},
+  publisher = {GitHub},
+  journal = {GitHub repository},
+  howpublished = {\url{https://github.com/huggingface/parler-tts}}
+}
+```
+
+```
+@misc{lyth2024natural,
+      title={Natural language guidance of high-fidelity text-to-speech with synthetic annotations},
+      author={Dan Lyth and Simon King},
+      year={2024},
+      eprint={2402.01912},
+      archivePrefix={arXiv},
+      primaryClass={cs.SD}
+}
+```
+
+## License
+
+This model is permissively licensed under the Apache 2.0 license.

config.json

@@ -0,0 +1,276 @@
+{
+  "_name_or_path": "sanchit-gandhi/parler-tts-mini-v0.1-expresso-combined",
+  "architectures": [
+    "ParlerTTSForConditionalGeneration"
+  ],
+  "audio_encoder": {
+    "_name_or_path": "ylacombe/dac_44khZ_8kbps",
+    "add_cross_attention": false,
+    "architectures": [
+      "DACModel"
+    ],
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": null,
+    "chunk_size_feed_forward": 0,
+    "codebook_size": 1024,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": null,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "frame_rate": 86,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "latent_dim": 1024,
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "min_length": 0,
+    "model_bitrate": 8,
+    "model_type": "dac",
+    "no_repeat_ngram_size": 0,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_codebooks": 9,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": null,
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "sampling_rate": 44100,
+    "sep_token_id": null,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": "float32",
+    "torchscript": false,
+    "typical_p": 1.0,
+    "use_bfloat16": false
+  },
+  "decoder": {
+    "_name_or_path": "/fsx/yoach/tmp/artefacts/decoder_400M/",
+    "activation_dropout": 0.0,
+    "activation_function": "gelu",
+    "add_cross_attention": true,
+    "architectures": [
+      "ParlerTTSForCausalLM"
+    ],
+    "attention_dropout": 0.0,
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": 1025,
+    "chunk_size_feed_forward": 0,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "dropout": 0.1,
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": 1024,
+    "exponential_decay_length_penalty": null,
+    "ffn_dim": 4096,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "hidden_size": 1024,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_factor": 0.02,
+    "is_decoder": true,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layerdrop": 0.0,
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "max_position_embeddings": 4096,
+    "min_length": 0,
+    "model_type": "parler_tts_decoder",
+    "no_repeat_ngram_size": 0,
+    "num_attention_heads": 16,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_codebooks": 9,
+    "num_hidden_layers": 24,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": 1024,
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "scale_embedding": false,
+    "sep_token_id": null,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": false,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": "float32",
+    "torchscript": false,
+    "typical_p": 1.0,
+    "use_bfloat16": false,
+    "use_cache": true,
+    "vocab_size": 1088
+  },
+  "decoder_start_token_id": 1025,
+  "is_encoder_decoder": true,
+  "model_type": "parler_tts",
+  "pad_token_id": 1024,
+  "text_encoder": {
+    "_name_or_path": "google/flan-t5-base",
+    "add_cross_attention": false,
+    "architectures": [
+      "T5ForConditionalGeneration"
+    ],
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": null,
+    "chunk_size_feed_forward": 0,
+    "classifier_dropout": 0.0,
+    "cross_attention_hidden_size": null,
+    "d_ff": 2048,
+    "d_kv": 64,
+    "d_model": 768,
+    "decoder_start_token_id": 0,
+    "dense_act_fn": "gelu_new",
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "dropout_rate": 0.1,
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": 1,
+    "exponential_decay_length_penalty": null,
+    "feed_forward_proj": "gated-gelu",
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_factor": 1.0,
+    "is_decoder": false,
+    "is_encoder_decoder": true,
+    "is_gated_act": true,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layer_norm_epsilon": 1e-06,
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "min_length": 0,
+    "model_type": "t5",
+    "n_positions": 512,
+    "no_repeat_ngram_size": 0,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_decoder_layers": 12,
+    "num_heads": 12,
+    "num_layers": 12,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_past": true,
+    "output_scores": false,
+    "pad_token_id": 0,
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "relative_attention_max_distance": 128,
+    "relative_attention_num_buckets": 32,
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "sep_token_id": null,
+    "suppress_tokens": null,
+    "task_specific_params": {
+      "summarization": {
+        "early_stopping": true,
+        "length_penalty": 2.0,
+        "max_length": 200,
+        "min_length": 30,
+        "no_repeat_ngram_size": 3,
+        "num_beams": 4,
+        "prefix": "summarize: "
+      },
+      "translation_en_to_de": {
+        "early_stopping": true,
+        "max_length": 300,
+        "num_beams": 4,
+        "prefix": "translate English to German: "
+      },
+      "translation_en_to_fr": {
+        "early_stopping": true,
+        "max_length": 300,
+        "num_beams": 4,
+        "prefix": "translate English to French: "
+      },
+      "translation_en_to_ro": {
+        "early_stopping": true,
+        "max_length": 300,
+        "num_beams": 4,
+        "prefix": "translate English to Romanian: "
+      }
+    },
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": false,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": null,
+    "torchscript": false,
+    "typical_p": 1.0,
+    "use_bfloat16": false,
+    "use_cache": true,
+    "vocab_size": 32128
+  },
+  "torch_dtype": "float32",
+  "transformers_version": "4.41.0.dev0",
+  "vocab_size": 32128
+}

generation_config.json

@@ -0,0 +1,12 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1025,
+  "decoder_start_token_id": 1025,
+  "do_sample": true,
+  "eos_token_id": 1024,
+  "guidance_scale": 1.0,
+  "max_length": 2580,
+  "min_new_tokens": 50,
+  "pad_token_id": 1024,
+  "transformers_version": "4.41.0.dev0"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ab690bfb1b231f648dc3c141200f7a753fe4c84b4e3e30307f4491bccd5e7ee7
+size 2588215392

preprocessor_config.json

@@ -0,0 +1,10 @@
+{
+  "chunk_length_s": null,
+  "feature_extractor_type": "EncodecFeatureExtractor",
+  "feature_size": 1,
+  "overlap": null,
+  "padding_side": "right",
+  "padding_value": 0.0,
+  "return_attention_mask": true,
+  "sampling_rate": 44100
+}

run_prompt_creation_expresso.py

@@ -0,0 +1,520 @@
+import json
+import logging
+import os
+import re
+import shutil
+import sys
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from accelerate import Accelerator, skip_first_batches
+from accelerate.logging import get_logger
+from datasets import DatasetDict, load_dataset
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    BitsAndBytesConfig,
+    HfArgumentParser,
+)
+
+
+logger = get_logger(__name__, log_level="INFO")
+
+
+@dataclass
+class ModelArguments:
+    """
+    Arguments pertaining to what data we are going to input our model for training and eval.
+    """
+
+    model_name_or_path: str = field(
+        metadata={"help": "The name of the model to use (via the transformers library) for the prompt annotation."},
+    )
+    per_device_eval_batch_size: int = field(
+        metadata={"help": "The per-device batch size to use for inference."},
+    )
+    model_variant: str = field(
+        default=None,
+        metadata={"help": "If specified load weights from `variant` filename, *e.g.* pytorch_model.<variant>.bin. "},
+    )
+    model_revision: str = field(
+        default="main",
+        metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."},
+    )
+    cache_dir: Optional[str] = field(
+        default=None,
+        metadata={"help": "Where to store the pretrained models downloaded from huggingface.co"},
+    )
+    torch_dtype: Optional[str] = field(
+        default="float16",
+        metadata={
+            "help": (
+                "Floating-point format in which the model weights should be initialized"
+                " and the computations run. Choose one of `[float32, float16, bfloat16]`."
+            )
+        },
+    )
+    attn_implementation: Optional[str] = field(
+        default="sdpa",
+        metadata={"help": "Which attn type to use: ['eager', 'sdpa', 'flash_attention_2']"},
+    )
+    load_in_8bit: Optional[bool] = field(
+        default=False, metadata={"help": "Whether to use 8-bit precision for inference."}
+    )
+    load_in_4bit: Optional[bool] = field(
+        default=False, metadata={"help": "Whether to use 4-bit precision for inference."}
+    )
+    bnb_4bit_quant_type: Optional[str] = field(
+        default="nf4", metadata={"help": "precise the quantization type (fp4 or nf4)"}
+    )
+    use_bnb_nested_quant: Optional[bool] = field(default=False, metadata={"help": "use nested quantization"})
+    trust_remote_code: Optional[bool] = field(
+        default=False,
+        metadata={
+            "help": (
+                "Whether or not to allow for custom models defined on the Hub in their own modeling files. This option "
+                "should only be set to `True` for repositories you trust and in which you have read the code, as it will "
+                "execute code present on the Hub on your local machine."
+            )
+        },
+    )
+    use_fast_tokenizer: Optional[bool] = field(
+        default=True, metadata={"help": "Use fast tokenizer for encoding/decoding input ids"}
+    )
+    token: Optional[bool] = field(
+        default=True,
+        metadata={
+            "help": "Whether or not to use an authentication token when loading/uploading from the Hugging Face Hub"
+        },
+    )
+    do_sample: Optional[bool] = field(default=True, metadata={"help": "Whether to use sampling mode for generation"})
+    temperature: Optional[float] = field(default=0.6, metadata={"help": "Temperature for sampling-based generation"})
+    max_new_tokens: Optional[int] = field(
+        default=256, metadata={"help": "Maximum number of new tokens during generation"}
+    )
+    torch_compile: Optional[bool] = field(
+        default=False,
+        metadata={
+            "help": "Whether to compile the forward pass (not sampling) in generate. Only compatible with Gemma and LlaMA."
+        },
+    )
+
+
+@dataclass
+class DataArguments:
+    """
+    Arguments pertaining to what data we are going to input our model for training and eval.
+    """
+
+    output_dir: str = field(
+        metadata={
+            "help": "Where to save the processed dataset to disk. If unspecified, uses a 'pretty' version of the "
+            "original dataset name. E.g. 'facebook/voxpopuli' will be saved under 'voxpopuli'."
+        },
+    )
+    dataset_name: str = field(
+        default=None,
+        metadata={"help": "The name of the dataset to use (via the datasets library)"},
+    )
+    dataset_config_name: Optional[str] = field(
+        default=None,
+        metadata={"help": "The configuration name of the dataset to use (via the datasets library)."},
+    )
+    dataset_split_name: Optional[str] = field(
+        default=None,
+        metadata={"help": "The split name of the dataset to use (via the datasets library)."},
+    )
+    dataset_cache_dir: Optional[str] = field(
+        default=None,
+        metadata={"help": "Path to cache directory for saving and loading datasets"},
+    )
+    max_eval_samples: Optional[int] = field(
+        default=None,
+        metadata={"help": "Maximum number of samples for generation - use for debugging purposes."},
+    )
+    overwrite_cache: bool = field(
+        default=False,
+        metadata={"help": "Overwrite the cached training and evaluation sets"},
+    )
+    preprocessing_num_workers: Optional[int] = field(
+        default=None,
+        metadata={"help": "The number of processes to use for the preprocessing."},
+    )
+    dataloader_num_workers: Optional[int] = field(
+        default=0,
+        metadata={"help": "The number of processes to use for the dataloader."},
+    )
+    push_to_hub: Optional[bool] = field(
+        default=False,
+        metadata={"help": "Whether or not to push the processed dataset to the Hub."},
+    )
+    hub_dataset_id: Optional[str] = field(
+        default=None,
+        metadata={"help": "Repository namespace if pushing to the Hugging Face Hub."},
+    )
+    overwrite_output_dir: Optional[bool] = field(
+        default=False,
+        metadata={"help": "Overwrite the content of the output directory each time the script is run."},
+    )
+    save_steps: Optional[int] = field(
+        default=500,
+        metadata={"help": "Save the generated prompts every save_steps."},
+    )
+    save_total_limit: Optional[int] = field(
+        default=1, metadata={"help": ("If a value is passed, will limit the total number of saved checkpoints")}
+    )
+
+    def __post_init__(self):
+        if self.push_to_hub and self.hub_dataset_id is None:
+            raise ValueError("You must specify the `hub_dataset_id` when setting `--push_to_hub=True`")
+
+
+def get_quantization_config(model_args: ModelArguments) -> Union[BitsAndBytesConfig, None]:
+    if model_args.load_in_4bit:
+        compute_dtype = torch.float16
+        if model_args.torch_dtype not in {"auto", None}:
+            compute_dtype = getattr(torch, model_args.torch_dtype)
+
+        quantization_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=compute_dtype,
+            bnb_4bit_quant_type=model_args.bnb_4bit_quant_type,
+            bnb_4bit_use_double_quant=model_args.use_bnb_nested_quant,
+        )
+    elif model_args.load_in_8bit:
+        quantization_config = BitsAndBytesConfig(
+            load_in_8bit=True,
+        )
+    else:
+        quantization_config = None
+
+    return quantization_config
+
+
+def get_current_device() -> int:
+    """Get the current device. For GPU we return the local process index to enable multiple GPU training."""
+    return Accelerator().local_process_index if torch.cuda.is_available() else "cpu"
+
+
+def get_kbit_device_map() -> Union[Dict[str, int], None]:
+    """Useful for running inference with quantized models by setting `device_map=get_peft_device_map()`"""
+    return {"": get_current_device()} if torch.cuda.is_available() else None
+
+
+CHECKPOINT_PREFIX = "checkpoint"
+_RE_CHECKPOINT = re.compile(r"^checkpoint-(\d+).json$")
+
+
+def save_checkpoint(output_dir, all_generated_ids, step):
+    checkpoint_path = f"{CHECKPOINT_PREFIX}-{step}.json"
+    output_path = os.path.join(output_dir, checkpoint_path)
+    all_generated_ids = [ids.tolist() for ids in all_generated_ids]
+    with open(output_path, "w") as file:
+        json.dump(all_generated_ids, file)
+
+
+def load_checkpoint(checkpoint_path):
+    with open(checkpoint_path, "r") as file:
+        all_generated_ids = json.load(file)
+    all_generated_ids = [np.array(lst) for lst in all_generated_ids]
+    return all_generated_ids
+
+
+def sorted_checkpoints(output_dir=None) -> List[str]:
+    """Helper function to sort saved checkpoints from oldest to newest."""
+    ordering_and_checkpoint_path = []
+
+    glob_checkpoints = [str(x) for x in Path(output_dir).glob(f"{CHECKPOINT_PREFIX}-*")]
+
+    for path in glob_checkpoints:
+        regex_match = re.match(f".*{CHECKPOINT_PREFIX}-([0-9]+)", path)
+        if regex_match is not None and regex_match.groups() is not None:
+            ordering_and_checkpoint_path.append((int(regex_match.groups()[0]), path))
+
+    checkpoints_sorted = sorted(ordering_and_checkpoint_path)
+    checkpoints_sorted = [checkpoint[1] for checkpoint in checkpoints_sorted]
+    return checkpoints_sorted
+
+
+def rotate_checkpoints(save_total_limit=None, output_dir=None) -> None:
+    """Helper function to delete old checkpoints."""
+    if save_total_limit is None or save_total_limit <= 0:
+        return
+    # Check if we should delete older checkpoint(s)
+    checkpoints_sorted = sorted_checkpoints(output_dir=output_dir)
+    if len(checkpoints_sorted) <= save_total_limit:
+        return
+
+    number_of_checkpoints_to_delete = max(0, len(checkpoints_sorted) - save_total_limit)
+    checkpoints_to_be_deleted = checkpoints_sorted[:number_of_checkpoints_to_delete]
+    for checkpoint in checkpoints_to_be_deleted:
+        logger.info(f"Deleting older checkpoint [{checkpoint}] due to args.save_total_limit")
+        os.remove(checkpoint)
+
+
+def get_last_checkpoint(folder) -> Tuple[List, int]:
+    if not os.path.exists(folder) or not os.path.isdir(folder):
+        os.makedirs(folder, exist_ok=True)
+        return [], 0
+    content = os.listdir(folder)
+    checkpoints = [path for path in content if _RE_CHECKPOINT.search(path) is not None]
+    if len(checkpoints) == 0:
+        return [], 0
+    last_checkpoint = os.path.join(folder, max(checkpoints, key=lambda x: int(_RE_CHECKPOINT.search(x).groups()[0])))
+    # Find num steps saved state string pattern
+    pattern = r"checkpoint-(\d+).json"
+    match = re.search(pattern, last_checkpoint)
+    cur_step = int(match.group(1))
+    # load corresponding generated ids
+    all_generated_ids = load_checkpoint(last_checkpoint)
+    return all_generated_ids, cur_step
+
+
+@dataclass
+class DataCollatorWithPadding:
+    """
+    Data collator that will dynamically pad the inputs received to the longest sequence in the batch.
+    """
+
+    tokenizer: Any
+
+    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:
+        # split inputs and labels since they have to be of different lengths and need
+        # different padding methods
+        input_ids = {"input_ids": [feature["input_ids"] for feature in features]}
+        batch = self.tokenizer.pad(input_ids, return_tensors="pt", padding="longest", return_attention_mask=True)
+        return batch
+
+id_to_name = {
+    "ex01": "Jerry",
+    "ex02": "Elisabeth",
+    "ex03": "Thomas",
+    "ex04": "Talia"
+}
+
+PROMPT = """You will be given a name and an enunciation style related to an audio sample of someone speaking.
+1. The name will be one of: Jerry, Elisabeth, Thomas, Talia.
+2. The enunciation style will be one of: 'enunciated', 'happy', 'confused', 'default' (meaning no particular emotion conveyed), 'laughing', 'sad', 'whisper', 'emphasis'.
+3. The pace of the speaker's delivery (e.g., very slowly, quite slowly, slightly slowly, moderate speed, slightly fast, quite fast, very fast)
+
+Your task is to create a simple text description using these keywords that accurately describes the audio sample. Ensure that the generated description is grammatically correct, easy to understand, and most importantly, concise.
+
+For example, given the following keywords: 'Talia', 'happy', 'quite slowly', a valid description would be: 'Talia speaks happily and quite slowly with high quality.'. Another valid description would be: 'Talia delivers her words happily and quite slowly with high quality audio.'. Another example, given the following keywords: 'Jerry', 'emphasis', 'slightly slowly': 'Jerry speaks with emphasis on certain words and slightly slowly with high quality audio.'
+
+Each description is appended with 'with high quality'.
+
+You are free to change the order of the information, and replace synonymous terms. Give one description and nothing else. No alternatives or repeating the task. Remember to prioritise conciseness and simplicity.
+
+For the information: '[speaker_id]', '[style]', '[speaking_rate]' the corresponding description is:"""
+
+
+def main():
+    # 1. Parse input arguments
+    parser = HfArgumentParser((ModelArguments, DataArguments))
+    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
+        # If we pass only one argument to the script and it's the path to a json file,
+        # let's parse it to get our arguments.
+        model_args, data_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
+    else:
+        model_args, data_args = parser.parse_args_into_dataclasses()
+
+    # 2. Setup logging
+    # Make one log on every process with the configuration for debugging.
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        handlers=[logging.StreamHandler(sys.stdout)],
+    )
+    
+    accelerator = Accelerator()
+
+    if data_args.overwrite_output_dir and os.path.exists(data_args.output_dir) and os.path.isdir(data_args.output_dir):
+        logger.info("Cleaning output dir from previous run...")
+        shutil.rmtree(data_args.output_dir)
+
+    # 3. Load annotated dataset
+    logger.info("*** Load annotated dataset ***")
+    if data_args.dataset_split_name is not None:
+        raw_datasets = DatasetDict()
+        data_splits = data_args.dataset_split_name.split("+")
+        # load on a split-wise basis
+        for split in data_splits:
+            with accelerator.local_main_process_first():
+                raw_datasets[split] = load_dataset(
+                    data_args.dataset_name,
+                    data_args.dataset_config_name,
+                    split=split,
+                    cache_dir=model_args.cache_dir,
+                    token=model_args.token,
+                    num_proc=data_args.preprocessing_num_workers,
+                )
+    else:
+        with accelerator.local_main_process_first():
+            # load all splits for annotation
+            raw_datasets = load_dataset(
+                data_args.dataset_name,
+                data_args.dataset_config_name,
+                cache_dir=model_args.cache_dir,
+                token=model_args.token,
+                num_proc=data_args.preprocessing_num_workers,
+            )
+
+    raw_datasets_features = set(raw_datasets[next(iter(raw_datasets))].features.keys())
+
+    if data_args.max_eval_samples is not None:
+        for split in raw_datasets:
+            raw_datasets[split] = raw_datasets[split].select(range(data_args.max_eval_samples))
+
+    # TODO(SG): add accent
+    EXPECTED_COLUMNS = {"speaker_id", "style", "speaking_rate"}
+    if not EXPECTED_COLUMNS.issubset(raw_datasets_features):
+        missing_columns = EXPECTED_COLUMNS - raw_datasets_features
+        raise ValueError(
+            f"Missing columns {missing_columns} from the dataset features. Got dataset features {raw_datasets_features}"
+        )
+
+    # 4. Load pre-trained model
+    logger.info("*** Load pretrained model ***")
+    torch_dtype = (
+        model_args.torch_dtype if model_args.torch_dtype in ["auto", None] else getattr(torch, model_args.torch_dtype)
+    )
+    quantization_config = get_quantization_config(model_args)
+
+    model = AutoModelForCausalLM.from_pretrained(
+        model_args.model_name_or_path,
+        revision=model_args.model_revision,
+        variant=model_args.model_variant,
+        trust_remote_code=model_args.trust_remote_code,
+        attn_implementation=model_args.attn_implementation,
+        torch_dtype=torch_dtype,
+        device_map=get_kbit_device_map() if quantization_config is not None else None,
+        quantization_config=quantization_config,
+        low_cpu_mem_usage=True,
+        token=model_args.token,
+    ).eval()
+
+    if model_args.torch_compile:
+        # torch compile only compatible with gemma and llama
+        if not callable(getattr(model, "_setup_cache", None)):
+            raise ValueError(
+                f"Static k/v cache is not compatible with the model {model.__class__.__name__}. Set `--torch_compile=False"
+                "for dynamic k/v cache"
+            )
+        model.generation_config.cache_implementation = "static"
+        # compile the forward pass (but not the top-{p,k} sampling)
+        model = torch.compile(model, mode="reduce-overhead", fullgraph=True)
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_args.model_name_or_path,
+        revision=model_args.model_revision,
+        trust_remote_code=model_args.trust_remote_code,
+        use_fast=model_args.use_fast_tokenizer,
+        padding_side="left",
+    )
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token_id = tokenizer.bos_token_id
+        model.generation_config.pad_token_id = model.generation_config.eos_token_id
+
+
+    def prepare_dataset(sample):
+        sample_prompt = PROMPT
+        sample["speaker_id"] = id_to_name[sample["speaker_id"]]
+        for key in EXPECTED_COLUMNS:
+            sample_prompt = sample_prompt.replace(f"[{key}]", sample[key])
+        sample_prompt = [{"role": "user", "content": sample_prompt}]
+        token_ids = tokenizer.apply_chat_template(sample_prompt)
+        sample["input_ids"] = token_ids
+        return sample
+
+    with accelerator.local_main_process_first():
+        vectorized_datasets = raw_datasets.map(
+            prepare_dataset, num_proc=data_args.preprocessing_num_workers, desc="Preparing prompts"
+        )
+
+    # Prepare everything with our `accelerator`
+    model = accelerator.prepare(model)
+    data_collator = DataCollatorWithPadding(tokenizer)
+
+    def generate_step(batch):
+        output_ids = accelerator.unwrap_model(model).generate(
+            batch["input_ids"],
+            attention_mask=batch["attention_mask"],
+            do_sample=model_args.do_sample,
+            temperature=model_args.temperature,
+            max_new_tokens=model_args.max_new_tokens,
+        )
+        output_ids = accelerator.pad_across_processes(output_ids, dim=1, pad_index=tokenizer.pad_token_id)
+        return output_ids
+
+    def postprocess_dataset(sample):
+        prompt_text = tokenizer.decode(sample["input_ids"], skip_special_tokens=True)
+        generated_text = tokenizer.decode(sample["generated_ids"], skip_special_tokens=True)
+        sample["text_description"] = generated_text[len(prompt_text) :]
+        return sample
+
+    for split in vectorized_datasets:
+        data_loader = DataLoader(
+            vectorized_datasets[split],
+            batch_size=model_args.per_device_eval_batch_size,
+            collate_fn=data_collator,
+            num_workers=data_args.dataloader_num_workers,
+            pin_memory=True,
+        )
+        data_loader = accelerator.prepare(data_loader)
+        total_inference_steps = len(data_loader)
+        progress_bar = tqdm(
+            range(total_inference_steps), desc=" ... ", position=0, disable=not accelerator.is_local_main_process
+        )
+
+        split_output_dir = os.path.join(data_args.output_dir, split)
+        all_generated_ids, cur_step = get_last_checkpoint(split_output_dir)
+
+        if cur_step > 0:
+            logger.info(f"Resuming {split} from step {cur_step}")
+            # efficiently skip the first n batches
+            data_loader = skip_first_batches(data_loader, cur_step)
+            progress_bar.update(cur_step)
+
+        while cur_step < total_inference_steps:
+            for batch in data_loader:
+                generated_ids = generate_step(batch)
+                generated_ids = accelerator.gather_for_metrics(generated_ids)
+                all_generated_ids.extend(generated_ids.cpu().numpy())
+
+                cur_step += 1
+                progress_bar.update(1)
+
+                if (cur_step % data_args.save_steps == 0) or (cur_step == total_inference_steps):
+                    save_checkpoint(split_output_dir, all_generated_ids, cur_step)
+                    rotate_checkpoints(data_args.save_total_limit, output_dir=split_output_dir)
+
+        vectorized_datasets[split] = vectorized_datasets[split].add_column("generated_ids", all_generated_ids)
+
+        if accelerator.is_main_process:
+            vectorized_datasets[split] = vectorized_datasets[split].map(
+                postprocess_dataset,
+                num_proc=data_args.preprocessing_num_workers,
+                desc="Postprocessing dataset",
+                remove_columns=["input_ids", "generated_ids"],
+            )
+        accelerator.wait_for_everyone()
+
+    if accelerator.is_main_process:
+        vectorized_datasets.save_to_disk(data_args.output_dir)
+        if data_args.push_to_hub:
+            vectorized_datasets.push_to_hub(
+                data_args.hub_dataset_id,
+                config_name=data_args.dataset_config_name if data_args.dataset_config_name is not None else "default",
+                token=model_args.token,
+            )
+    accelerator.wait_for_everyone()
+    accelerator.end_training()
+
+
+if __name__ == "__main__":
+    main()

special_tokens_map.json

@@ -0,0 +1,125 @@
+{
+  "additional_special_tokens": [
+    "<extra_id_0>",
+    "<extra_id_1>",
+    "<extra_id_2>",
+    "<extra_id_3>",
+    "<extra_id_4>",
+    "<extra_id_5>",
+    "<extra_id_6>",
+    "<extra_id_7>",
+    "<extra_id_8>",
+    "<extra_id_9>",
+    "<extra_id_10>",
+    "<extra_id_11>",
+    "<extra_id_12>",
+    "<extra_id_13>",
+    "<extra_id_14>",
+    "<extra_id_15>",
+    "<extra_id_16>",
+    "<extra_id_17>",
+    "<extra_id_18>",
+    "<extra_id_19>",
+    "<extra_id_20>",
+    "<extra_id_21>",
+    "<extra_id_22>",
+    "<extra_id_23>",
+    "<extra_id_24>",
+    "<extra_id_25>",
+    "<extra_id_26>",
+    "<extra_id_27>",
+    "<extra_id_28>",
+    "<extra_id_29>",
+    "<extra_id_30>",
+    "<extra_id_31>",
+    "<extra_id_32>",
+    "<extra_id_33>",
+    "<extra_id_34>",
+    "<extra_id_35>",
+    "<extra_id_36>",
+    "<extra_id_37>",
+    "<extra_id_38>",
+    "<extra_id_39>",
+    "<extra_id_40>",
+    "<extra_id_41>",
+    "<extra_id_42>",
+    "<extra_id_43>",
+    "<extra_id_44>",
+    "<extra_id_45>",
+    "<extra_id_46>",
+    "<extra_id_47>",
+    "<extra_id_48>",
+    "<extra_id_49>",
+    "<extra_id_50>",
+    "<extra_id_51>",
+    "<extra_id_52>",
+    "<extra_id_53>",
+    "<extra_id_54>",
+    "<extra_id_55>",
+    "<extra_id_56>",
+    "<extra_id_57>",
+    "<extra_id_58>",
+    "<extra_id_59>",
+    "<extra_id_60>",
+    "<extra_id_61>",
+    "<extra_id_62>",
+    "<extra_id_63>",
+    "<extra_id_64>",
+    "<extra_id_65>",
+    "<extra_id_66>",
+    "<extra_id_67>",
+    "<extra_id_68>",
+    "<extra_id_69>",
+    "<extra_id_70>",
+    "<extra_id_71>",
+    "<extra_id_72>",
+    "<extra_id_73>",
+    "<extra_id_74>",
+    "<extra_id_75>",
+    "<extra_id_76>",
+    "<extra_id_77>",
+    "<extra_id_78>",
+    "<extra_id_79>",
+    "<extra_id_80>",
+    "<extra_id_81>",
+    "<extra_id_82>",
+    "<extra_id_83>",
+    "<extra_id_84>",
+    "<extra_id_85>",
+    "<extra_id_86>",
+    "<extra_id_87>",
+    "<extra_id_88>",
+    "<extra_id_89>",
+    "<extra_id_90>",
+    "<extra_id_91>",
+    "<extra_id_92>",
+    "<extra_id_93>",
+    "<extra_id_94>",
+    "<extra_id_95>",
+    "<extra_id_96>",
+    "<extra_id_97>",
+    "<extra_id_98>",
+    "<extra_id_99>"
+  ],
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<pad>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

spiece.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d60acb128cf7b7f2536e8f38a5b18a05535c9e14c7a355904270e15b0945ea86
+size 791656

tokenizer_config.json

@@ -0,0 +1,941 @@
+{
+  "add_prefix_space": true,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32000": {
+      "content": "<extra_id_99>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32001": {
+      "content": "<extra_id_98>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32002": {
+      "content": "<extra_id_97>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32003": {
+      "content": "<extra_id_96>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32004": {
+      "content": "<extra_id_95>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32005": {
+      "content": "<extra_id_94>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32006": {
+      "content": "<extra_id_93>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32007": {
+      "content": "<extra_id_92>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32008": {
+      "content": "<extra_id_91>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32009": {
+      "content": "<extra_id_90>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32010": {
+      "content": "<extra_id_89>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32011": {
+      "content": "<extra_id_88>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32012": {
+      "content": "<extra_id_87>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32013": {
+      "content": "<extra_id_86>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32014": {
+      "content": "<extra_id_85>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32015": {
+      "content": "<extra_id_84>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32016": {
+      "content": "<extra_id_83>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32017": {
+      "content": "<extra_id_82>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32018": {
+      "content": "<extra_id_81>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32019": {
+      "content": "<extra_id_80>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32020": {
+      "content": "<extra_id_79>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32021": {
+      "content": "<extra_id_78>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32022": {
+      "content": "<extra_id_77>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32023": {
+      "content": "<extra_id_76>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32024": {
+      "content": "<extra_id_75>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32025": {
+      "content": "<extra_id_74>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32026": {
+      "content": "<extra_id_73>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32027": {
+      "content": "<extra_id_72>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32028": {
+      "content": "<extra_id_71>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32029": {
+      "content": "<extra_id_70>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32030": {
+      "content": "<extra_id_69>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32031": {
+      "content": "<extra_id_68>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32032": {
+      "content": "<extra_id_67>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32033": {
+      "content": "<extra_id_66>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32034": {
+      "content": "<extra_id_65>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32035": {
+      "content": "<extra_id_64>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32036": {
+      "content": "<extra_id_63>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32037": {
+      "content": "<extra_id_62>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32038": {
+      "content": "<extra_id_61>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32039": {
+      "content": "<extra_id_60>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32040": {
+      "content": "<extra_id_59>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32041": {
+      "content": "<extra_id_58>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32042": {
+      "content": "<extra_id_57>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32043": {
+      "content": "<extra_id_56>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32044": {
+      "content": "<extra_id_55>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32045": {
+      "content": "<extra_id_54>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32046": {
+      "content": "<extra_id_53>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32047": {
+      "content": "<extra_id_52>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32048": {
+      "content": "<extra_id_51>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32049": {
+      "content": "<extra_id_50>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32050": {
+      "content": "<extra_id_49>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32051": {
+      "content": "<extra_id_48>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32052": {
+      "content": "<extra_id_47>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32053": {
+      "content": "<extra_id_46>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32054": {
+      "content": "<extra_id_45>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32055": {
+      "content": "<extra_id_44>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32056": {
+      "content": "<extra_id_43>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32057": {
+      "content": "<extra_id_42>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32058": {
+      "content": "<extra_id_41>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32059": {
+      "content": "<extra_id_40>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32060": {
+      "content": "<extra_id_39>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32061": {
+      "content": "<extra_id_38>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32062": {
+      "content": "<extra_id_37>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32063": {
+      "content": "<extra_id_36>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32064": {
+      "content": "<extra_id_35>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32065": {
+      "content": "<extra_id_34>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32066": {
+      "content": "<extra_id_33>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32067": {
+      "content": "<extra_id_32>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32068": {
+      "content": "<extra_id_31>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32069": {
+      "content": "<extra_id_30>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32070": {
+      "content": "<extra_id_29>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32071": {
+      "content": "<extra_id_28>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32072": {
+      "content": "<extra_id_27>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32073": {
+      "content": "<extra_id_26>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32074": {
+      "content": "<extra_id_25>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32075": {
+      "content": "<extra_id_24>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32076": {
+      "content": "<extra_id_23>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32077": {
+      "content": "<extra_id_22>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32078": {
+      "content": "<extra_id_21>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32079": {
+      "content": "<extra_id_20>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32080": {
+      "content": "<extra_id_19>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32081": {
+      "content": "<extra_id_18>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32082": {
+      "content": "<extra_id_17>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32083": {
+      "content": "<extra_id_16>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32084": {
+      "content": "<extra_id_15>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32085": {
+      "content": "<extra_id_14>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32086": {
+      "content": "<extra_id_13>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32087": {
+      "content": "<extra_id_12>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32088": {
+      "content": "<extra_id_11>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32089": {
+      "content": "<extra_id_10>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32090": {
+      "content": "<extra_id_9>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32091": {
+      "content": "<extra_id_8>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32092": {
+      "content": "<extra_id_7>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32093": {
+      "content": "<extra_id_6>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32094": {
+      "content": "<extra_id_5>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32095": {
+      "content": "<extra_id_4>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32096": {
+      "content": "<extra_id_3>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32097": {
+      "content": "<extra_id_2>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32098": {
+      "content": "<extra_id_1>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32099": {
+      "content": "<extra_id_0>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<extra_id_0>",
+    "<extra_id_1>",
+    "<extra_id_2>",
+    "<extra_id_3>",
+    "<extra_id_4>",
+    "<extra_id_5>",
+    "<extra_id_6>",
+    "<extra_id_7>",
+    "<extra_id_8>",
+    "<extra_id_9>",
+    "<extra_id_10>",
+    "<extra_id_11>",
+    "<extra_id_12>",
+    "<extra_id_13>",
+    "<extra_id_14>",
+    "<extra_id_15>",
+    "<extra_id_16>",
+    "<extra_id_17>",
+    "<extra_id_18>",
+    "<extra_id_19>",
+    "<extra_id_20>",
+    "<extra_id_21>",
+    "<extra_id_22>",
+    "<extra_id_23>",
+    "<extra_id_24>",
+    "<extra_id_25>",
+    "<extra_id_26>",
+    "<extra_id_27>",
+    "<extra_id_28>",
+    "<extra_id_29>",
+    "<extra_id_30>",
+    "<extra_id_31>",
+    "<extra_id_32>",
+    "<extra_id_33>",
+    "<extra_id_34>",
+    "<extra_id_35>",
+    "<extra_id_36>",
+    "<extra_id_37>",
+    "<extra_id_38>",
+    "<extra_id_39>",
+    "<extra_id_40>",
+    "<extra_id_41>",
+    "<extra_id_42>",
+    "<extra_id_43>",
+    "<extra_id_44>",
+    "<extra_id_45>",
+    "<extra_id_46>",
+    "<extra_id_47>",
+    "<extra_id_48>",
+    "<extra_id_49>",
+    "<extra_id_50>",
+    "<extra_id_51>",
+    "<extra_id_52>",
+    "<extra_id_53>",
+    "<extra_id_54>",
+    "<extra_id_55>",
+    "<extra_id_56>",
+    "<extra_id_57>",
+    "<extra_id_58>",
+    "<extra_id_59>",
+    "<extra_id_60>",
+    "<extra_id_61>",
+    "<extra_id_62>",
+    "<extra_id_63>",
+    "<extra_id_64>",
+    "<extra_id_65>",
+    "<extra_id_66>",
+    "<extra_id_67>",
+    "<extra_id_68>",
+    "<extra_id_69>",
+    "<extra_id_70>",
+    "<extra_id_71>",
+    "<extra_id_72>",
+    "<extra_id_73>",
+    "<extra_id_74>",
+    "<extra_id_75>",
+    "<extra_id_76>",
+    "<extra_id_77>",
+    "<extra_id_78>",
+    "<extra_id_79>",
+    "<extra_id_80>",
+    "<extra_id_81>",
+    "<extra_id_82>",
+    "<extra_id_83>",
+    "<extra_id_84>",
+    "<extra_id_85>",
+    "<extra_id_86>",
+    "<extra_id_87>",
+    "<extra_id_88>",
+    "<extra_id_89>",
+    "<extra_id_90>",
+    "<extra_id_91>",
+    "<extra_id_92>",
+    "<extra_id_93>",
+    "<extra_id_94>",
+    "<extra_id_95>",
+    "<extra_id_96>",
+    "<extra_id_97>",
+    "<extra_id_98>",
+    "<extra_id_99>"
+  ],
+  "clean_up_tokenization_spaces": true,
+  "eos_token": "</s>",
+  "extra_ids": 100,
+  "legacy": true,
+  "model_max_length": 512,
+  "pad_token": "<pad>",
+  "padding_side": "left",
+  "sp_model_kwargs": {},
+  "tokenizer_class": "T5Tokenizer",
+  "unk_token": "<unk>"
+}