add training notebooks

README.md

@@ -37,7 +37,16 @@ the decoder from [dbmdz/german-gpt2](https://huggingface.co/dbmdz/german-gpt2).
 
 It was trained using a two step process:
 * fine-tuning only the cross-attention weights and the decoder using the pre-computed outputs of the Wav2Vec-Modell
+  * relatively fast training
+  * also works on small GPU (eg. 8 GB)
+  * but may take a lot of disk space
+  * should already yield decent results
 * fine-tuning the model end-to-end
+  * much slower
+  * needs a bigger GPU
 
 There is also one trick, which seemed to improve performance significantly: adding position embeddings to the
-encoder outputs and initializing them with the pre-trained position embeddings of the GPT2 model (See `eval.py`).
+encoder outputs and initializing them with the pre-trained position embeddings of the GPT2 model (See `eval.py`).
+
+The training notebooks are still early drafts. Also results can probably improved a lot by using for example a learning
+rate schedule.

training/data_loading.py

@@ -0,0 +1,41 @@
+import torch
+from torch.utils.data import Dataset
+from pathlib import Path
+
+
+class S2TDataset(Dataset):
+    def __init__(self, data_path):
+        self.path = Path(data_path)
+        self.files = list(self.path.iterdir())
+
+    def __len__(self):
+        return len(self.files)
+
+    def __getitem__(self, idx):
+        file_path = self.files[idx]
+        eg = torch.load(file_path)
+        eg['file_path'] = file_path
+        return eg
+
+
+# TODO: Somehow masks do not work yet (bad performace), but Training also works w/o using the mask.
+def make_collate_fn(tokenizer):
+    def collate_fn(examples):
+        wav2vec_feats = [eg['wave2vec_features'] for eg in examples]
+        max_len = len(max(wav2vec_feats, key=len))
+        padded_feats, attention_masks = [], []
+        for feats in wav2vec_feats:
+            num_pads = max_len - len(feats)
+            padded_feats.append(torch.cat([feats, torch.zeros((num_pads, feats.shape[-1]), device=feats.device)]))
+            if num_pads > 0:
+                mask = torch.zeros((max_len,), device=feats.device).long()
+                mask[:-num_pads] = 1
+            else:
+                mask = torch.ones((max_len,), device=feats.device).long()
+            attention_masks.append(mask)
+
+        encoder_hidden_states = torch.stack(padded_feats, dim=0)
+        encoder_attention_masks = torch.stack(attention_masks, dim=0).bool()
+        input_ids = tokenizer([eg['sentence'] for eg in examples], return_tensors='pt', padding=True).input_ids
+        return encoder_hidden_states, encoder_attention_masks, input_ids
+    return collate_fn

training/decoder_only_training.ipynb

@@ -0,0 +1,371 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "521e21ab",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# This notebook is currently designed for a GPU using fp16. Hyperparameters however are barely tuned."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1732f970",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import random\n",
+    "import torch\n",
+    "from pathlib import Path"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f55f4047",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "EXPERIMENT_NAME = '00'\n",
+    "DATA_PATH = Path('../data/common_voice/de')\n",
+    "\n",
+    "model_dir = Path('decoder_only/de') / EXPERIMENT_NAME\n",
+    "log_dir = model_dir / 'logs'\n",
+    "log_dir.mkdir(exist_ok=True, parents=True)\n",
+    "\n",
+    "config = {\n",
+    "    'use_train_frac': 1.0, # When using all samples the wav2vec-outputs take up ~275GB disk space!!(~360,000 samples)\n",
+    "    'use_val_frac': 0.2,\n",
+    "    'encoder_id': 'jonatasgrosman/wav2vec2-large-xlsr-53-german',\n",
+    "    'decoder_id': 'dbmdz/german-gpt2',\n",
+    "    'decoder_pad_token': '_',\n",
+    "    'decoder_bos_token': '~',\n",
+    "    'num_beams': 1,\n",
+    "    'batch_size': 16,\n",
+    "    'weight_decay': 0.,\n",
+    "    'accumulate_grad': 2,\n",
+    "    'max_epochs': 10,\n",
+    "    'max_len': 36  # len(max(tokenizer(common_voice['validation']['sentence'] + common_voice['test']['sentence']).input_ids, key=len))\n",
+    "}"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "eb3de6a4",
+   "metadata": {},
+   "source": [
+    "# Feature Extraction"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b176328e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from huggingface_hub import notebook_login\n",
+    "from datasets import load_dataset\n",
+    "from datasets.features import Audio\n",
+    "from transformers import Wav2Vec2Model, Wav2Vec2FeatureExtractor"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "54e70696",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "notebook_login()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f0d22752",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def extract_features_to_files(model, feature_extractor, dataset_split, batch_size, output_path):\n",
+    "    output_path = Path(output_path)\n",
+    "    output_path.mkdir(parents=True, exist_ok=True)\n",
+    "\n",
+    "    model.eval().cuda()\n",
+    "    for i in range(0, len(dataset_split), batch_size):\n",
+    "        batch = dataset_split[i:i+batch_size]\n",
+    "        sent_batch = batch['sentence']\n",
+    "        audio_batch = batch['audio']\n",
+    "        for i, eg in enumerate(audio_batch):\n",
+    "            # Remove the longest examples, should be only three and these may lead to OOM- or Index-Errors.\n",
+    "            if len(eg['array']) > 300_000:\n",
+    "                print('Too Long.')\n",
+    "                audio_batch.pop(i)\n",
+    "                sent_batch.pop(i)\n",
+    "        features = feature_extractor([eg['array'] for eg in audio_batch],\n",
+    "                                     sampling_rate=16_000,\n",
+    "                                     return_tensors='pt',\n",
+    "                                     padding='longest')\n",
+    "\n",
+    "        with torch.no_grad():\n",
+    "            out = model(features.input_values.cuda(), attention_mask=features.attention_mask.cuda())\n",
+    "\n",
+    "        assert len(sent_batch) == len(audio_batch) == len(out.last_hidden_state)\n",
+    "        for sent, audio, hs in zip(sent_batch, audio_batch, out.last_hidden_state.bfloat16().cpu()):\n",
+    "            file_name = audio['path'].split('/')[-1]\n",
+    "            torch.save(\n",
+    "                # .clone() is necessary: https://github.com/pytorch/pytorch/issues/1995\n",
+    "                {'sentence': sent, 'wave2vec_features': hs.clone()},\n",
+    "                output_path / file_name\n",
+    "            )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "06324b6f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "if not DATA_PATH.exists():\n",
+    "    \n",
+    "    common_voice = load_dataset('mozilla-foundation/common_voice_7_0', 'de', use_auth_token=True)\n",
+    "    \n",
+    "    random.seed(419)\n",
+    "    train_inds = list(range(len(common_voice['train'])))\n",
+    "    random.shuffle(train_inds)\n",
+    "    val_inds = list(range(len(common_voice['validation'])))\n",
+    "    random.shuffle(val_inds)\n",
+    "    \n",
+    "    train_inds = train_inds[:int(config['use_train_frac'] * len(train_inds))]\n",
+    "    train = common_voice['train'].select(train_inds)\n",
+    "    train = train.cast_column('audio', Audio(sampling_rate=16_000))\n",
+    "    \n",
+    "    val_inds = val_inds[:int(config['use_val_frac'] * len(val_inds))]\n",
+    "    val = common_voice['validation'].select(val_inds)\n",
+    "    val = val.cast_column('audio', Audio(sampling_rate=16_000))\n",
+    "    \n",
+    "    # Load Model for feature extraction.\n",
+    "    wave2vec_extractor = Wav2Vec2FeatureExtractor.from_pretrained(config['encoder_id'])\n",
+    "    wave2vec = Wav2Vec2Model.from_pretrained(config['encoder_id'])\n",
+    "    wave2vec.eval().cuda()\n",
+    "    \n",
+    "    extract_features_to_files(wave2vec, wave2vec_extractor, train, batch_size=8, output_path=DATA_PATH / 'train')\n",
+    "    extract_features_to_files(wave2vec, wave2vec_extractor, val, batch_size=8, output_path=DATA_PATH / 'val')\n",
+    "    \n",
+    "    wave2vec.cpu()\n",
+    "    torch.cuda.empty_cache()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b2ae2a47",
+   "metadata": {},
+   "source": [
+    "# Training"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "188ef54f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import json\n",
+    "from accelerate import Accelerator\n",
+    "from torch.utils.data import DataLoader\n",
+    "from torch.optim import AdamW\n",
+    "from torch.utils.tensorboard import SummaryWriter\n",
+    "from transformers import AutoTokenizer, Wav2Vec2FeatureExtractor\n",
+    "from transformers.models.wav2vec2.modeling_wav2vec2 import Wav2Vec2BaseModelOutput\n",
+    "from data_loading import make_collate_fn, S2TDataset\n",
+    "from wer import calculate_wer  # Not what's used in eval.py.\n",
+    "from model import Wav2VecGPT2Model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "41518c81",
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "tokenizer = AutoTokenizer.from_pretrained(config['decoder_id'])\n",
+    "tokenizer.add_special_tokens({'pad_token': config['decoder_pad_token'], 'bos_token': config['decoder_bos_token']})\n",
+    "\n",
+    "model = Wav2VecGPT2Model.from_encoder_decoder_pretrained(\n",
+    "    config['encoder_id'], config['decoder_id'], max_length=config['max_len'], num_beams=config['num_beams']\n",
+    ")\n",
+    "\n",
+    "model.config.decoder_start_token_id = tokenizer.bos_token_id\n",
+    "model.config.pad_token_id = tokenizer.pad_token_id"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a95ec028",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "collate_fn = make_collate_fn(tokenizer)\n",
+    "\n",
+    "train_ds = S2TDataset(DATA_PATH / 'train')\n",
+    "train_dl = DataLoader(train_ds, batch_size=config['batch_size'], shuffle=True, collate_fn=collate_fn, num_workers=4)\n",
+    "\n",
+    "val_ds = S2TDataset(DATA_PATH / 'val')\n",
+    "val_dl = DataLoader(val_ds, batch_size=config['batch_size'], shuffle=False, collate_fn=collate_fn, num_workers=4)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0aaeeced",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "high_lr_modules = ['cross_attn', 'crossattention', 'enc_to_dec_proj', 'encoder_outputs_pos_emb']\n",
+    "high_lr_params = [p for n, p in model.named_parameters() if any(m in n for m in high_lr_modules)]\n",
+    "\n",
+    "optimizer_grouped_parameters = [\n",
+    "    {\n",
+    "        \"params\": high_lr_params,\n",
+    "        \"lr\": 5e-4,\n",
+    "    },\n",
+    "    {\n",
+    "        \"params\": [p for n, p in model.decoder.named_parameters() if not any(m in n for m in high_lr_modules)],\n",
+    "        \"lr\": 1e-6,\n",
+    "    },\n",
+    "]\n",
+    "optimizer = AdamW(optimizer_grouped_parameters, weight_decay=0.)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cf98d090",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "accelerator = Accelerator(fp16=True)\n",
+    "print(f'Using {accelerator.device}.')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "da9e928e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model, optimizer, train_dl, val_dl = accelerator.prepare(model, optimizer, train_dl, val_dl)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f191f256",
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "with open(log_dir / 'config.json', 'w') as config_file:\n",
+    "    json.dump(config, config_file, indent=4)\n",
+    "    \n",
+    "writer = SummaryWriter(log_dir)\n",
+    "val_golds = [eg['sentence'] for eg in val_ds]\n",
+    "best_val_wer = 10.\n",
+    "global_train_step = 0\n",
+    "\n",
+    "for epoch in range(config['max_epochs']):\n",
+    "    \n",
+    "    model.train()\n",
+    "    model.encoder.cpu()  # Model gets moved to gpu for evaluation (see below).\n",
+    "    torch.cuda.empty_cache()\n",
+    "    for batch_step, (encoder_hidden_states, att_mask, input_ids) in enumerate(train_dl):\n",
+    "        if encoder_hidden_states.shape[1] > 1024:\n",
+    "            # That's too long for the position embeddings. \n",
+    "            # TODO: handle this in model code.\n",
+    "            print(f'SKIPPED: {encoder_hidden_states.shape}')\n",
+    "            continue\n",
+    "        global_train_step += 1\n",
+    "        \n",
+    "        out = model(labels=input_ids, encoder_outputs=Wav2Vec2BaseModelOutput(encoder_hidden_states))\n",
+    "        accelerator.backward(out.loss)\n",
+    "        writer.add_scalar('train_loss', out.loss.item(), global_train_step)\n",
+    "        \n",
+    "        if (batch_step + 1) % config['accumulate_grad'] == 0:\n",
+    "            optimizer.step()\n",
+    "            optimizer.zero_grad()\n",
+    "            \n",
+    "        if batch_step % 300 == 0:\n",
+    "            print(out.loss.item())\n",
+    "            \n",
+    "    model.eval()\n",
+    "    model.cuda()  # Necessary for input_ids to be initialized on the correct device.\n",
+    "    val_preds = []\n",
+    "    for encoder_hidden_states, att_mask, _ in val_dl:\n",
+    "        with torch.no_grad():\n",
+    "            generated = model.generate(\n",
+    "                encoder_outputs=Wav2Vec2BaseModelOutput(last_hidden_state=encoder_hidden_states)\n",
+    "            )\n",
+    "        val_preds += tokenizer.batch_decode(generated)\n",
+    "    val_preds = [pred.lstrip('~').rstrip('_') for pred in val_preds]\n",
+    "    wer = calculate_wer(val_preds, val_golds)\n",
+    "    writer.add_scalar('val_wer', wer, epoch)\n",
+    "    print('WER: ', wer)\n",
+    "    \n",
+    "    if wer < best_val_wer:\n",
+    "        torch.save(model.state_dict(), model_dir / 'model.pt')\n",
+    "        print('Saved Model.')\n",
+    "        best_val_wer = wer"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d84a7e5c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# # Load saved pytorch model and save with all necessary model files.\n",
+    "# output_path = model_dir /'full_model'\n",
+    "# \n",
+    "# model.load_state_dict(torch.load(model_dir / 'model.pt'))\n",
+    "# \n",
+    "# tokenizer.save_pretrained(output_path)\n",
+    "# wave2vec_extractor = Wav2Vec2FeatureExtractor.from_pretrained(config['encoder_id'])\n",
+    "# wave2vec_extractor.save_pretrained(output_path)\n",
+    "# model.save_pretrained(output_path)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.7"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

training/end2end_training.ipynb

@@ -0,0 +1,269 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9e852db9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# This notebook is currently designed for a GPU using fp16. Hyperparameters however are barely tuned."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e730080b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import json\n",
+    "import random\n",
+    "import torch\n",
+    "from pathlib import Path\n",
+    "from accelerate import Accelerator\n",
+    "from datasets import load_dataset, concatenate_datasets\n",
+    "from datasets.features import Audio\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "from torch.optim import AdamW\n",
+    "from torch.utils.tensorboard import SummaryWriter\n",
+    "from transformers import AutoTokenizer, Wav2Vec2FeatureExtractor\n",
+    "from wer import calculate_wer  # Not what's used in eval.py.\n",
+    "from model import Wav2VecGPT2Model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "72af6337",
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
+    "common_voice = load_dataset('mozilla-foundation/common_voice_7_0', 'de', use_auth_token=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6396e61d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "EXPERIMENT_NAME = '00'\n",
+    "\n",
+    "model_dir = Path('end2end/de') / EXPERIMENT_NAME\n",
+    "log_dir = model_dir / 'logs'\n",
+    "log_dir.mkdir(exist_ok=True, parents=True)\n",
+    "\n",
+    "config = {\n",
+    "    'encoder_id': 'jonatasgrosman/wav2vec2-large-xlsr-53-german',\n",
+    "    'decoder_id': 'dbmdz/german-gpt2',\n",
+    "    'decoder_pad_token': '_',\n",
+    "    'decoder_bos_token': '~',\n",
+    "    'num_beams': 1,\n",
+    "    'num_val_examples': 1500,\n",
+    "    'batch_size': 8,\n",
+    "    'base_lr': 3e-4,\n",
+    "    'weight_decay': 0.,\n",
+    "    'accumulate_grad': 4,\n",
+    "    'max_epochs': 10,\n",
+    "    'max_len': 36  # len(max(tokenizer(common_voice['validation']['sentence'] + common_voice['test']['sentence']).input_ids, key=len))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6c632a61",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tokenizer = AutoTokenizer.from_pretrained(config['decoder_id'])\n",
+    "tokenizer.add_special_tokens({'pad_token': config['decoder_pad_token'], 'bos_token': config['decoder_bos_token']})\n",
+    "\n",
+    "wave2vec_extractor = Wav2Vec2FeatureExtractor.from_pretrained(config['encoder_id'])\n",
+    "\n",
+    "model = SpeechEncoderDecoderModel.from_encoder_decoder_pretrained(\n",
+    "    config['encoder_id'], config['decoder_id'], max_length=config['max_len'], num_beams=config['num_beams']\n",
+    ")\n",
+    "\n",
+    "model.config.decoder_start_token_id = tokenizer.bos_token_id\n",
+    "model.config.pad_token_id = tokenizer.pad_token_id"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "30e5b73c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Load model from decoder-only training.\n",
+    "model.load_state_dict(torch.load('decoder_only/de/00/model.pt'))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5466e908",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class AudioDataset(Dataset):\n",
+    "    \n",
+    "    def __init__(self, ds):\n",
+    "        self.ds = ds\n",
+    "        \n",
+    "    def __len__(self):\n",
+    "        return len(self.ds)\n",
+    "    \n",
+    "    def __getitem__(self, idx):\n",
+    "        eg = self.ds[idx]\n",
+    "        return eg['audio']['array'], eg['sentence']\n",
+    "    \n",
+    "def collate_fn(examples):\n",
+    "    # Remove the longest examples, should be only three and these may lead to OOM- or Index-Errors.\n",
+    "    examples = [eg for eg in examples if len(eg[0]) < 300_000]\n",
+    "    \n",
+    "    audio_features = wave2vec_extractor(\n",
+    "        [eg[0] for eg in examples], sampling_rate=16_000, return_tensors='pt', padding='longest'\n",
+    "    ).input_values\n",
+    "    \n",
+    "    input_ids = tokenizer(\n",
+    "        [eg[1] for eg in examples], return_tensors='pt', padding=True\n",
+    "    ).input_ids\n",
+    "    \n",
+    "    return audio_features, input_ids"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0453ccc1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train = common_voice['train'].cast_column('audio', Audio(sampling_rate=16_000))\n",
+    "val = common_voice['validation'].cast_column('audio', Audio(sampling_rate=16_000))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ad81c9ab",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "random.seed(419)\n",
+    "val_inds = list(range(len(common_voice['validation'])))\n",
+    "random.shuffle(val_inds)\n",
+    "\n",
+    "train_ds = AudioDataset(concatenate_datasets([train, val.select(val_inds[config['num_val_examples']:])]))\n",
+    "val_ds = AudioDataset(val.select(val_inds[:config['num_val_examples']]))\n",
+    "\n",
+    "train_dl = DataLoader(train_ds, batch_size=config['batch_size'], shuffle=True, collate_fn=collate_fn, num_workers=4)\n",
+    "val_dl = DataLoader(val_ds, batch_size=config['batch_size'], shuffle=False, collate_fn=collate_fn, num_workers=4)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f0d1c290",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "accelerator = Accelerator(fp16=True)\n",
+    "print(f'Using {accelerator.device}.')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2af1f2f1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "optimizer = AdamW(model.parameters(), lr=config['base_lr'], weight_decay=config['weight_decay'])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6921d32c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model, optimizer, train_dl, val_dl = accelerator.prepare(model, optimizer, train_dl, val_dl)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d699c404",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "with open(log_dir / 'config.json', 'w') as config_file:\n",
+    "    json.dump(config, config_file, indent=4)\n",
+    "    \n",
+    "writer = SummaryWriter(log_dir)\n",
+    "val_golds = common_voice['validation'].select(val_inds[:config['num_val_examples']])['sentence']\n",
+    "best_val_wer = 10.\n",
+    "global_train_step = 0\n",
+    "\n",
+    "for epoch in range(config['max_epochs']):\n",
+    "    model.train()\n",
+    "    for batch_step, (audio_features, input_ids) in enumerate(train_dl):\n",
+    "        global_train_step += 1\n",
+    "        \n",
+    "        out = model(labels=input_ids, input_values=audio_features)\n",
+    "        accelerator.backward(out.loss)\n",
+    "        writer.add_scalar('train_loss', out.loss.item(), global_train_step)\n",
+    "        \n",
+    "        if (batch_step + 1) % config['accumulate_grad'] == 0:\n",
+    "            optimizer.step()\n",
+    "            optimizer.zero_grad()\n",
+    "        if batch_step % 300 == 0:\n",
+    "            print(out.loss.item())\n",
+    "            \n",
+    "    model.eval()\n",
+    "    val_preds = []\n",
+    "    for audio_features, input_ids in val_dl:\n",
+    "        with torch.no_grad():\n",
+    "            generated = model.generate(audio_features)\n",
+    "        val_preds += tokenizer.batch_decode(generated)\n",
+    "    val_preds = [pred.lstrip('~').rstrip('_') for pred in val_preds]\n",
+    "    wer = calculate_wer(val_preds, val_golds)\n",
+    "    writer.add_scalar('val_wer', wer, epoch)\n",
+    "    print('WER: ', wer)\n",
+    "    \n",
+    "    if wer < best_val_wer:\n",
+    "        torch.save(model.state_dict(), model_dir / 'model.pt')\n",
+    "        print('Saved Model.')\n",
+    "        best_val_wer = wer"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.7"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

training/model.py

@@ -0,0 +1,137 @@
+import torch
+from transformers import SpeechEncoderDecoderModel
+from torch import nn
+from torch.nn import CrossEntropyLoss
+from transformers.models.encoder_decoder.modeling_encoder_decoder import shift_tokens_right
+from transformers.modeling_outputs import Seq2SeqLMOutput
+
+class Wav2VecGPT2Model(SpeechEncoderDecoderModel):
+    """
+    Basically the same as `SpeechEncoderDecoderModel` but position embeddings (initialized with GPT2's position
+    embeddings) are added to encoder output
+    """
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.encoder_outputs_pos_emb = nn.Embedding(1024, self.decoder.config.hidden_size)
+        with torch.no_grad():
+            self.encoder_outputs_pos_emb.weight.copy_(self.decoder.transformer.wpe.weight)
+        self.enc_to_dec_proj_ln = nn.LayerNorm(self.decoder.config.hidden_size,
+                                               eps=self.decoder.config.layer_norm_epsilon)
+
+    def __getattribute__(self, name):
+        # Fake class so it is recognized as seq2seq model.
+        if name == '__class__':
+            return SpeechEncoderDecoderModel
+        return SpeechEncoderDecoderModel.__getattribute__(self, name)
+
+    def forward(
+        self,
+        inputs=None,
+        attention_mask=None,
+        decoder_input_ids=None,
+        decoder_attention_mask=None,
+        encoder_outputs=None,
+        past_key_values=None,
+        decoder_inputs_embeds=None,
+        labels=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        input_values=None,
+        input_features=None,
+        return_dict=None,
+        **kwargs,
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        kwargs_encoder = {argument: value for argument, value in kwargs.items() if not argument.startswith("decoder_")}
+
+        kwargs_decoder = {
+            argument[len("decoder_") :]: value for argument, value in kwargs.items() if argument.startswith("decoder_")
+        }
+
+        if encoder_outputs is None and inputs is None:
+            if input_values is not None and input_features is not None:
+                raise ValueError("You cannot specify both input_values and input_features at the same time")
+            elif input_values is not None:
+                inputs = input_values
+            elif input_features is not None:
+                inputs = input_features
+            else:
+                raise ValueError("You have to specify either input_values or input_features")
+
+            encoder_outputs = self.encoder(
+                inputs,
+                attention_mask=attention_mask,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+                **kwargs_encoder,
+            )
+
+        encoder_hidden_states = encoder_outputs[0]
+
+        # optionally project encoder_hidden_states
+        if (
+            self.encoder_output_dim != self.decoder.config.hidden_size
+            and self.decoder.config.cross_attention_hidden_size is None
+        ):
+            # TODO: Truncate and warn if the sequence length is greater than 1024!
+            encoder_hidden_states = self.enc_to_dec_proj(encoder_hidden_states)
+            encoder_hidden_states += self.encoder_outputs_pos_emb(
+                torch.arange(0, encoder_hidden_states.shape[1], device=encoder_hidden_states.device)
+            )
+            encoder_hidden_states = self.enc_to_dec_proj_ln(encoder_hidden_states)
+
+        # compute correct encoder attention mask
+        if attention_mask is not None:
+            encoder_attention_mask = self.encoder._get_feature_vector_attention_mask(
+                encoder_hidden_states.shape[1], attention_mask
+            )
+        else:
+            encoder_attention_mask = None
+
+        if (labels is not None) and (decoder_input_ids is None and decoder_inputs_embeds is None):
+            decoder_input_ids = shift_tokens_right(
+                labels, self.config.pad_token_id, self.config.decoder_start_token_id
+            )
+
+        # Decode
+        decoder_outputs = self.decoder(
+            input_ids=decoder_input_ids,
+            attention_mask=decoder_attention_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            inputs_embeds=decoder_inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            use_cache=use_cache,
+            past_key_values=past_key_values,
+            return_dict=return_dict,
+            **kwargs_decoder,
+        )
+
+        # Compute loss independent from decoder (as some shift the logits inside them)
+        loss = None
+        if labels is not None:
+            logits = decoder_outputs.logits if return_dict else decoder_outputs[0]
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(logits.reshape(-1, self.decoder.config.vocab_size), labels.view(-1))
+
+        if not return_dict:
+            if loss is not None:
+                return (loss,) + decoder_outputs + encoder_outputs
+            else:
+                return decoder_outputs + encoder_outputs
+
+        return Seq2SeqLMOutput(
+            loss=loss,
+            logits=decoder_outputs.logits,
+            past_key_values=decoder_outputs.past_key_values,
+            decoder_hidden_states=decoder_outputs.hidden_states,
+            decoder_attentions=decoder_outputs.attentions,
+            cross_attentions=decoder_outputs.cross_attentions,
+            encoder_last_hidden_state=encoder_outputs[0],
+            encoder_hidden_states=getattr(encoder_outputs, 'hidden_states', None),  # TODO: only temporary (inconsistant)
+            encoder_attentions=getattr(encoder_outputs, 'attentions', None),
+        )

training/wer.py

@@ -0,0 +1,12 @@
+import jiwer
+
+def calculate_wer(predictions, golds):
+
+    transformation = jiwer.Compose([
+        jiwer.ToLowerCase(),
+        jiwer.RemoveWhiteSpace(replace_by_space=True),
+        jiwer.RemoveMultipleSpaces(),
+        jiwer.Strip(),
+        jiwer.ReduceToListOfListOfWords(word_delimiter=" ")
+    ])
+    return jiwer.wer(golds, predictions, truth_transform=transformation, hypothesis_transform=transformation)