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](http://www.inf.ed.ac.uk/teaching/courses/asr/2020-21/asr05-hmm-algorithms.pdf).  Remember that the Viterbi algorithm is used to find the joint probability of an observation sequence $X$ and the **best** path 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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install numpy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "partner = 's1608733'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "lex = parse_lexicon('lexicon.txt')\n",
    "lex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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 symbol 𝜖 to all tables\n",
    "    state_table.add_symbol('𝜖')\n",
    "    phone_table.add_symbol('𝜖')\n",
    "    word_table.add_symbol('𝜖')\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",
    "word_table, phone_table, state_table = generate_symbol_tables(lex)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_phone_wfst(f, start_state, phone, n):\n",
    "    \"\"\"\n",
    "    Generate a WFST representing 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, assumed to exist already\n",
    "        phone (str): the phone label \n",
    "        n (int): number of states of the 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",
    "        # self-loop back to current state\n",
    "        f.add_arc(current_state, fst.Arc(in_label, 0, fst.Weight('log', -math.log(0.1)), 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 = phone_table.find('𝜖')   # output empty label 𝜖\n",
    "            \n",
    "        next_state = f.add_state()\n",
    "        f.add_arc(current_state, fst.Arc(in_label, out_label, fst.Weight('log', -math.log(0.9)), next_state))\n",
    "       \n",
    "        current_state = next_state\n",
    "    return current_state\n",
    "\n",
    "f = fst.Fst('log')\n",
    "start = f.add_state()\n",
    "f.set_start(start)\n",
    "\n",
    "last_state = generate_phone_wfst(f, start, 'p', 3)\n",
    "\n",
    "f.set_input_symbols(state_table)\n",
    "f.set_output_symbols(phone_table)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_word_wfst(f, start_state, word, n):\n",
    "    \"\"\" Generate a WFST for any word in the lexicon, composed of n-state phone WFSTs.\n",
    "        This will currently output phone labels.\n",
    "    \n",
    "    Args:\n",
    "        f (fst.Fst()): an FST object, assumed to exist already\n",
    "        start_state (int): the index of the first state, assumed to exist already\n",
    "        word (str): the word to generate\n",
    "        n (int): states per phone HMM\n",
    "        \n",
    "    Returns:\n",
    "        the constructed WFST\n",
    "    \n",
    "    \"\"\"\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",
    "        # your code here\n",
    "        \n",
    "        current_state = generate_phone_wfst(f, current_state, phone, n)\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",
    "# f = fst.Fst('log')\n",
    "# start = f.add_state()\n",
    "# f.set_start(start)\n",
    "\n",
    "# generate_word_wfst(f, start, 'peppers', 3)\n",
    "# f.set_input_symbols(state_table)\n",
    "# f.set_output_symbols(phone_table)\n",
    "\n",
    "# from subprocess import check_call\n",
    "# from IPython.display import Image\n",
    "# f.draw('tmp.dot', portrait=True)\n",
    "# check_call(['dot','-Tpng','-Gdpi=200','tmp.dot','-o','tmp.png'])\n",
    "# Image(filename='tmp.png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import openfst_python as fst\n",
    "import observation_model\n",
    "import math\n",
    "import numpy as np\n",
    "\n",
    "class MyViterbiDecoder:\n",
    "    \n",
    "    NLL_ZERO = 1e10  # define a constant representing -log(0).  This is 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",
    "        \n",
    "    def initialise_decoding(self):\n",
    "        \"\"\"set up the values for V_j(0) (as negative log-likelihoods)\n",
    "        \n",
    "        \"\"\"\n",
    "        \n",
    "        self.V = []   # stores likelihood along best path reaching state j\n",
    "        self.B = []   # stores identity of best previous state reaching state j\n",
    "        self.W = []   # stores output labels sequence along arc reaching j - this removes need for \n",
    "                      # extra code to read the output sequence along the best path\n",
    "        \n",
    "        for t in range(self.om.observation_length()+1):\n",
    "            self.V.append([self.NLL_ZERO]*self.f.num_states())\n",
    "            self.B.append([-1]*self.f.num_states())\n",
    "            self.W.append([[] for i in range(self.f.num_states())])  #  multiplying the empty list doesn't make multiple\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][self.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(self.f.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",
    "                        # save backtrace information.  In the case of an epsilon transition, \n",
    "                        # we save the identity of the best state at t-1.  This means we may not\n",
    "                        # be able to fully recover the best path, but to do otherwise would\n",
    "                        # require a more complicated way of storing backtrace information\n",
    "                        self.B[t][j] = self.B[t][i] \n",
    "                        \n",
    "                        # and save the output labels encountered - this is a list, because\n",
    "                        # there could be multiple output labels (in the case of <eps> arcs)\n",
    "                        if arc.olabel != 0:\n",
    "                            self.W[t][j] = self.W[t][i] + [arc.olabel]\n",
    "                        else:\n",
    "                            self.W[t][j] = self.W[t][i]\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",
    "        for i in self.f.states():\n",
    "            \n",
    "            if not self.V[t-1][i] == self.NLL_ZERO:   # no point in propagating states with zero probability\n",
    "                \n",
    "                for arc in self.f.arcs(i):\n",
    "                    \n",
    "                    if arc.ilabel != 0: # <eps> transitions don't emit an observation\n",
    "                        j = arc.nextstate\n",
    "                        tp = float(arc.weight)  # transition prob\n",
    "                        ep = -self.om.log_observation_probability(self.f.input_symbols().find(arc.ilabel), t)  # emission negative log prob\n",
    "                        prob = tp + ep + self.V[t-1][i] # they're logs\n",
    "                        if prob < self.V[t][j]:\n",
    "                            self.V[t][j] = prob\n",
    "                            self.B[t][j] = i\n",
    "                            \n",
    "                            # store the output labels encountered too\n",
    "                            if arc.olabel !=0:\n",
    "                                self.W[t][j] = [arc.olabel]\n",
    "                            else:\n",
    "                                self.W[t][j] = []\n",
    "                            \n",
    "    \n",
    "    def finalise_decoding(self):\n",
    "        \"\"\" this incorporates the probability of terminating at each state\n",
    "        \"\"\"\n",
    "        \n",
    "        for state in self.f.states():\n",
    "            final_weight = float(self.f.final(state))\n",
    "            if self.V[-1][state] != self.NLL_ZERO:\n",
    "                if final_weight == math.inf:\n",
    "                    self.V[-1][state] = self.NLL_ZERO  # effectively says that we can't end in this state\n",
    "                else:\n",
    "                    self.V[-1][state] += final_weight\n",
    "                    \n",
    "        # get a list of all states where there was a path ending with non-zero probability\n",
    "        finished = [x for x in self.V[-1] if x < self.NLL_ZERO]\n",
    "        if not finished:  # if empty\n",
    "            print(\"No path got to the end of the observations.\")\n",
    "        \n",
    "        \n",
    "    def decode(self):\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",
    "        self.finalise_decoding()\n",
    "    \n",
    "    def backtrace(self):\n",
    "        \n",
    "        best_state_sequence, output_sequence = [], []\n",
    "        \n",
    "        # no backtrace at start state, iterate from t = 1\n",
    "        for t, _ in enumerate(self.B[1:], 1):\n",
    "            \n",
    "            # update backtrace until start state is reached\n",
    "            state = self.B[-t][-1] if t == 1 else self.B[-t][state]\n",
    "\n",
    "            # store best path so far at current timestep\n",
    "            best_state_sequence.append(state)\n",
    "\n",
    "            # store output phones\n",
    "            if self.W[-t][state + 1]: # skip epsilon transitions (empty list)\n",
    "                phone = phone_table.find(self.W[-t][state + 1][0])\n",
    "                \n",
    "                # add phone if not a self-transition/loop\n",
    "                if not output_sequence or output_sequence[-1] != phone:\n",
    "                    output_sequence.append(phone)\n",
    "\n",
    "        return best_state_sequence[::-1], output_sequence[::-1]\n",
    "\n",
    "f = fst.Fst('log')\n",
    "start = f.add_state()\n",
    "f.set_start(start)\n",
    "\n",
    "generate_word_wfst(f, start, 'peppers', 3)\n",
    "f.set_input_symbols(state_table)\n",
    "f.set_output_symbols(phone_table)\n",
    "\n",
    "# to call the decoder (in a dummy example)\n",
    "# f will be a WFST that you have created in a previous lab\n",
    "decoder = MyViterbiDecoder(f, '')   # empty string '' just means use dummy probabilities for testing\n",
    "decoder.decode()\n",
    "print(decoder.backtrace())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "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.\n",
    "\n",
    "\n",
    "## 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!\n",
    "\n",
    "### 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.\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}