asr_lab3_4.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ASR Labs 3 and 4 &ndash; Viterbi decoding\n",
    "\n",
    "In these labs we'll use the experience &ndash; and code &ndash; you have developed in previous labs to develop your own Viterbi decoder.  You'll want to refer back to [Lecture 5](https://www.inf.ed.ac.uk/teaching/courses/asr/2022-23/asr05-hmm-algorithms.pdf).  Remember that the Viterbi algorithm is used to find the joint probability of an observation sequence $X$ and the **best** single path $Q$, allowing this best path to be efficiently discovered. \n",
    "\n",
    "Your decoder will find the best path through the WFST representations HMMs that you developed in Labs 1 and 2.  In the case that your WFST is just a linear chain, the best path will provide an *alignment* of the observation sequence to the HMM states.  If your WFST is a set of word or phone loops (or any other structure) then the best path will allow you to recover the most likely transcription for the observed acoustic features, subject to the constraints of your grammar and vocabulary.\n",
    "\n",
    "Observation probabilities $b_j(t)$ will be supplied for you &ndash; you do not need to use the observations $(x_1, \\dotsc x_T)$ directly.  These are supplied via the `observation_model` module.  You can use this as follows:\n",
    "\n",
    "```python\n",
    "import observation_model\n",
    "\n",
    "my_om = observation_model.ObservationModel()\n",
    "\n",
    "my_om.load_audio('filename.wav')  \n",
    "\n",
    "# or use dummy audio for debugging\n",
    "my_om.load_dummy_audio()  # will generate dummy observations as seen in Lab 2, \n",
    "                          # useful for testing\n",
    "    \n",
    "my_om.observation_length()  # returns the sequence length, T  \n",
    "\n",
    "my_om.log_observation_probability(hmm_label, t)  # returns log b_j(t) given HMM label in string form\n",
    "                                                 # raises IndexError if t > T\n",
    "                                                 # raises KeyError if hmm_label is not known\n",
    "```\n",
    "\n",
    "\n",
    "It's easiest to write your decoder as a Python class, and we will supply a template below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: numpy in /afs/inf.ed.ac.uk/user/s20/s2070789/miniconda3/lib/python3.7/site-packages (1.20.1)\r\n"
     ]
    }
   ],
   "source": [
    "!pip install numpy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import observation_model\n",
    "import math\n",
    "import openfst_python as fst\n",
    "\n",
    "class MyViterbiDecoder:\n",
    "    \n",
    "    NLL_ZERO = 1e10  # define a constant representing -log(0).  This is not really infinite, but approximate\n",
    "                     # it here with a very large number\n",
    "    \n",
    "    def __init__(self, f, audio_file_name):\n",
    "        \"\"\"Set up the decoder class with an audio file and WFST f\n",
    "        \"\"\"\n",
    "        self.om = observation_model.ObservationModel()\n",
    "        self.f = f\n",
    "        \n",
    "        if audio_file_name:\n",
    "            self.om.load_audio(audio_file_name)\n",
    "        else:\n",
    "            self.om.load_dummy_audio()\n",
    "        \n",
    "        self.initialise_decoding()\n",
    "    \n",
    "    def initialise_decoding(self):\n",
    "        \"\"\"set up the values for V_j(0) (as negative log-likelihoods)\n",
    "        \n",
    "        \"\"\"\n",
    "        \n",
    "        self.V = []\n",
    "        for t in range(self.om.observation_length()+1):\n",
    "            self.V.append([self.NLL_ZERO]*self.f.num_states())\n",
    "        \n",
    "        # The above code means that self.V[t][j] for t = 0, ... T gives the Viterbi cost\n",
    "        # of state j, time t (in negative log-likelihood form)\n",
    "        # Initialising the costs to NLL_ZERO effectively means zero probability    \n",
    "        \n",
    "        # give the WFST start state a probability of 1.0   (NLL = 0.0)\n",
    "        self.V[0][f.start()] = 0.0\n",
    "        \n",
    "        # some WFSTs might have arcs with epsilon on the input (you might have already created \n",
    "        # examples of these in earlier labs) these correspond to non-emitting states, \n",
    "        # which means that we need to process them without stepping forward in time.  \n",
    "        # Don't worry too much about this!  \n",
    "        self.traverse_epsilon_arcs(0)\n",
    "        \n",
    "    def traverse_epsilon_arcs(self, t):\n",
    "        \"\"\"Traverse arcs with <eps> on the input at time t\n",
    "        \n",
    "        These correspond to transitions that don't emit an observation\n",
    "        \n",
    "        We've implemented this function for you as it's slightly trickier than\n",
    "        the normal case.  You might like to look at it to see what's going on, but\n",
    "        don't worry if you can't fully follow it.\n",
    "        \n",
    "        \"\"\"\n",
    "        \n",
    "        states_to_traverse = list(range(self.f.num_states())) # traverse all states\n",
    "        while states_to_traverse:\n",
    "            \n",
    "            # Set i to the ID of the current state, the first \n",
    "            # item in the list (and remove it from the list)\n",
    "            i = states_to_traverse.pop(0)   \n",
    "        \n",
    "            # don't bother traversing states which have zero probability\n",
    "            if self.V[t][i] == self.NLL_ZERO:\n",
    "                    continue\n",
    "        \n",
    "            for arc in self.f.arcs(i):\n",
    "                \n",
    "                if arc.ilabel == 0:     # if <eps> transition\n",
    "                  \n",
    "                    j = arc.nextstate   # ID of next state  \n",
    "                \n",
    "                    if self.V[t][j] > self.V[t][i] + float(arc.weight):\n",
    "                        \n",
    "                        # this means we've found a lower-cost path to\n",
    "                        # state j at time t.  We might need to add it\n",
    "                        # back to the processing queue.\n",
    "                        self.V[t][j] = self.V[t][i] + float(arc.weight)\n",
    "                  \n",
    "                        if j not in states_to_traverse:\n",
    "                            states_to_traverse.append(j)\n",
    "\n",
    "    \n",
    "    def forward_step(self, t):\n",
    "        \n",
    "        # TODO - exercise\n",
    "        pass\n",
    "    \n",
    "    def finalise_decoding(self):\n",
    "        \n",
    "        # TODO - exercise\n",
    "        pass\n",
    "    \n",
    "    def decode(self):\n",
    "        \n",
    "        self.initialise_decoding()\n",
    "        t = 1\n",
    "        while t <= self.om.observation_length():\n",
    "            self.forward_step(t)\n",
    "            self.traverse_epsilon_arcs(t)\n",
    "            t += 1\n",
    "        \n",
    "        self.finalise_decoding()\n",
    "    \n",
    "    def backtrace(self):\n",
    "        \n",
    "        # TODO - exercise \n",
    "        \n",
    "        # complete code to trace back through the\n",
    "        # best state sequence\n",
    "        \n",
    "        # You'll need to create a structure B_j(t) to store the \n",
    "        # back-pointers (see lectures), and amend the functions above to fill it.\n",
    "        best_state_sequence = []\n",
    "        return best_state_sequence\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copied from asr_lab2_solutions.ipynb\n",
    "# Just for you to debug your code\n",
    "\n",
    "def parse_lexicon(lex_file):\n",
    "    \"\"\"\n",
    "    Parse the lexicon file and return it in dictionary form.\n",
    "    \n",
    "    Args:\n",
    "        lex_file (str): filename of lexicon file with structure '<word> <phone1> <phone2>...'\n",
    "                        eg. peppers p eh p er z\n",
    "\n",
    "    Returns:\n",
    "        lex (dict): dictionary mapping words to list of phones\n",
    "    \"\"\"\n",
    "    \n",
    "    lex = {}  # create a dictionary for the lexicon entries (this could be a problem with larger lexica)\n",
    "    with open(lex_file, 'r') as f:\n",
    "        for line in f:\n",
    "            line = line.split()  # split at each space\n",
    "            lex[line[0]] = line[1:]  # first field the word, the rest is the phones\n",
    "    return lex\n",
    "\n",
    "def generate_symbol_tables(lexicon, n=3):\n",
    "    '''\n",
    "    Return word, phone and state symbol tables based on the supplied lexicon\n",
    "        \n",
    "    Args:\n",
    "        lexicon (dict): lexicon to use, created from the parse_lexicon() function\n",
    "        n (int): number of states for each phone HMM\n",
    "        \n",
    "    Returns:\n",
    "        word_table (fst.SymbolTable): table of words\n",
    "        phone_table (fst.SymbolTable): table of phones\n",
    "        state_table (fst.SymbolTable): table of HMM phone-state IDs\n",
    "    '''\n",
    "    \n",
    "    state_table = fst.SymbolTable()\n",
    "    phone_table = fst.SymbolTable()\n",
    "    word_table = fst.SymbolTable()\n",
    "    \n",
    "    # add empty <eps> symbol to all tables\n",
    "    state_table.add_symbol('<eps>')\n",
    "    phone_table.add_symbol('<eps>')\n",
    "    word_table.add_symbol('<eps>')\n",
    "    \n",
    "    for word, phones  in lexicon.items():\n",
    "        \n",
    "        word_table.add_symbol(word)\n",
    "        \n",
    "        for p in phones: # for each phone\n",
    "            \n",
    "            phone_table.add_symbol(p)\n",
    "            for i in range(1,n+1): # for each state 1 to n\n",
    "                state_table.add_symbol('{}_{}'.format(p, i))\n",
    "            \n",
    "    return word_table, phone_table, state_table\n",
    "\n",
    "\n",
    "# call these two functions\n",
    "lex = parse_lexicon('lexicon.txt')\n",
    "word_table, phone_table, state_table = generate_symbol_tables(lex)\n",
    "\n",
    "def generate_phone_wfst(f, start_state, phone, n):\n",
    "    \"\"\"\n",
    "    Generate a WFST representating an n-state left-to-right phone HMM\n",
    "    \n",
    "    Args:\n",
    "        f (fst.Fst()): an FST object, assumed to exist already\n",
    "        start_state (int): the index of the first state, assmed to exist already\n",
    "        phone (str): the phone label \n",
    "        n (int): number of states for each phone HMM\n",
    "        \n",
    "    Returns:\n",
    "        the final state of the FST\n",
    "    \"\"\"\n",
    "    \n",
    "    current_state = start_state\n",
    "    \n",
    "    for i in range(1, n+1):\n",
    "        \n",
    "        in_label = state_table.find('{}_{}'.format(phone, i))\n",
    "        \n",
    "        sl_weight = fst.Weight('log', -math.log(0.1))  # weight for self-loop\n",
    "        # self-loop back to current state\n",
    "        f.add_arc(current_state, fst.Arc(in_label, 0, sl_weight, current_state))\n",
    "        \n",
    "        # transition to next state\n",
    "        \n",
    "        # we want to output the phone label on the final state\n",
    "        # note: if outputting words instead this code should be modified\n",
    "        if i == n:\n",
    "            out_label = phone_table.find(phone)\n",
    "        else:\n",
    "            out_label = 0   # output empty <eps> label\n",
    "            \n",
    "        next_state = f.add_state()\n",
    "        next_weight = fst.Weight('log', -math.log(0.9)) # weight to next state\n",
    "        f.add_arc(current_state, fst.Arc(in_label, out_label, next_weight, next_state))    \n",
    "       \n",
    "        current_state = next_state\n",
    "        \n",
    "    return current_state\n",
    "\n",
    "def generate_word_wfst(word):\n",
    "    \"\"\" Generate a WFST for any word in the lexicon, composed of 3-state phone WFSTs.\n",
    "        This will currently output word labels.  \n",
    "        Exercise: could you modify this function and the one above to output a single phone label instead?\n",
    "    \n",
    "    Args:\n",
    "        word (str): the word to generate\n",
    "        \n",
    "    Returns:\n",
    "        the constructed WFST\n",
    "    \n",
    "    \"\"\"\n",
    "    f = fst.Fst('log')\n",
    "    \n",
    "    # create the start state\n",
    "    start_state = f.add_state()\n",
    "    f.set_start(start_state)\n",
    "    \n",
    "    current_state = start_state\n",
    "    \n",
    "    # iterate over all the phones in the word\n",
    "    for phone in lex[word]:   # will raise an exception if word is not in the lexicon\n",
    "        \n",
    "        current_state = generate_phone_wfst(f, current_state, phone, 3)\n",
    "    \n",
    "        # note: new current_state is now set to the final state of the previous phone WFST\n",
    "        \n",
    "    f.set_final(current_state)\n",
    "    \n",
    "    return f\n",
    "\n",
    "# example\n",
    "f = generate_word_wfst('peppers')\n",
    "f.set_input_symbols(state_table)\n",
    "f.set_output_symbols(phone_table)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercises &ndash; Lab 3\n",
    "\n",
    "The `__init__()`, `initialise_decoding()` and `decode()` functions have been completed in the template above.\n",
    "You should aim to complete the `forward_step()` and `finalise_decoding()` functions for Lab 3.  Don't worry about implementing the back-trace in Lab 3.\n",
    "\n",
    "You should draw on your solutions to Lab 2 &ndash; the main difference now is that, rather than simply sampling a single path and computing its likelihood, you'll need to compute and store, at every time step $t$, and for every state in the WFST, the likelihood of the best path reaching that state after $t$ time steps.  For how to do this using the Viterbi algorithm, see Lecture 5, slides 11 onwards.\n",
    "\n",
    "Test your algorithms on an WFST that recognises the word \"*pepper*\" and one that recognises any word in the vocabulary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here is the code for you to debug\n",
    "# To call the decoder (in a dummy example)\n",
    "# f is the WFST generated with the code from asr_lab2_solutions.ipynb\n",
    "decoder = MyViterbiDecoder(f, '')   # empty string '' just means use dummy probabilities for testing\n",
    "decoder.decode()\n",
    "# Check the last time step\n",
    "print(decoder.V[-1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With the dummy example and the word WFST for 'peppers', you should see the result as below,\n",
    "\n",
    "[10000000000.0, 10000000000.0, 10000000000.0, 10000000000.0, 10000000000.0, 10000000000.0, 10000000000.0, 10000000000.0, 10000000000.0, 10000000000.0, 10000000000.0, 10000000000.0, 10000000000.0, 10000000000.0, 10000000000.0, 71.71724553457089]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercises &ndash; Lab 4\n",
    "\n",
    "Now complete the `backtrace()` function to allow the best path to be recovered.  As noted in the code, you'll need to create a structure to store $B_j(t)$, which stores the identity of the best preceding state reaching state $j$ at time $t$ .  You can follow a similar method to storing $V_j(t)$, given in the code already.  You'll need to add it to the `initialise_decoding()`, `forward_step()` and `finalise_decoding()` functions.\n",
    "\n",
    "Once you are happy that your function works, you should amend your code so that you can also recover the sequence of *output* symbols on your WFST's best path as well.  This should allow you to produce your first word-recognition result!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here is the code for you to debug\n",
    "# To call the decoder (in a dummy example)\n",
    "# f is the WFST generated with the code from asr_lab2_solutions.ipynb\n",
    "decoder = MyViterbiDecoder(f, '')   # empty string '' just means use dummy probabilities for testing\n",
    "decoder.decode()\n",
    "# Check the backtrace result\n",
    "print(decoder.backtrace())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With the dummy example and the word WFST for 'peppers', you should see the result as below,\n",
    "\n",
    "([-1, 0, 1, 2, 3, 3, 3, 3, 4, 5, 5, 6, 7, 8, 8, 8, 9, 10, 11, 11, 11, 11, 11, 12, 13, 14, 14, 14, 14, 14, 14, 15], 'p eh p er z')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Working on real speech\n",
    "\n",
    "You can now find the first batch of recordings made by ASR students in `/group/teaching/asr/labs/recordings`.  You can test your decoder on real speech data by passing the full path to the WAV file when you create your `MyViterbiDecoder` object.  The transcriptions are also available in the same folder.\n",
    "\n",
    "When working with real speech, you may want to modify your code to print only the word output labels! \n",
    "\n",
    "If you are interested, the observation model (kindly supplied by Andrea) is a monophone time-delay neural network trained using the lattice-free MMI criterion, on the WSJ corpus of read speech."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}