asr_lab2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ASR Lab 2 - Computing HMM probabilities\n",
    "\n",
    "To begin with, we'll use your function to generate a Word WFST for the word \"*peppers*\", using `generate_word_wfst('peppers')`.  By viewing this as an HMM, you'll be able to sample possible paths through the model and also generate the likelihood of an observation sequence $(x_1, \\dotsc, x_T)$.\n",
    "\n",
    "We'll build on this to implement the basics of the Viterbi algorithm, which can later be used for word recognition.\n",
    "\n",
    "First, copy your code from Lab 1 into the space below.  You can use the official solutions if you like.\n",
    "If you want to extract the code-only parts of your previous notebook, on the terminal command line you can type:\n",
    "\n",
    "```bash\n",
    "jupyter nbconvert --to python <notebook-name.ipynb>\n",
    "```\n",
    "\n",
    "where <notebook-name.ipynb> indicates the path of the notebook file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "partner = 's1608733'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openfst_python as fst\n",
    "import math\n",
    "\n",
    "input_sym = fst.SymbolTable()\n",
    "output_sym = fst.SymbolTable()\n",
    "\n",
    "input_sym.add_symbol('𝜖') # by convention, 𝜖 always\n",
    "                              # has symbol zero\n",
    "input_sym.add_symbol('a')  # input symbols\n",
    "input_sym.add_symbol('b')\n",
    "\n",
    "output_sym.add_symbol('𝜖')  # output symbols\n",
    "output_sym.add_symbol('d')\n",
    "output_sym.add_symbol('c')\n",
    "\n",
    "f = fst.Fst('log')\n",
    "\n",
    "f.set_input_symbols(input_sym)\n",
    "f.set_output_symbols(output_sym)\n",
    "\n",
    "s0 = f.add_state()\n",
    "s1 = f.add_state()\n",
    "s2 = f.add_state()\n",
    "s3 = f.add_state()\n",
    "\n",
    "a = input_sym.find('a')\n",
    "b = input_sym.find('b')\n",
    "c = output_sym.find('c')\n",
    "d = output_sym.find('d')\n",
    "\n",
    "f.add_arc(s0, fst.Arc(a, c, fst.Weight('log', -math.log(0.5)), s1))\n",
    "f.add_arc(s0, fst.Arc(b, d, fst.Weight('log', -math.log(0.5)), s2))\n",
    "f.add_arc(s1, fst.Arc(a, c, fst.Weight('log', -math.log(1)), s3))\n",
    "f.add_arc(s2, fst.Arc(b, d, fst.Weight('log', -math.log(1)), s3))\n",
    "\n",
    "f.set_start(s0)\n",
    "f.set_final(s3)"
   ]
  },
  {
   "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_parallel_path_wfst(f, start_state, n):\n",
    "    \"\"\"\n",
    "    Generate a WFST representing an n-state parallel-path left-to-right 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",
    "        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",
    "    next_state = f.add_state()\n",
    "    for i in range(n):\n",
    "\n",
    "        # self-loop back to current state\n",
    "        f.add_arc(current_state, fst.Arc(0, 0, fst.Weight('log', -math.log(0.1)), current_state))\n",
    "        # only one arc out of penultimate state, else two\n",
    "        weight = 0.9 if i == n-1 else 0.45\n",
    "        f.add_arc(current_state, fst.Arc(0, 0, fst.Weight('log', -math.log(weight)), next_state))\n",
    "        if i != n-1:\n",
    "            next_next_state = f.add_state()\n",
    "            f.add_arc(current_state, fst.Arc(0, 0, fst.Weight('log', -math.log(0.45)), next_next_state))\n",
    "        current_state = next_state\n",
    "        next_state = next_next_state\n",
    "\n",
    "    return current_state\n",
    "\n",
    "\n",
    "f = fst.Fst('log')\n",
    "start = f.add_state()\n",
    "f.set_start(start)\n",
    "\n",
    "last_state = generate_parallel_path_wfst(f, start, 4)\n",
    "f"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_ergodic_wfst(f, start_state, n):\n",
    "    \"\"\"\n",
    "    Generate a WFST representing an n-state ergodic 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",
    "        n (int): number of states of the HMM excluding start and end\n",
    "        \n",
    "    Returns:\n",
    "        the final state of the FST\n",
    "    \"\"\"    \n",
    "    \n",
    "    current_state = start_state\n",
    "    \n",
    "    for i in range(n):\n",
    "        f.add_state()\n",
    "        \n",
    "    for i in range(n+1): # +1 is start state\n",
    "        for j in range(n+1):\n",
    "            # weight is uniformly distributed for all arcs except self-transitions\n",
    "            weight = 0.1 if i == j else 0.9 / (n - 1)\n",
    "            f.add_arc(i, fst.Arc(0, 0, fst.Weight('log', -math.log(weight)), j))\n",
    "\n",
    "    return current_state\n",
    "\n",
    "\n",
    "f = fst.Fst('log')\n",
    "start = f.add_state()\n",
    "f.set_start(start)\n",
    "\n",
    "last_state = generate_ergodic_wfst(f, start, 5)\n",
    "f\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": {},
   "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": {},
   "outputs": [],
   "source": [
    "def generate_phone_recognition_wfst(n):\n",
    "    \"\"\" generate a HMM to recognise any single phone in the lexicon\n",
    "    \n",
    "    Args:\n",
    "        n (int): states per phone HMM\n",
    "\n",
    "    Returns:\n",
    "        the constructed WFST\n",
    "    \n",
    "    \"\"\"\n",
    "    \n",
    "    f = fst.Fst('log')\n",
    "    \n",
    "    # create a single start state\n",
    "    start_state = f.add_state()\n",
    "    f.set_start(start_state)\n",
    "    \n",
    "    # get a list of all the phones in the lexicon\n",
    "    # there are lots of way to do this.  Here, we use the set() object\n",
    "\n",
    "    # will contain all unique phones in the lexicon\n",
    "    phone_set = set()\n",
    "    \n",
    "    for pronunciation in lex.values():\n",
    "        phone_set = phone_set.union(pronunciation)\n",
    "    \n",
    "    # weight is uniformly distributed for all phones\n",
    "    weight = 1 / len(phone_set)\n",
    "    \n",
    "    for phone in phone_set:\n",
    "        \n",
    "        # we need to add an empty arc from the start state to where the actual phone HMM\n",
    "        # will begin.  If you can't see why this is needed, try without it!\n",
    "        current_state = f.add_state()\n",
    "        f.add_arc(start_state, fst.Arc(0, 0, fst.Weight('log', -math.log(weight)), current_state))\n",
    "    \n",
    "        end_state = generate_phone_wfst(f, current_state, phone, n)\n",
    "    \n",
    "        f.set_final(end_state)\n",
    "\n",
    "    return f\n",
    "\n",
    "f = generate_phone_recognition_wfst(3)\n",
    "f.set_input_symbols(state_table)\n",
    "f.set_output_symbols(phone_table)\n",
    "\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": {},
   "outputs": [],
   "source": [
    "def generate_phone_sequence_recognition_wfst(n):\n",
    "    \"\"\" generate a HMM to recognise any single phone sequence in the lexicon\n",
    "    \n",
    "    Args:\n",
    "        n (int): states per phone HMM\n",
    "\n",
    "    Returns:\n",
    "        the constructed WFST\n",
    "    \n",
    "    \"\"\"\n",
    "    \n",
    "    f = fst.Fst('log')\n",
    "    \n",
    "    # create a single start state\n",
    "    start_state = f.add_state()\n",
    "    f.set_start(start_state)\n",
    "    \n",
    "    phone_set = set()\n",
    "    \n",
    "    for pronunciation in lex.values():\n",
    "        phone_set = phone_set.union(pronunciation)\n",
    "\n",
    "    # weight is uniformly distributed for all phones\n",
    "    weight = 1 / len(phone_set)\n",
    "        \n",
    "    for phone in phone_set:\n",
    "        current_state = f.add_state()\n",
    "        f.add_arc(start_state, fst.Arc(0, 0, fst.Weight('log', -math.log(weight)), current_state))\n",
    "    \n",
    "        end_state = generate_phone_wfst(f, current_state, phone, n)\n",
    "        \n",
    "        f.add_arc(end_state, fst.Arc(0,0, fst.Weight('log', -math.log(1)), start_state))\n",
    "        f.set_final(end_state)\n",
    "\n",
    "    return f\n",
    "\n",
    "f = generate_phone_sequence_recognition_wfst(3)\n",
    "f.set_input_symbols(state_table)\n",
    "f.set_output_symbols(phone_table)\n",
    "\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": {},
   "outputs": [],
   "source": [
    "def generate_word_sequence_recognition_wfst(n):\n",
    "    \"\"\" generate a HMM to recognise any single word sequence for words in the lexicon\n",
    "    \n",
    "    Args:\n",
    "        n (int): states per phone HMM\n",
    "\n",
    "    Returns:\n",
    "        the constructed WFST\n",
    "    \n",
    "    \"\"\"\n",
    "    \n",
    "    f = fst.Fst('log')\n",
    "    \n",
    "    # create a single start state\n",
    "    start_state = f.add_state()\n",
    "    f.set_start(start_state)\n",
    "    \n",
    "    # weight is uniformly distributed for all word sequences\n",
    "    weight = 1 / len(lex.keys())\n",
    "\n",
    "    for word, phones in lex.items():\n",
    "        current_state = f.add_state()\n",
    "        f.add_arc(start_state, fst.Arc(0, 0, fst.Weight('log', -math.log(weight)), current_state))\n",
    "        \n",
    "        for phone in phones: \n",
    "            current_state = generate_phone_wfst(f, current_state, phone, 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",
    "        f.add_arc(current_state, fst.Arc(0, 0, fst.Weight('log', -math.log(0.1)), start_state))\n",
    "        \n",
    "    return f\n",
    "\n",
    "f = generate_word_sequence_recognition_wfst(3)\n",
    "f.set_input_symbols(state_table)\n",
    "f.set_output_symbols(phone_table)\n",
    "\n",
    "f.draw('tmp.dot', portrait=True)\n",
    "check_call(['dot','-Tpng','-Gdpi=300','tmp.dot','-o','tmp.png'])\n",
    "Image(filename='tmp.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that the WFST has been constructed, we can traverse over the states and arcs.  This example (taken from [OpenFst](http://www.openfst.org/twiki/bin/view/FST/PythonExtension)) shows how you can do this:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for state in f.states():\n",
    "    \n",
    "    # iterate over all arcs leaving this state    \n",
    "    for arc in f.arcs(state):\n",
    "         print(state, arc.ilabel, arc.olabel, arc.weight, arc.nextstate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Alternatively, we could begin at the start state, and traverse in a depth-first manner.  **Warning**: the code below specifically handles self-loops, but won't work if your WFST has larger cycles in it!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def traverse_arcs(state):\n",
    "    \"\"\"Traverse every arc leaving a particular state\n",
    "    \"\"\"\n",
    "    for arc in f.arcs(state):\n",
    "        print(state, arc.ilabel, arc.olabel, arc.weight, arc.nextstate)\n",
    "        \n",
    "        if arc.nextstate != state:   # don't follow the self-loops or we'll get stuck forever!\n",
    "            traverse_arcs(arc.nextstate)\n",
    "\n",
    "s = f.start()\n",
    "traverse_arcs(s)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For a more readable table, you could find the indexes of the input and output labels in your symbol tables and print the string instead."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercises\n",
    "\n",
    "1. Write code to randomly generate (sample) a path through your word HMM for \"*peppers*\".  You should output the sequence of input and output labels along the path.  To sample from a list of arcs, you can use code like\n",
    "\n",
    "```python\n",
    "import random\n",
    "\n",
    "arc_list = list(f.arcs(state))\n",
    "sampled_arc = random.sample(arc_list,1)[0]\n",
    "```\n",
    "\n",
    "  Notice that if you repeat your random sampling by running the code multiple times, you'll get paths of different lengths due to the self-loops\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import random\n",
    "\n",
    "def sample_random_path(f):\n",
    "    '''Given an FST, randomly sample a path through it.\n",
    "    \n",
    "        Args:\n",
    "            f (fst.Fst()): an FST\n",
    "        \n",
    "        Returns:\n",
    "            input_label_seq (list(str)): the list of input labels from the arcs that were sampled\n",
    "            output_label_seq (list(str)): the list of output labels from the arcs that were sampled\n",
    "        '''\n",
    "    curr_state = f.start() # start from beginning\n",
    "    weight_type = f.weight_type() # type of weights used in the fst\n",
    "    input_label_seq = []\n",
    "    output_label_seq = []\n",
    "\n",
    "    while f.final(curr_state) == fst.Weight(weight_type, 'inf'): # the .final method returns the probability of a state being final\n",
    "                                                             # it's infinite when the state is NOT final\n",
    "        # get all candidate arcs from current state\n",
    "        arc_list = list(f.arcs(curr_state))\n",
    "        # select one (word) randomly\n",
    "        sampled_arc = random.sample(arc_list,1)[0]\n",
    "        # store input label (phone_state) and output label (phone)\n",
    "        input_label_seq.append(state_table.find(sampled_arc.ilabel))\n",
    "        output_label_seq.append(phone_table.find(sampled_arc.olabel))\n",
    "        # advance to next state\n",
    "        curr_state = sampled_arc.nextstate\n",
    "\n",
    "    return input_label_seq, output_label_seq\n",
    "\n",
    "f = fst.Fst('log')\n",
    "start = f.add_state()\n",
    "f.set_start(start)\n",
    "generate_word_wfst(f, start, 'peppers', 3)\n",
    "\n",
    "input_label_seq, output_label_seq = sample_random_path(f)\n",
    "\n",
    "print('\\n'.join(['{} {}'.format(input_label_seq[i], output_label_seq[i]) for i in range(len(input_label_seq))]))\n",
    "\n",
    "# for ilab, olab in zip(input_label_seq, output_label_seq):\n",
    "#     print(ilab, state_table.find(ilab))\n",
    "#     print(olab, phone_table.find(olab))\n",
    "#     print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2. Now it's time to add probabilities to your WFST.  As mentioned at the end of Lab 1, probabilities in WFSTs are traditionally expressed in negative log format, that is, the weight $w$ on an arc transitioning between states $i$ and $j$ is given by $w=-\\log a_{ij}$, where $a_{ij}$ is the HMM transition probability.  Remember that you can add weights using the third argument to `fst.Arc()`.\n",
    "\n",
    "  You should now modify your code above to add weights to your word and phone recognition WFSTs from Lab 1, corresponding to transition probabilities.  Assume that the probability of a self-loop is $0.1$, and that when transitioning *between* separate multiple sets of phones (or words), the probabilities are uniform over all transitions.\n",
    "\n",
    "  Remember to set your fst to use log probabilities and use log weights:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "import math\n",
    "f = fst.Fst('log')\n",
    "\n",
    "s1 = f.add_state()\n",
    "s2 = f.add_state()\n",
    "weight = fst.Weight('log', -math.log(0.1))\n",
    "f.add_arc(s1, fst.Arc(0, 0, weight, s2))\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "3. Modify your answer to exercise 1 to sample a path through the word HMM *and* also compute the negative log probability of the path.  This gives you $-\\log p(Q)$ in the lecture notation.  (Recall that $\\log ab = \\log a + \\log b$)\n",
    "\n",
    "  **Note**: Internally OpenFst stores weights in a special object that you will need to convert to a float, using the `float()` function, before adding your negative log probabilities.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_random_path_prob(f):\n",
    "    '''Given an FST, randomly sample a path through it and compute the negative log probability.\n",
    "    \n",
    "        Args:\n",
    "            f (fst.Fst()): an FST\n",
    "        \n",
    "        Returns:\n",
    "            input_label_seq (list(str)): the list of input labels from the arcs that were sampled\n",
    "            output_label_seq (list(str)): the list of output labels from the arcs that were sampled\n",
    "            neg_log_prob (float): negative log probability of the sampled path\n",
    "        '''\n",
    "    curr_state = f.start() # start from beginning\n",
    "    weight_type = f.weight_type() # type of weights used in the fst\n",
    "    input_label_seq = []\n",
    "    output_label_seq = []\n",
    "    neg_log_prob = 0\n",
    "\n",
    "    while f.final(curr_state) == fst.Weight(weight_type, 'inf'): # the .final method returns the probability of a state being final\n",
    "                                                             # it's infinite when the state is NOT final\n",
    "        # get all candidate arcs from current state\n",
    "        arc_list = list(f.arcs(curr_state))\n",
    "        # select one (word) randomly\n",
    "        sampled_arc = random.sample(arc_list,1)[0]\n",
    "        # store input label (phone_state) and output label (phone)\n",
    "        input_label_seq.append(state_table.find(sampled_arc.ilabel))\n",
    "        output_label_seq.append(phone_table.find(sampled_arc.olabel))\n",
    "        # accumulate neg log prob\n",
    "        neg_log_prob += float(sampled_arc.weight)\n",
    "        # advance to next state\n",
    "        curr_state = sampled_arc.nextstate\n",
    "            \n",
    "    return input_label_seq, output_label_seq, neg_log_prob\n",
    "\n",
    "input_label_seq, output_label_seq, neg_log_prob = sample_random_path_prob(f)\n",
    "\n",
    "print('\\n'.join(['{} {}'.format(input_label_seq[i], output_label_seq[i]) for i in range(len(input_label_seq))]))\n",
    "print(neg_log_prob)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "4. You are now given a set of observations, ($x_1, \\dotsc, x_t, \\dotsc$).  Can you use your WFST for the word \"*peppers*\" to compute $p(X,Q)$ for a randomly sampled path $Q$ through the HMM?  For now, we won't use real samples $x_t$, and will instead assume that you already have a function `observation_probability(state, t)` that computes $b_j(t) = p(x_t|q_t=j)$, provided here:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def observation_probability(hmm_label, t):\n",
    "    \"\"\" Computes b_j(t) where j is the current state\n",
    "    \n",
    "    This is just a dummy version!  In later labs we'll generate \n",
    "    probabilities for real speech frames.\n",
    "    \n",
    "    You don't need to look at this function in detail.\n",
    "    \n",
    "    Args: hmm_label (str): the HMM state label, j.  We'll use string form: \"p_1\", \"p_2\", \"eh_1\" etc  \n",
    "          t (int) : current time step, starting at 1\n",
    "          \n",
    "    Returns: \n",
    "          p (float): the observation probability p(x_t | q_t = hmm_label)\n",
    "    \"\"\"\n",
    "    \n",
    "    p = {} # dictionary of probabilities\n",
    "    \n",
    "    assert(t>0)\n",
    "    \n",
    "    # this is just a simulation!\n",
    "    if t < 4:\n",
    "        p = {'p_1': 1.0, 'p_2':1.0, 'p_3': 1.0, 'eh_1':0.2}\n",
    "    elif t < 9:\n",
    "        p = {'p_3': 0.5, 'eh_1':1.0, 'eh_2': 1.0, 'eh_3': 1.0}\n",
    "    elif t < 13:\n",
    "        p = {'eh_3': 1.0, 'p_1': 1.0, 'p_2': 1.0, 'p_3':1.0, 'er_1':0.5}\n",
    "    elif t < 18:\n",
    "        p = {'p_3': 1.0, 'er_1': 1.0, 'er_2': 1.0, 'er_3':0.7}\n",
    "    elif t < 25:\n",
    "        p = {'er_3': 1.0, 'z_1': 1.0, 'z_2': 1.0, 'z_3':1.0}\n",
    "    else:\n",
    "        p = {'z_2': 0.5, 'z_3': 1.0}\n",
    "        \n",
    "    for label in ['p_1', 'p_2', 'p_3', 'eh_1', 'eh_2', 'eh_3', 'er_1', 'er_2', 'er_3', 'z_1', 'z_2', 'z_3']:        \n",
    "        if label not in p:\n",
    "            p[label] = 0.01  # give all other states a small probability to avoid zero probability\n",
    "            \n",
    "    # normalise the probabilities:\n",
    "    scale = sum(p.values())\n",
    "    for k in p:\n",
    "        p[k] = p[k]/scale\n",
    "        \n",
    "    return p[hmm_label]\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Enter your code below.  You might want to convert the observation probabilities into negative log probabilities.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_random_path_obs_prob(f):\n",
    "    '''Given an FST and observation probabilities, randomly sample a path\n",
    "        through it and compute the negative log probability.\n",
    "    \n",
    "        Args:\n",
    "            f (fst.Fst()): an FST\n",
    "        \n",
    "        Returns:\n",
    "            input_label_seq (list(str)): the list of input labels from the arcs that were sampled\n",
    "            output_label_seq (list(str)): the list of output labels from the arcs that were sampled\n",
    "            neg_log_prob (float): negative log probability of the sampled path\n",
    "        '''\n",
    "    t = 1\n",
    "    curr_state = f.start() # start from beginning\n",
    "    weight_type = f.weight_type() # type of weights used in the fst\n",
    "    input_label_seq = []\n",
    "    output_label_seq = []\n",
    "    neg_log_prob = 0.0 # log(1) = 0\n",
    "\n",
    "    while f.final(curr_state) == fst.Weight(weight_type, 'inf'):\n",
    "        \n",
    "        # get all candidate arcs from current state\n",
    "        arc_list = list(f.arcs(curr_state))\n",
    "        # select one (word) randomly\n",
    "        sampled_arc = random.sample(arc_list,1)[0]\n",
    "        # store input label (phone_state) and output label (phone)\n",
    "        input_label_seq.append(state_table.find(sampled_arc.ilabel))\n",
    "        output_label_seq.append(phone_table.find(sampled_arc.olabel))\n",
    "        # accumulate neg log prob\n",
    "        neg_log_prob += float(sampled_arc.weight)\n",
    "        # accumulate observation probability b_j(x)\n",
    "        neg_log_prob += -math.log(observation_probability(state_table.find(sampled_arc.ilabel), t))\n",
    "        # advance to next state\n",
    "        curr_state = sampled_arc.nextstate\n",
    "        \n",
    "        t += 1\n",
    "    \n",
    "    return input_label_seq, output_label_seq, neg_log_prob\n",
    "\n",
    "input_label_seq, output_label_seq, neg_log_prob = sample_random_path_obs_prob(f)\n",
    "print('\\n'.join(['{} {}'.format(input_label_seq[i], output_label_seq[i]) for i in range(len(input_label_seq))]))\n",
    "print(neg_log_prob)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You might have noticed that the dummy observation probability function above effectively allows the observation sequence $x_t$ to be arbitrarily long.  This is simply to allow it to match the length of your sampled path $Q$.  In real use, the observation sequence will have a fixed length $T$, and any matching path through the HMM will have to have the same length.  We'll explore this more when writing the Viterbi decoder in the next lab.\n",
    "\n",
    "## If you have more time\n",
    "\n",
    "You might like to start thinking about how to implement the Viterbi algorithm over HMMs in WFST form.  Try working with the \"*peppers*\" example above.  You'll need to write functions to compute and store the probabilities $V_j(t)$, giving the probability up to time step $t$ of the observation sequence $(x_1, \\dotsc, x_t)$ along the most likely path $(q_1, \\dotsc, q_t)$."
   ]
  },
  {
   "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
}