up

create_model.py

@@ -0,0 +1,56 @@
+import numpy as np
+from transformers import AutoFeatureExtractor, AutoTokenizer, FlaxSpeechEncoderDecoderModel, GPT2Tokenizer
+
+encoder_id = "facebook/wav2vec2-large-lv60"
+decoder_id = "gpt2-medium"
+
+model = FlaxSpeechEncoderDecoderModel.from_encoder_decoder_pretrained(encoder_id, decoder_id, encoder_add_adapter=True)
+
+# set all encoder regularisation to zero
+model.config.encoder.feat_proj_dropout = 0.0
+model.config.encoder.final_dropout = 0.0
+model.config.encoder.activation_dropout = 0.0
+model.config.encoder.apply_spec_augment = False
+model.config.encoder.attention_dropout = 0.0
+model.config.encoder.feat_extract_dropout = 0.0
+model.config.encoder.feat_proj_dropout = 0.0
+model.config.encoder.hidden_dropout = 0.0
+model.config.encoder.hidden_dropout_prob = 0.0
+model.config.encoder.layerdrop = 0.0
+model.config.encoder.mask_feature_prob = 0.0
+model.config.encoder.mask_time_prob = 0.0
+
+# set all decoder regularisation to zero
+model.config.decoder.attn_pdrop = 0.0
+model.config.decoder.embd_pdrop = 0.0
+model.config.decoder.resid_pdrop = 0.0
+model.config.decoder.summary_first_dropout = 0.0
+
+# force GPT2 to append EOS to begin and end of seq
+def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
+    outputs = [self.bos_token_id] + token_ids_0 + [self.eos_token_id]
+    return outputs
+
+GPT2Tokenizer.build_inputs_with_special_tokens = build_inputs_with_special_tokens
+gpt2_tokenizer = GPT2Tokenizer.from_pretrained(decoder_id)
+# set pad_token_id to unk_token_id, note: unk_token_id == eos_token_id == bos_token_id
+gpt2_tokenizer.pad_token = gpt2_tokenizer.unk_token
+gpt2_tokenizer.save_pretrained("./")
+
+model.config.pad_token_id = gpt2_tokenizer.pad_token_id
+model.config.decoder_start_token_id = model.config.decoder.bos_token_id
+model.config.eos_token_id = model.config.decoder.eos_token_id
+model.config.max_length = 50
+model.config.num_beams = 1
+
+model.config.use_cache = False
+model.config.decoder.use_cache = False
+model.config.processor_class = "Wav2Vec2Processor"
+
+# check if generation works
+out = model.generate(np.ones((1, 2000)))
+
+model.save_pretrained("./")
+
+feature_extractor = AutoFeatureExtractor.from_pretrained(encoder_id)
+feature_extractor.save_pretrained("./")

run_flax_speech_recognition_seq2seq.py

@@ -0,0 +1 @@
+/home/sanchitgandhi/transformers/examples/flax/speech-recognition/run_flax_speech_recognition_seq2seq.py

run_librispeech.sh

@@ -0,0 +1,31 @@
+#!/usr/bin/env bash
+python run_flax_speech_recognition_seq2seq.py \
+        --dataset_name="librispeech_asr" \
+        --model_name_or_path="./" \
+        --dataset_config_name="clean" \
+        --train_split_name="train.100" \
+        --eval_split_name="validation" \
+        --output_dir="./" \
+        --preprocessing_num_workers="16" \
+        --length_column_name="input_length" \
+        --overwrite_output_dir \
+        --num_train_epochs="1" \
+        --per_device_train_batch_size="4" \
+        --per_device_eval_batch_size="4" \
+        --logging_steps="25" \
+        --max_duration_in_seconds="15" \
+        --max_target_length="64" \
+        --generation_max_length="40" \
+        --generation_num_beams="1" \
+        --learning_rate="3e-4" \
+        --warmup_steps="100" \
+        --text_column_name="text" \
+        --save_total_limit="1" \
+        --freeze_feature_encoder \
+        --predict_with_generate \
+        --do_lower_case \
+        --do_eval \
+        --do_train \
+        --push_to_hub \
+        --use_auth_token 	
+