sandbox_script.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a20ef248",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import datasets\n",
    "from transformers import Wav2Vec2CTCTokenizer, Wav2Vec2FeatureExtractor, Wav2Vec2Processor\n",
    "import string"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "id": "8fe38573",
   "metadata": {},
   "outputs": [],
   "source": [
    "# model_name_or_path = 'jonatasgrosman/wav2vec2-large-xlsr-53-english'\n",
    "model_name_or_path = 'ydshieh/wav2vec2-large-xlsr-53-chinese-zh-cn-gpt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "id": "83f27a41",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/samuel/anaconda3/envs/env_ascend/lib/python3.9/site-packages/transformers/configuration_utils.py:348: UserWarning: Passing `gradient_checkpointing` to a config initialization is deprecated and will be removed in v5 Transformers. Using `model.gradient_checkpointing_enable()` instead, or if you are using the `Trainer` API, pass `gradient_checkpointing=True` in your `TrainingArguments`.\n",
      "  warnings.warn(\n",
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "pretrained_tokenizer = Wav2Vec2CTCTokenizer.from_pretrained(model_name_or_path)\n",
    "pretrained_vocab = list(pretrained_tokenizer.get_vocab().keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "id": "de6e9877",
   "metadata": {},
   "outputs": [],
   "source": [
    "new_vocab_list = ['我', '去', '学', '爱', '你'] + list(string.ascii_lowercase) + list(string.ascii_uppercase) + list(string.digits)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "id": "6a023d48",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "len vocab dict 21157\n"
     ]
    }
   ],
   "source": [
    "all_vocab = list(dict.fromkeys(pretrained_vocab + new_vocab_list))\n",
    "vocab_dict = {v: k for k, v in enumerate(all_vocab)}\n",
    "\n",
    "def _assign_id_to_special_tokens(vocab_dict):\n",
    "    bos_token = \"<s>\"\n",
    "    eos_token = \"</s>\"\n",
    "    unk_token = \"[UNK]\"\n",
    "    pad_token = \"<pad>\"\n",
    "\n",
    "    if bos_token not in vocab_dict:\n",
    "        vocab_dict[bos_token] = len(vocab_dict)\n",
    "\n",
    "    if eos_token not in vocab_dict:\n",
    "        vocab_dict[eos_token] = len(vocab_dict)\n",
    "\n",
    "    if unk_token not in vocab_dict:\n",
    "        vocab_dict[unk_token] = len(vocab_dict)\n",
    "\n",
    "    if pad_token not in vocab_dict:\n",
    "        vocab_dict[pad_token] = len(vocab_dict)\n",
    "\n",
    "    return vocab_dict\n",
    "\n",
    "vocab_dict = _assign_id_to_special_tokens(vocab_dict)\n",
    "print(\"len vocab dict\", len(vocab_dict))\n",
    "\n",
    "with open(\"tmp/all_vocab.json\", \"w\") as vocab_file:\n",
    "    json.dump(vocab_dict, vocab_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "id": "278efb3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = Wav2Vec2CTCTokenizer(\"tmp/all_vocab.json\", unk_token=\"[UNK]\")\n",
    "feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(model_name_or_path)\n",
    "processor = Wav2Vec2Processor(feature_extractor=feature_extractor, tokenizer=tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "id": "d5058038",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = {y:x for x,y in pretrained_tokenizer.get_vocab().items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "id": "a3aad5c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "cache_dir_path = '/home/samuel/ascend-corpus/cache/ydshieh/wav2vec2-large-xlsr-53-chinese-zh-cn-gpt'\n",
    "vectorized_datasets = datasets.load_from_disk('{}/preprocess_data.arrow'.format(cache_dir_path))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 177,
   "id": "8cbad40b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# vocab_path = '/home/holy/projects/ascend-corpus/save/ydshieh/wav2vec2-large-xlsr-53-chinese-zh-cn-gpt/'\n",
    "tokenizer = Wav2Vec2CTCTokenizer(f\"tmp/all_vocab.json\", unk_token=\"[UNK]\", do_lower_case=False)\n",
    "# tokenizer = Wav2Vec2CTCTokenizer(f\"{vocab_path}/all_vocab.json\", unk_token=\"[UNK]\")\n",
    "# model_name_or_path = 'jonatasgrosman/wav2vec2-large-xlsr-53-english'\n",
    "# tokenizer = Wav2Vec2CTCTokenizer.from_pretrained(model_name_or_path)\n",
    "feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(model_name_or_path)\n",
    "processor = Wav2Vec2Processor(feature_extractor=feature_extractor, tokenizer=tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 179,
   "id": "9864394b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# tokenizer.tokens_trie = None\n",
    "tokenizer.tokens_trie.__dict__['data'] = {}\n",
    "# del tokenizer.tokens_trie.__dict__['data']['m']\n",
    "# del tokenizer.tokens_trie.__dict__['data']['i']\n",
    "# del tokenizer.tokens_trie.__dict__['data']['a']\n",
    "# del tokenizer.tokens_trie.__dict__['data']['n']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 180,
   "id": "0ead2a9a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['b', '|', 'b', 'i', 'p', '|', 'b', 'o', 'p', '|', 'b', 'a', 'p']"
      ]
     },
     "execution_count": 180,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.tokenize('b bip bop bap')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 167,
   "id": "2abb7f0c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['i', '|', 'a', 'm', '|', 'a', '|', 'm', 'a', 'n']"
      ]
     },
     "execution_count": 167,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.tokenize('i am a man')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 170,
   "id": "0f72c1e9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['I', '|', 'A', 'M', '|', 'A', '|', 'M', 'A', 'N']"
      ]
     },
     "execution_count": 170,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.tokenize('I AM A MAN')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "id": "c39aa1cd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['##骨##耦',\n",
       " '##骨##耦',\n",
       " '翹##孢##castsimple##其髒霜##馈玑英##初遺茵てした##淌##染##柒淪book玑##燔茵てした##姣髒##大##姣##荒##涤##龟##鸯##耦']"
      ]
     },
     "execution_count": 95,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "processor.batch_decode(vectorized_datasets['test'][:3]['labels'], group_tokens=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af1baa88",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "id": "6372a9e3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[5876,\n",
       " 20049,\n",
       " 15480,\n",
       " 7295,\n",
       " 18985,\n",
       " 17220,\n",
       " 18121,\n",
       " 18831,\n",
       " 18440,\n",
       " 18121,\n",
       " 1548,\n",
       " 5876,\n",
       " 14576,\n",
       " 18510]"
      ]
     },
     "execution_count": 73,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vectorized_datasets['train'][0]['labels']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03f658f1",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2c226358",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using custom data configuration CAiRE--ASCEND-cfec671ab34ecfe4\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading and preparing dataset csv/CAiRE--ASCEND to /home/samuel/.cache/huggingface/datasets/csv/CAiRE--ASCEND-cfec671ab34ecfe4/0.0.0/6b9057d9e23d9d8a2f05b985917a0da84d70c5dae3d22ddd8a3f22fb01c69d9e...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cdf36b07f7c6433fa0a1423881faca47",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "98085ab229ca455d95c4ee2c5500a575",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset csv downloaded and prepared to /home/samuel/.cache/huggingface/datasets/csv/CAiRE--ASCEND-cfec671ab34ecfe4/0.0.0/6b9057d9e23d9d8a2f05b985917a0da84d70c5dae3d22ddd8a3f22fb01c69d9e. Subsequent calls will reuse this data.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e6b63250535b4b958a421cf5be037efb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "dataset = load_dataset(\"CAiRE/ASCEND\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "cfa8fa70",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "csv-test.arrow\t csv-validation.arrow  LICENSE\n",
      "csv-train.arrow  dataset_info.json\n"
     ]
    }
   ],
   "source": [
    "!ls /home/samuel/.cache/huggingface/datasets/csv/CAiRE___ASCEND-cfec671ab34ecfe4/0.0.0/bf68a4c4aefa545d0712b2fcbb1b327f905bbe2f6425fbc5e8c25234acb9e14a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "20300228",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mSignature:\u001b[0m\n",
       "\u001b[0mdataset\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mfunction\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mwith_indices\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mbool\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0minput_columns\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mUnion\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mstr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mList\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mstr\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mNoneType\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mbatched\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mbool\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mbatch_size\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mOptional\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mint\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m1000\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mremove_columns\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mUnion\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mstr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mList\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mstr\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mNoneType\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mkeep_in_memory\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mbool\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mload_from_cache_file\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mbool\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mcache_file_names\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mOptional\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mDict\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mstr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mOptional\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mstr\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mwriter_batch_size\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mOptional\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mint\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m1000\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mfeatures\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mOptional\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mdatasets\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfeatures\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfeatures\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mFeatures\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mdisable_nullable\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mbool\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mfn_kwargs\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mOptional\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mdict\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mnum_proc\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mOptional\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mint\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mdesc\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mOptional\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mstr\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0;34m'DatasetDict'\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
       "\u001b[0;31mSource:\u001b[0m   \n",
       "    \u001b[0;32mdef\u001b[0m \u001b[0mmap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0mfunction\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0mwith_indices\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mbool\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0minput_columns\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mOptional\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mUnion\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mstr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mList\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mstr\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0mbatched\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mbool\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0mbatch_size\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mOptional\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mint\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m1000\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0mremove_columns\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mOptional\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mUnion\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mstr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mList\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mstr\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0mkeep_in_memory\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mbool\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0mload_from_cache_file\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mbool\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0mcache_file_names\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mOptional\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mDict\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mstr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mOptional\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mstr\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0mwriter_batch_size\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mOptional\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mint\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m1000\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0mfeatures\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mOptional\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mFeatures\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0mdisable_nullable\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mbool\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0mfn_kwargs\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mOptional\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mdict\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0mnum_proc\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mOptional\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mint\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0mdesc\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mOptional\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mstr\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0;34m\"DatasetDict\"\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0;34m\"\"\"Apply a function to all the elements in the table (individually or in batches)\u001b[0m\n",
       "\u001b[0;34m        and update the table (if function does updated examples).\u001b[0m\n",
       "\u001b[0;34m        The transformation is applied to all the datasets of the dataset dictionary.\u001b[0m\n",
       "\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m        Args:\u001b[0m\n",
       "\u001b[0;34m            function (`callable`): with one of the following signature:\u001b[0m\n",
       "\u001b[0;34m                - `function(example: Dict) -> Union[Dict, Any]` if `batched=False` and `with_indices=False`\u001b[0m\n",
       "\u001b[0;34m                - `function(example: Dict, indices: int) -> Union[Dict, Any]` if `batched=False` and `with_indices=True`\u001b[0m\n",
       "\u001b[0;34m                - `function(batch: Dict[List]) -> Union[Dict, Any]` if `batched=True` and `with_indices=False`\u001b[0m\n",
       "\u001b[0;34m                - `function(batch: Dict[List], indices: List[int]) -> Union[Dict, Any]` if `batched=True` and `with_indices=True`\u001b[0m\n",
       "\u001b[0;34m            with_indices (`bool`, defaults to `False`): Provide example indices to `function`. Note that in this case the signature of `function` should be `def function(example, idx): ...`.\u001b[0m\n",
       "\u001b[0;34m            input_columns (`Optional[Union[str, List[str]]]`, defaults to `None`): The columns to be passed into `function` as\u001b[0m\n",
       "\u001b[0;34m                positional arguments. If `None`, a dict mapping to all formatted columns is passed as one argument.\u001b[0m\n",
       "\u001b[0;34m            batched (`bool`, defaults to `False`): Provide batch of examples to `function`\u001b[0m\n",
       "\u001b[0;34m            batch_size (`Optional[int]`, defaults to `1000`): Number of examples per batch provided to `function` if `batched=True`\u001b[0m\n",
       "\u001b[0;34m                `batch_size <= 0` or `batch_size == None`: Provide the full dataset as a single batch to `function`\u001b[0m\n",
       "\u001b[0;34m            remove_columns (`Optional[Union[str, List[str]]]`, defaults to `None`): Remove a selection of columns while doing the mapping.\u001b[0m\n",
       "\u001b[0;34m                Columns will be removed before updating the examples with the output of `function`, i.e. if `function` is adding\u001b[0m\n",
       "\u001b[0;34m                columns with names in `remove_columns`, these columns will be kept.\u001b[0m\n",
       "\u001b[0;34m            keep_in_memory (`bool`, defaults to `False`): Keep the dataset in memory instead of writing it to a cache file.\u001b[0m\n",
       "\u001b[0;34m            load_from_cache_file (`bool`, defaults to `True`): If a cache file storing the current computation from `function`\u001b[0m\n",
       "\u001b[0;34m                can be identified, use it instead of recomputing.\u001b[0m\n",
       "\u001b[0;34m            cache_file_names (`Optional[Dict[str, str]]`, defaults to `None`): Provide the name of a path for the cache file. It is used to store the\u001b[0m\n",
       "\u001b[0;34m                results of the computation instead of the automatically generated cache file name.\u001b[0m\n",
       "\u001b[0;34m                You have to provide one :obj:`cache_file_name` per dataset in the dataset dictionary.\u001b[0m\n",
       "\u001b[0;34m            writer_batch_size (:obj:`int`, default `1000`): Number of rows per write operation for the cache file writer.\u001b[0m\n",
       "\u001b[0;34m                This value is a good trade-off between memory usage during the processing, and processing speed.\u001b[0m\n",
       "\u001b[0;34m                Higher value makes the processing do fewer lookups, lower value consume less temporary memory while running `.map()`.\u001b[0m\n",
       "\u001b[0;34m            features (`Optional[datasets.Features]`, defaults to `None`): Use a specific Features to store the cache file\u001b[0m\n",
       "\u001b[0;34m                instead of the automatically generated one.\u001b[0m\n",
       "\u001b[0;34m            disable_nullable (`bool`, defaults to `False`): Disallow null values in the table.\u001b[0m\n",
       "\u001b[0;34m            fn_kwargs (`Optional[Dict]`, defaults to `None`): Keyword arguments to be passed to `function`\u001b[0m\n",
       "\u001b[0;34m            num_proc (`Optional[int]`, defaults to `None`): Number of processes for multiprocessing. By default it doesn't\u001b[0m\n",
       "\u001b[0;34m                use multiprocessing.\u001b[0m\n",
       "\u001b[0;34m            desc (`Optional[str]`, defaults to `None`): Meaningful description to be displayed alongside with the progress bar while mapping examples.\u001b[0m\n",
       "\u001b[0;34m        \"\"\"\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_check_values_type\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0;32mif\u001b[0m \u001b[0mcache_file_names\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m            \u001b[0mcache_file_names\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m{\u001b[0m\u001b[0mk\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;32mNone\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mk\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m}\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0;32mreturn\u001b[0m \u001b[0mDatasetDict\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m            \u001b[0;34m{\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m                \u001b[0mk\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mdataset\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m                    \u001b[0mfunction\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mfunction\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m                    \u001b[0mwith_indices\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mwith_indices\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m                    \u001b[0minput_columns\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0minput_columns\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m                    \u001b[0mbatched\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mbatched\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m                    \u001b[0mbatch_size\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mbatch_size\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m                    \u001b[0mremove_columns\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mremove_columns\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m                    \u001b[0mkeep_in_memory\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mkeep_in_memory\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m                    \u001b[0mload_from_cache_file\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mload_from_cache_file\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m                    \u001b[0mcache_file_name\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mcache_file_names\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mk\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m                    \u001b[0mwriter_batch_size\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mwriter_batch_size\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m                    \u001b[0mfeatures\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mfeatures\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m                    \u001b[0mdisable_nullable\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdisable_nullable\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m                    \u001b[0mfn_kwargs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mfn_kwargs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m                    \u001b[0mnum_proc\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mnum_proc\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m                    \u001b[0mdesc\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdesc\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m                \u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m                \u001b[0;32mfor\u001b[0m \u001b[0mk\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdataset\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mitems\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m            \u001b[0;34m}\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m        \u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
       "\u001b[0;31mFile:\u001b[0m      ~/anaconda3/envs/env_ascend/lib/python3.9/site-packages/datasets/dataset_dict.py\n",
       "\u001b[0;31mType:\u001b[0m      method\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "??dataset.map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "cc798356",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import logging\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import argparse\n",
    "\n",
    "import torchaudio\n",
    "import torch\n",
    "import re\n",
    "import json \n",
    "import librosa\n",
    "from datasets import load_from_disk, load_dataset, load_metric\n",
    "\n",
    "from transformers import (\n",
    "    Wav2Vec2Processor, \n",
    "    Wav2Vec2CTCTokenizer,\n",
    "    Wav2Vec2FeatureExtractor,\n",
    "    Wav2Vec2ForCTC,\n",
    "    Trainer,\n",
    "    TrainingArguments,\n",
    "    HfArgumentParser,\n",
    "    EarlyStoppingCallback\n",
    ")\n",
    "\n",
    "from datasets import DatasetDict, load_metric, load_from_disk\n",
    "from dataclasses import dataclass, field\n",
    "from typing import Any, Dict, List, Optional, Union\n",
    "import datasets\n",
    "import pickle\n",
    "\n",
    "import editdistance\n",
    "import jieba\n",
    "from itertools import chain"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba9a2c6f",
   "metadata": {},
   "source": [
    "## Test Metric \n",
    "notes: WER from jiwer feels a bit weird, we just implement our own MER and CER"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "56a960d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenize_for_mer(text):\n",
    "    tokens = list(filter(lambda tok: len(tok.strip()) > 0, jieba.lcut(text)))\n",
    "    tokens = [[tok] if tok.isascii() else list(tok) for tok in tokens]\n",
    "    return list(chain(*tokens))\n",
    "\n",
    "def tokenize_for_cer(text):\n",
    "    tokens = list(filter(lambda tok: len(tok.strip()) > 0, list(text)))\n",
    "    return tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "cd8bf073",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Building prefix dict from the default dictionary ...\n",
      "Loading model from cache /tmp/jieba.cache\n",
      "Loading model cost 0.651 seconds.\n",
      "Prefix dict has been built successfully.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "0.5"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pred = tokenize_for_mer('我爱你 anak gembala')\n",
    "ref = tokenize_for_mer('我爱你 agum')\n",
    "editdistance.distance(pred, ref) / len(ref)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "83f49b48",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.1428571428571428"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pred = tokenize_for_cer('我爱你 anak gembala')\n",
    "ref = tokenize_for_cer('我爱你 agum')\n",
    "editdistance.distance(pred, ref) / len(ref)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5efae62c",
   "metadata": {},
   "source": [
    "# Code starts here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "fb72eac6",
   "metadata": {},
   "outputs": [],
   "source": [
    "#####\n",
    "# Common Functions\n",
    "#####\n",
    "CHARS_TO_IGNORE = [\",\", \"?\", \"¿\", \".\", \"!\", \"¡\", \";\", \"；\", \":\", '\"\"', \"%\", '\"', \"�\", \"ʿ\", \"·\", \"჻\", \"~\", \"՞\",\n",
    "                   \"؟\", \"،\", \"।\", \"॥\", \"«\", \"»\", \"„\", \"“\", \"”\", \"「\", \"」\", \"‘\", \"’\", \"《\", \"》\", \"(\", \")\", \"[\", \"]\",\n",
    "                   \"{\", \"}\", \"=\", \"`\", \"_\", \"+\", \"<\", \">\", \"…\", \"–\", \"°\", \"´\", \"ʾ\", \"‹\", \"›\", \"©\", \"®\", \"—\", \"→\", \"。\",\n",
    "                   \"、\", \"﹂\", \"﹁\", \"‧\", \"～\", \"﹏\", \"，\", \"｛\", \"｝\", \"（\", \"）\", \"［\", \"］\", \"【\", \"】\", \"‥\", \"〽\",\n",
    "                   \"『\", \"』\", \"〝\", \"〟\", \"⟨\", \"⟩\", \"〜\", \"：\", \"！\", \"？\", \"♪\", \"؛\", \"/\", \"\\\\\", \"º\", \"−\", \"^\", \"ʻ\", \"ˆ\"]\n",
    "\n",
    "def remove_special_characters(batch):\n",
    "    batch[\"sentence\"] = re.sub(chars_to_ignore_regex, '', batch[\"sentence\"]).lower() + \" \"\n",
    "    return batch\n",
    "\n",
    "def extract_all_chars(batch):\n",
    "    all_text = \" \".join(batch[\"sentence\"])\n",
    "    vocab = list(set(all_text))\n",
    "    return {\"vocab\": [vocab], \"all_text\": [all_text]}\n",
    "\n",
    "#####\n",
    "# Data Loading Function\n",
    "#####\n",
    "def speech_file_to_array_fn(batch):\n",
    "    speech_array, sampling_rate = torchaudio.load(batch[\"audio_path\"])\n",
    "    batch[\"speech_sample\"] = speech_array[0].numpy()\n",
    "    batch[\"sampling_rate\"] = sampling_rate\n",
    "    return batch\n",
    "\n",
    "def load_dataset(manifest_file, num_proc):\n",
    "    batches = {\"path\": [], \"text\": [], \"target_sample_rate\": []}\n",
    "    base_path = '/'.join(manifest_file.split('/')[:-1])\n",
    "    \n",
    "    manifest_df = pd.read_csv(manifest_file)\n",
    "    manifest_df = manifest_df.rename({'text': 'target_text'}, axis=1)\n",
    "    manifest_df['audio_path'] = manifest_df['audio_path'].apply(lambda path: f'{base_path}/{path}')\n",
    "        \n",
    "    batches = Dataset.from_pandas(manifest_df)\n",
    "    batches = batches.map(speech_file_to_array_fn, num_proc=num_proc)\n",
    "    return batches\n",
    "\n",
    "#####\n",
    "# Data Collator\n",
    "#####\n",
    "@dataclass\n",
    "class DataCollatorCTCWithPadding:\n",
    "    \"\"\"\n",
    "    Data collator that will dynamically pad the inputs received.\n",
    "    Args:\n",
    "        processor (:class:`~transformers.Wav2Vec2Processor`)\n",
    "            The processor used for proccessing the data.\n",
    "        padding (:obj:`bool`, :obj:`str` or :class:`~transformers.tokenization_utils_base.PaddingStrategy`, `optional`, defaults to :obj:`True`):\n",
    "            Select a strategy to pad the returned sequences (according to the model's padding side and padding index)\n",
    "            among:\n",
    "            * :obj:`True` or :obj:`'longest'`: Pad to the longest sequence in the batch (or no padding if only a single\n",
    "              sequence if provided).\n",
    "            * :obj:`'max_length'`: Pad to a maximum length specified with the argument :obj:`max_length` or to the\n",
    "              maximum acceptable input length for the model if that argument is not provided.\n",
    "            * :obj:`False` or :obj:`'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of\n",
    "              different lengths).\n",
    "        max_length (:obj:`int`, `optional`):\n",
    "            Maximum length of the ``input_values`` of the returned list and optionally padding length (see above).\n",
    "        max_length_labels (:obj:`int`, `optional`):\n",
    "            Maximum length of the ``labels`` returned list and optionally padding length (see above).\n",
    "        pad_to_multiple_of (:obj:`int`, `optional`):\n",
    "            If set will pad the sequence to a multiple of the provided value.\n",
    "            This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability >=\n",
    "            7.5 (Volta).\n",
    "    \"\"\"\n",
    "\n",
    "    processor: Wav2Vec2Processor\n",
    "    padding: Union[bool, str] = True\n",
    "    max_length: Optional[int] = None\n",
    "    max_length_labels: Optional[int] = None\n",
    "    pad_to_multiple_of: Optional[int] = None\n",
    "    pad_to_multiple_of_labels: Optional[int] = None\n",
    "\n",
    "    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:\n",
    "        # split inputs and labels since they have to be of different lenghts and need\n",
    "        # different padding methods\n",
    "        input_features = [{\"input_values\": feature[\"input_values\"]} for feature in features]\n",
    "        label_features = [{\"input_ids\": feature[\"labels\"]} for feature in features]\n",
    "\n",
    "        batch = self.processor.pad(\n",
    "            input_features,\n",
    "            padding=self.padding,\n",
    "            max_length=self.max_length,\n",
    "            pad_to_multiple_of=self.pad_to_multiple_of,\n",
    "            return_tensors=\"pt\",\n",
    "        )\n",
    "        with self.processor.as_target_processor():\n",
    "            labels_batch = self.processor.pad(\n",
    "                label_features,\n",
    "                padding=self.padding,\n",
    "                max_length=self.max_length_labels,\n",
    "                pad_to_multiple_of=self.pad_to_multiple_of_labels,\n",
    "                return_tensors=\"pt\",\n",
    "            )\n",
    "\n",
    "        # replace padding with -100 to ignore loss correctly\n",
    "        labels = labels_batch[\"input_ids\"].masked_fill(labels_batch.attention_mask.ne(1), -100)\n",
    "        batch[\"labels\"] = labels\n",
    "\n",
    "        return batch\n",
    "\n",
    "#####\n",
    "# Compute Metric Function\n",
    "#####\n",
    "wer_metric = load_metric(\"wer\")\n",
    "cer_metric = load_metric(\"cer\")\n",
    "\n",
    "def tokenize_for_mer(text):\n",
    "    tokens = list(filter(lambda tok: len(tok.strip()) > 0, jieba.lcut(text)))\n",
    "    tokens = [[tok] if tok.isascii() else list(tok) for tok in tokens]\n",
    "    return list(chain(*tokens))\n",
    "\n",
    "def tokenize_for_cer(text):\n",
    "    tokens = list(filter(lambda tok: len(tok.strip()) > 0, list(text)))\n",
    "    return tokens\n",
    "\n",
    "def compute_metrics(pred):\n",
    "    pred_logits = pred.predictions\n",
    "    pred_ids = np.argmax(pred_logits, axis=-1)\n",
    "    pred.label_ids[pred.label_ids == -100] = processor.tokenizer.pad_token_id\n",
    "\n",
    "    pred_strs = processor.batch_decode(pred_ids)\n",
    "\n",
    "    # we do not want to group tokens when computing the metrics\n",
    "    label_strs = processor.batch_decode(pred.label_ids, group_tokens=False)\n",
    "    mixed_distance, mixed_tokens = 0, 0\n",
    "    char_distance, char_tokens = 0, 0\n",
    "    for pred_str, label_str in zip(pred_strs, label_strs)\n",
    "        # Calculate \n",
    "        m_pred = tokenize_for_mer(pred_str)\n",
    "        m_ref = tokenize_for_mer(label_str)\n",
    "        mixed_distance += editdistance.distance(m_pred, m_ref)\n",
    "        mixed_tokens += len(m_ref)\n",
    "    \n",
    "        c_pred = tokenize_for_cer(pred_str)\n",
    "        c_ref = tokenize_for_cer(label_str)\n",
    "        char_distance += editdistance.distance(c_pred, c_ref)\n",
    "        char_tokens += len(c_ref)\n",
    "    mer = mixed_distance / mixed_tokens\n",
    "    cer = char_distance / char_tokens\n",
    "#     wer = wer_metric.compute(predictions=pred_str, references=label_str)\n",
    "#     cer = cer_metric.compute(predictions=pred_str, references=label_str)\n",
    "\n",
    "    return {\"mer\": mer, \"cer\": cer}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "337f7dff",
   "metadata": {},
   "outputs": [],
   "source": [
    "#####\n",
    "# Main Functions\n",
    "#####\n",
    "def run(model_args, data_args, training_args):\n",
    "    ###\n",
    "    # Prepare Dataset\n",
    "    ###\n",
    "    raw_datasets = DatasetDict()\n",
    "    raw_datasets[\"train\"] = load_dataset(data_args.train_manifest_path, data_args.num_proc)\n",
    "    raw_datasets[\"valid\"] = load_dataset(data_args.valid_manifest_path, data_args.num_proc)\n",
    "    raw_datasets[\"test\"] = load_dataset(data_args.test_manifest_path, data_args.num_proc)\n",
    "\n",
    "    ###\n",
    "    # Prepare Processor & Model    \n",
    "    ###\n",
    "    processor = Wav2Vec2Processor.from_pretrained(model_args.model_name_or_path)\n",
    "    model = Wav2Vec2ForCTC.from_pretrained(model_args.model_name_or_path)\n",
    "    model.cuda()\n",
    "    \n",
    "    ###\n",
    "    # Preprocessing datasets\n",
    "    ###\n",
    "    # Remove ignorable characters\n",
    "    chars_to_ignore_regex = f\"[{re.escape(''.join(CHARS_TO_IGNORE))}]\"\n",
    "    def remove_special_characters(batch):\n",
    "        if chars_to_ignore_regex is not None:\n",
    "            batch[\"target_text\"] = re.sub(chars_to_ignore_regex, \"\", batch[data_args.text_column_name]).lower() + \" \"\n",
    "        else:\n",
    "            batch[\"target_text\"] = batch[data_args.text_column_name].lower() + \" \"\n",
    "        return batch\n",
    "\n",
    "    with training_args.main_process_first(desc=\"dataset map special characters removal\"):\n",
    "        raw_datasets = raw_datasets.map(\n",
    "            remove_special_characters,\n",
    "            remove_columns=[data_args.text_column_name],\n",
    "            desc=\"remove special characters from datasets\",\n",
    "        )\n",
    "        \n",
    "    # Preprocess audio sample and label text\n",
    "    def prepare_dataset(batch):\n",
    "        # Preprocess audio\n",
    "        batch[\"input_values\"] = processor(batch[\"speech_sample\"]).input_values[0]\n",
    "\n",
    "        # Preprocess text\n",
    "        with processor.as_target_processor():\n",
    "            batch[\"labels\"] = processor(batch[\"target_text\"]).input_ids\n",
    "            \n",
    "        return batch\n",
    "\n",
    "    with training_args.main_process_first(desc=\"dataset map preprocessing\"):\n",
    "        vectorized_datasets = raw_datasets.map(\n",
    "            prepare_dataset,\n",
    "            remove_columns=raw_datasets[\"train\"].column_names,\n",
    "            num_proc=data_args.preprocessing_num_workers,\n",
    "            desc=\"preprocess datasets\",\n",
    "        )\n",
    "\n",
    "    if data_args.preprocessing_only:\n",
    "        logger.info(f\"Data preprocessing finished. Files cached at {vectorized_datasets.cache_files}\")\n",
    "        return\n",
    "    \n",
    "    ###\n",
    "    # Prepare Data Collator and Trainer\n",
    "    ###\n",
    "    # Instantiate custom data collator\n",
    "    data_collator = DataCollatorCTCWithPadding(processor=processor)\n",
    "\n",
    "    # Initialize Trainer\n",
    "    trainer = Trainer(\n",
    "        model=model,\n",
    "        data_collator=data_collator,\n",
    "        args=training_args,\n",
    "        compute_metrics=compute_metrics,\n",
    "        train_dataset=vectorized_datasets[\"train\"] if training_args.do_train else None,\n",
    "        eval_dataset=vectorized_datasets[\"valid\"] if training_args.do_eval else None,\n",
    "        tokenizer=processor.feature_extractor,\n",
    "    )\n",
    "\n",
    "    ###\n",
    "    # Training Phase\n",
    "    ###\n",
    "    if training_args.do_train:\n",
    "        # use last checkpoint if exist\n",
    "        if last_checkpoint is not None:\n",
    "            checkpoint = last_checkpoint\n",
    "        elif os.path.isdir(model_args.model_name_or_path):\n",
    "            checkpoint = model_args.model_name_or_path\n",
    "        else:\n",
    "            checkpoint = None\n",
    "\n",
    "        # Save the feature_extractor and the tokenizer\n",
    "        if is_main_process(training_args.local_rank):\n",
    "            processor.save_pretrained(training_args.output_dir)\n",
    "\n",
    "        train_result = trainer.train(resume_from_checkpoint=checkpoint)\n",
    "        trainer.save_model()\n",
    "\n",
    "        metrics = train_result.metrics\n",
    "        max_train_samples = (\n",
    "            data_args.max_train_samples\n",
    "            if data_args.max_train_samples is not None\n",
    "            else len(vectorized_datasets[\"train\"])\n",
    "        )\n",
    "        metrics[\"train_samples\"] = min(max_train_samples, len(vectorized_datasets[\"train\"]))\n",
    "\n",
    "        trainer.log_metrics(\"train\", metrics)\n",
    "        trainer.save_metrics(\"train\", metrics)\n",
    "        trainer.save_state()\n",
    "\n",
    "    ###\n",
    "    # Evaluation Phase\n",
    "    ###\n",
    "    results = {}\n",
    "    if training_args.do_eval:\n",
    "        logger.info(\"*** Evaluate ***\")\n",
    "        metrics = trainer.evaluate(eval_dataset=vectorized_datasets[\"test\"])\n",
    "        metrics[\"eval_samples\"] = len(vectorized_datasets[\"test\"])\n",
    "\n",
    "        trainer.log_metrics(\"eval\", metrics)\n",
    "        trainer.save_metrics(\"eval\", metrics)\n",
    "\n",
    "    # Write model card and (optionally) push to hub\n",
    "    kwargs = {\n",
    "        \"finetuned_from\": model_args.model_name_or_path,\n",
    "        \"tasks\": \"speech-recognition\",\n",
    "        \"tags\": [\"automatic-speech-recognition\", \"ASCEND\"],\n",
    "        \"dataset_args\": f\"Config: na\",\n",
    "        \"dataset\": f\"ASCEND\",\n",
    "        \"language\": 'zh-en'\n",
    "    }\n",
    "\n",
    "    if training_args.push_to_hub:\n",
    "        trainer.push_to_hub(**kwargs)\n",
    "    else:\n",
    "        trainer.create_model_card(**kwargs)\n",
    "\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8d0df5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def main():\n",
    "    ###\n",
    "    # Parsing & Initialization\n",
    "    ###\n",
    "    # Parse argument\n",
    "    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))\n",
    "    if len(sys.argv) == 2 and sys.argv[1].endswith(\".json\"):\n",
    "        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))\n",
    "    else:\n",
    "        model_args, data_args, training_args = parser.parse_args_into_dataclasses()\n",
    "\n",
    "    # Set random seed\n",
    "    set_seed(training_args.seed)\n",
    "    \n",
    "    # Detect last checkpoint\n",
    "    last_checkpoint = None\n",
    "    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:\n",
    "        last_checkpoint = get_last_checkpoint(training_args.output_dir)\n",
    "        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:\n",
    "            raise ValueError(\n",
    "                f\"Output directory ({training_args.output_dir}) already exists and is not empty. \"\n",
    "                \"Use --overwrite_output_dir to overcome.\"\n",
    "            )\n",
    "        elif last_checkpoint is not None:\n",
    "            logger.info(\n",
    "                f\"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change \"\n",
    "                \"the `--output_dir` or add `--overwrite_output_dir` to train from scratch.\"\n",
    "            )\n",
    "\n",
    "    ###\n",
    "    # Prepare logger\n",
    "    ###\n",
    "    # Init logging\n",
    "    logging.basicConfig(\n",
    "        format=\"%(asctime)s - %(levelname)s - %(name)s - %(message)s\",\n",
    "        datefmt=\"%m/%d/%Y %H:%M:%S\",\n",
    "        handlers=[logging.StreamHandler(sys.stdout)],\n",
    "    )\n",
    "    logger.setLevel(logging.INFO if is_main_process(training_args.local_rank) else logging.WARN)\n",
    "\n",
    "    # Log on each process the small summary:\n",
    "    logger.warning(\n",
    "        f\"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}\"\n",
    "        f\"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}\"\n",
    "    )\n",
    "    # Set the verbosity to info of the Transformers logger (on main process only):\n",
    "    if is_main_process(training_args.local_rank):\n",
    "        transformers.utils.logging.set_verbosity_info()\n",
    "    logger.info(\"Training/evaluation parameters %s\", training_args)\n",
    "    \n",
    "    ###\n",
    "    # RUN RUN RUN!!!\n",
    "    ###\n",
    "    run(model_args, data_args, training_args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "598ecb37",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env_ascend",
   "language": "python",
   "name": "env_ascend"
  },
  "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
}