# Chatbot
We are going to use the [Reformer](https://arxiv.org/abs/2001.04451), also known as the efficient Transformer, to generate a dialogue between two bots. You will feed conversations to your model and it will learn how to understand the context of each one. Not only will it learn how to answer questions but it will also know how to ask questions if it needs more info. For example, after a customer asks for a train ticket, the chatbot can ask what time the said customer wants to leave. You can use this concept to automate call centers, hotel receptions, personal trainers, or any type of customer service.

* Understand how the Reformer works
* Explore the [MultiWoz](https://arxiv.org/abs/1810.00278) dataset
* Process the data to feed it into the model
* Train your model
* Generate a dialogue by feeding a question to the model

```python

from google.colab import drive
drive.mount('/content/drive')
```

    Mounted at /content/drive

# Part 1:   Exploring the MultiWoz dataset

We will start by exploring the MultiWoz dataset. The dataset we are about to use has more than 10,000 human annotated dialogues and spans multiple domains and topics. Some dialogues include multiple domains and others include single domains.

Let's first import the modules we will be using:

```python
import json
import random
import numpy as np
from termcolor import colored
!pip install trax
import trax
from trax import layers as tl
from trax.supervised import training
!pip list | grep trax
```

    Collecting trax
      Downloading https://files.pythonhosted.org/packages/85/1d/c0a3aeed127c26a0c3f0925fc9cc7278c272e52310eedfc322477e854972/trax-1.3.6-py2.py3-none-any.whl (468kB)
         |████████████████████████████████| 471kB 5.6MB/s
    Requirement already satisfied: gin-config in /usr/local/lib/python3.6/dist-packages (from trax) (0.3.0)
    Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from trax) (1.18.5)
    Collecting tensorflow-text
      Downloading https://files.pythonhosted.org/packages/28/b2/2dbd90b93913afd07e6101b8b84327c401c394e60141c1e98590038060b3/tensorflow_text-2.3.0-cp36-cp36m-manylinux1_x86_64.whl (2.6MB)
         |████████████████████████████████| 2.6MB 16.1MB/s
    Requirement already satisfied: scipy in /usr/local/lib/python3.6/dist-packages (from trax) (1.4.1)
    Requirement already satisfied: gym in /usr/local/lib/python3.6/dist-packages (from trax) (0.17.3)
    Requirement already satisfied: six in /usr/local/lib/python3.6/dist-packages (from trax) (1.15.0)
    Requirement already satisfied: jax in /usr/local/lib/python3.6/dist-packages (from trax) (0.2.6)
    Requirement already satisfied: tensorflow-datasets in /usr/local/lib/python3.6/dist-packages (from trax) (4.0.1)
    Collecting funcsigs
      Downloading https://files.pythonhosted.org/packages/69/cb/f5be453359271714c01b9bd06126eaf2e368f1fddfff30818754b5ac2328/funcsigs-1.0.2-py2.py3-none-any.whl
    Requirement already satisfied: jaxlib in /usr/local/lib/python3.6/dist-packages (from trax) (0.1.57+cuda101)
    Requirement already satisfied: absl-py in /usr/local/lib/python3.6/dist-packages (from trax) (0.10.0)
    Collecting t5
      Downloading https://files.pythonhosted.org/packages/20/39/a607d2450190af7675e4f77c5eff0cc9a83f82fe63fb396872ef2004106b/t5-0.7.1-py3-none-any.whl (172kB)
         |████████████████████████████████| 174kB 33.6MB/s
    Requirement already satisfied: tensorflow<2.4,>=2.3.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow-text->trax) (2.3.0)
    Requirement already satisfied: pyglet<=1.5.0,>=1.4.0 in /usr/local/lib/python3.6/dist-packages (from gym->trax) (1.5.0)
    Requirement already satisfied: cloudpickle<1.7.0,>=1.2.0 in /usr/local/lib/python3.6/dist-packages (from gym->trax) (1.3.0)
    Requirement already satisfied: opt-einsum in /usr/local/lib/python3.6/dist-packages (from jax->trax) (3.3.0)
    Requirement already satisfied: attrs>=18.1.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow-datasets->trax) (20.3.0)
    Requirement already satisfied: importlib-resources; python_version < "3.9" in /usr/local/lib/python3.6/dist-packages (from tensorflow-datasets->trax) (3.3.0)
    Requirement already satisfied: future in /usr/local/lib/python3.6/dist-packages (from tensorflow-datasets->trax) (0.16.0)
    Requirement already satisfied: tqdm in /usr/local/lib/python3.6/dist-packages (from tensorflow-datasets->trax) (4.41.1)
    Requirement already satisfied: dataclasses; python_version < "3.7" in /usr/local/lib/python3.6/dist-packages (from tensorflow-datasets->trax) (0.8)
    Requirement already satisfied: promise in /usr/local/lib/python3.6/dist-packages (from tensorflow-datasets->trax) (2.3)
    Requirement already satisfied: tensorflow-metadata in /usr/local/lib/python3.6/dist-packages (from tensorflow-datasets->trax) (0.25.0)
    Requirement already satisfied: dill in /usr/local/lib/python3.6/dist-packages (from tensorflow-datasets->trax) (0.3.3)
    Requirement already satisfied: termcolor in /usr/local/lib/python3.6/dist-packages (from tensorflow-datasets->trax) (1.1.0)
    Requirement already satisfied: protobuf>=3.6.1 in /usr/local/lib/python3.6/dist-packages (from tensorflow-datasets->trax) (3.12.4)
    Requirement already satisfied: dm-tree in /usr/local/lib/python3.6/dist-packages (from tensorflow-datasets->trax) (0.1.5)
    Requirement already satisfied: requests>=2.19.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow-datasets->trax) (2.23.0)
    Requirement already satisfied: flatbuffers in /usr/local/lib/python3.6/dist-packages (from jaxlib->trax) (1.12)
    Collecting sentencepiece
      Downloading https://files.pythonhosted.org/packages/e5/2d/6d4ca4bef9a67070fa1cac508606328329152b1df10bdf31fb6e4e727894/sentencepiece-0.1.94-cp36-cp36m-manylinux2014_x86_64.whl (1.1MB)
         |████████████████████████████████| 1.1MB 38.0MB/s
    Requirement already satisfied: torch in /usr/local/lib/python3.6/dist-packages (from t5->trax) (1.7.0+cu101)
    Collecting transformers>=2.7.0
      Downloading https://files.pythonhosted.org/packages/99/84/7bc03215279f603125d844bf81c3fb3f2d50fe8e511546eb4897e4be2067/transformers-4.0.0-py3-none-any.whl (1.4MB)
         |████████████████████████████████| 1.4MB 29.4MB/s
    Collecting mesh-tensorflow[transformer]>=0.1.13
      Downloading https://files.pythonhosted.org/packages/16/8b/553deb763ce8d00afb17debab7cb14a87b209cd4c6f0e8ecfc8d884cb12a/mesh_tensorflow-0.1.17-py3-none-any.whl (342kB)
         |████████████████████████████████| 348kB 44.3MB/s
    Collecting rouge-score
      Downloading https://files.pythonhosted.org/packages/1f/56/a81022436c08b9405a5247b71635394d44fe7e1dbedc4b28c740e09c2840/rouge_score-0.0.4-py2.py3-none-any.whl
    Requirement already satisfied: babel in /usr/local/lib/python3.6/dist-packages (from t5->trax) (2.9.0)
    Requirement already satisfied: scikit-learn in /usr/local/lib/python3.6/dist-packages (from t5->trax) (0.22.2.post1)
    Collecting sacrebleu
      Downloading https://files.pythonhosted.org/packages/a3/c4/8e948f601a4f9609e8b2b58f31966cb13cf17b940b82aa3e767f01c42c52/sacrebleu-1.4.14-py3-none-any.whl (64kB)
         |████████████████████████████████| 71kB 6.1MB/s
    Requirement already satisfied: pandas in /usr/local/lib/python3.6/dist-packages (from t5->trax) (1.1.4)
    Collecting tfds-nightly
      Downloading https://files.pythonhosted.org/packages/0b/b5/4b1d9cf76b8fb047563eca49c86e0fc74dc867c3f7ac9122486b939421c0/tfds_nightly-4.1.0.dev202012080107-py3-none-any.whl (3.7MB)
         |████████████████████████████████| 3.7MB 41.6MB/s
    Requirement already satisfied: nltk in /usr/local/lib/python3.6/dist-packages (from t5->trax) (3.2.5)
    Requirement already satisfied: grpcio>=1.8.6 in /usr/local/lib/python3.6/dist-packages (from tensorflow<2.4,>=2.3.0->tensorflow-text->trax) (1.33.2)
    Requirement already satisfied: wrapt>=1.11.1 in /usr/local/lib/python3.6/dist-packages (from tensorflow<2.4,>=2.3.0->tensorflow-text->trax) (1.12.1)
    Requirement already satisfied: h5py<2.11.0,>=2.10.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow<2.4,>=2.3.0->tensorflow-text->trax) (2.10.0)
    Requirement already satisfied: tensorboard<3,>=2.3.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow<2.4,>=2.3.0->tensorflow-text->trax) (2.3.0)
    Requirement already satisfied: gast==0.3.3 in /usr/local/lib/python3.6/dist-packages (from tensorflow<2.4,>=2.3.0->tensorflow-text->trax) (0.3.3)
    Requirement already satisfied: wheel>=0.26 in /usr/local/lib/python3.6/dist-packages (from tensorflow<2.4,>=2.3.0->tensorflow-text->trax) (0.35.1)
    Requirement already satisfied: astunparse==1.6.3 in /usr/local/lib/python3.6/dist-packages (from tensorflow<2.4,>=2.3.0->tensorflow-text->trax) (1.6.3)
    Requirement already satisfied: tensorflow-estimator<2.4.0,>=2.3.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow<2.4,>=2.3.0->tensorflow-text->trax) (2.3.0)
    Requirement already satisfied: keras-preprocessing<1.2,>=1.1.1 in /usr/local/lib/python3.6/dist-packages (from tensorflow<2.4,>=2.3.0->tensorflow-text->trax) (1.1.2)
    Requirement already satisfied: google-pasta>=0.1.8 in /usr/local/lib/python3.6/dist-packages (from tensorflow<2.4,>=2.3.0->tensorflow-text->trax) (0.2.0)
    Requirement already satisfied: zipp>=0.4; python_version < "3.8" in /usr/local/lib/python3.6/dist-packages (from importlib-resources; python_version < "3.9"->tensorflow-datasets->trax) (3.4.0)
    Requirement already satisfied: googleapis-common-protos<2,>=1.52.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow-metadata->tensorflow-datasets->trax) (1.52.0)
    Requirement already satisfied: setuptools in /usr/local/lib/python3.6/dist-packages (from protobuf>=3.6.1->tensorflow-datasets->trax) (50.3.2)
    Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.6/dist-packages (from requests>=2.19.0->tensorflow-datasets->trax) (2020.11.8)
    Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.6/dist-packages (from requests>=2.19.0->tensorflow-datasets->trax) (1.24.3)
    Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.6/dist-packages (from requests>=2.19.0->tensorflow-datasets->trax) (3.0.4)
    Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.6/dist-packages (from requests>=2.19.0->tensorflow-datasets->trax) (2.10)
    Requirement already satisfied: typing-extensions in /usr/local/lib/python3.6/dist-packages (from torch->t5->trax) (3.7.4.3)
    Collecting sacremoses
      Downloading https://files.pythonhosted.org/packages/7d/34/09d19aff26edcc8eb2a01bed8e98f13a1537005d31e95233fd48216eed10/sacremoses-0.0.43.tar.gz (883kB)
         |████████████████████████████████| 890kB 38.5MB/s
    Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.6/dist-packages (from transformers>=2.7.0->t5->trax) (2019.12.20)
    Collecting tokenizers==0.9.4
      Downloading https://files.pythonhosted.org/packages/0f/1c/e789a8b12e28be5bc1ce2156cf87cb522b379be9cadc7ad8091a4cc107c4/tokenizers-0.9.4-cp36-cp36m-manylinux2010_x86_64.whl (2.9MB)
         |████████████████████████████████| 2.9MB 41.5MB/s
    Requirement already satisfied: filelock in /usr/local/lib/python3.6/dist-packages (from transformers>=2.7.0->t5->trax) (3.0.12)
    Requirement already satisfied: packaging in /usr/local/lib/python3.6/dist-packages (from transformers>=2.7.0->t5->trax) (20.4)
    Requirement already satisfied: pytz>=2015.7 in /usr/local/lib/python3.6/dist-packages (from babel->t5->trax) (2018.9)
    Requirement already satisfied: joblib>=0.11 in /usr/local/lib/python3.6/dist-packages (from scikit-learn->t5->trax) (0.17.0)
    Collecting portalocker
      Downloading https://files.pythonhosted.org/packages/89/a6/3814b7107e0788040870e8825eebf214d72166adf656ba7d4bf14759a06a/portalocker-2.0.0-py2.py3-none-any.whl
    Requirement already satisfied: python-dateutil>=2.7.3 in /usr/local/lib/python3.6/dist-packages (from pandas->t5->trax) (2.8.1)
    Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /usr/local/lib/python3.6/dist-packages (from tensorboard<3,>=2.3.0->tensorflow<2.4,>=2.3.0->tensorflow-text->trax) (0.4.2)
    Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /usr/local/lib/python3.6/dist-packages (from tensorboard<3,>=2.3.0->tensorflow<2.4,>=2.3.0->tensorflow-text->trax) (1.7.0)
    Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.6/dist-packages (from tensorboard<3,>=2.3.0->tensorflow<2.4,>=2.3.0->tensorflow-text->trax) (3.3.3)
    Requirement already satisfied: werkzeug>=0.11.15 in /usr/local/lib/python3.6/dist-packages (from tensorboard<3,>=2.3.0->tensorflow<2.4,>=2.3.0->tensorflow-text->trax) (1.0.1)
    Requirement already satisfied: google-auth<2,>=1.6.3 in /usr/local/lib/python3.6/dist-packages (from tensorboard<3,>=2.3.0->tensorflow<2.4,>=2.3.0->tensorflow-text->trax) (1.17.2)
    Requirement already satisfied: click in /usr/local/lib/python3.6/dist-packages (from sacremoses->transformers>=2.7.0->t5->trax) (7.1.2)
    Requirement already satisfied: pyparsing>=2.0.2 in /usr/local/lib/python3.6/dist-packages (from packaging->transformers>=2.7.0->t5->trax) (2.4.7)
    Requirement already satisfied: requests-oauthlib>=0.7.0 in /usr/local/lib/python3.6/dist-packages (from google-auth-oauthlib<0.5,>=0.4.1->tensorboard<3,>=2.3.0->tensorflow<2.4,>=2.3.0->tensorflow-text->trax) (1.3.0)
    Requirement already satisfied: importlib-metadata; python_version < "3.8" in /usr/local/lib/python3.6/dist-packages (from markdown>=2.6.8->tensorboard<3,>=2.3.0->tensorflow<2.4,>=2.3.0->tensorflow-text->trax) (2.0.0)
    Requirement already satisfied: rsa<5,>=3.1.4; python_version >= "3" in /usr/local/lib/python3.6/dist-packages (from google-auth<2,>=1.6.3->tensorboard<3,>=2.3.0->tensorflow<2.4,>=2.3.0->tensorflow-text->trax) (4.6)
    Requirement already satisfied: cachetools<5.0,>=2.0.0 in /usr/local/lib/python3.6/dist-packages (from google-auth<2,>=1.6.3->tensorboard<3,>=2.3.0->tensorflow<2.4,>=2.3.0->tensorflow-text->trax) (4.1.1)
    Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.6/dist-packages (from google-auth<2,>=1.6.3->tensorboard<3,>=2.3.0->tensorflow<2.4,>=2.3.0->tensorflow-text->trax) (0.2.8)
    Requirement already satisfied: oauthlib>=3.0.0 in /usr/local/lib/python3.6/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard<3,>=2.3.0->tensorflow<2.4,>=2.3.0->tensorflow-text->trax) (3.1.0)
    Requirement already satisfied: pyasn1>=0.1.3 in /usr/local/lib/python3.6/dist-packages (from rsa<5,>=3.1.4; python_version >= "3"->google-auth<2,>=1.6.3->tensorboard<3,>=2.3.0->tensorflow<2.4,>=2.3.0->tensorflow-text->trax) (0.4.8)
    Building wheels for collected packages: sacremoses
      Building wheel for sacremoses (setup.py) ... done
      Created wheel for sacremoses: filename=sacremoses-0.0.43-cp36-none-any.whl size=893257 sha256=e1995006bc7b33a7495fcea495804b86fd414a9746dcb2740bdfe316575233ef
      Stored in directory: /root/.cache/pip/wheels/29/3c/fd/7ce5c3f0666dab31a50123635e6fb5e19ceb42ce38d4e58f45
    Successfully built sacremoses
    Installing collected packages: tensorflow-text, funcsigs, sentencepiece, sacremoses, tokenizers, transformers, mesh-tensorflow, rouge-score, portalocker, sacrebleu, tfds-nightly, t5, trax
    Successfully installed funcsigs-1.0.2 mesh-tensorflow-0.1.17 portalocker-2.0.0 rouge-score-0.0.4 sacrebleu-1.4.14 sacremoses-0.0.43 sentencepiece-0.1.94 t5-0.7.1 tensorflow-text-2.3.0 tfds-nightly-4.1.0.dev202012080107 tokenizers-0.9.4 transformers-4.0.0 trax-1.3.6
    trax                          1.3.6

Let's also declare some constants we will be using in the exercises.

```python
# filename of the MultiWOZ dialogue dataset
DATA_FILE = 'data.json'

# data directory
DATA_DIR = './data'

# dictionary where we will load the dialogue dataset
DIALOGUE_DB = {}

# vocabulary filename
VOCAB_FILE = 'en_32k.subword'

# vocabulary file directory
VOCAB_DIR = 'data/vocabs'
```

Let's now load the MultiWOZ 2.1 dataset. We have already provided it for you in your workspace. It is in JSON format so we should load it as such:

```python
# help function to load a JSON file
def load_json(directory, file):
    with open(f'{directory}/{file}') as file:
        db = json.load(file)
    return db

# load the dialogue data set into our dictionary
DIALOGUE_DB = load_json(DATA_DIR, DATA_FILE)
```

Let's see how many dialogues we have in the dictionary. 1 key-value pair is one dialogue so we can just get the dictionary's length.

```python
print(f'The number of dialogues is: {len(DIALOGUE_DB)}')
```

    The number of dialogues is: 10438

The dialogues are composed of multiple files and the filenames are used as keys in our dictionary. Those with multi-domain dialogues have "MUL" in their filenames while single domain dialogues have either "SNG" or "WOZ".

```python
# print 7 keys from the dataset to see the filenames
print(list(DIALOGUE_DB.keys())[0:7])
```

    ['SNG01856.json', 'SNG0129.json', 'PMUL1635.json', 'MUL2168.json', 'SNG0073.json', 'SNG01445.json', 'MUL2105.json']

As you can see from the cells above, there are 10,438 conversations, each in its own file.  You will train your model on all those conversations. Each file is also loaded into a dictionary and each has two keys which are the following:

```python
# get keys of the fifth file in the list above
print(DIALOGUE_DB['SNG0073.json'].keys())
```

    dict_keys(['goal', 'log'])

The `goal` also points to a dictionary and it contains several keys pertaining to the objectives of the conversation. For example below, we can see that the conversation will be about booking a taxi.

```python
DIALOGUE_DB['SNG0073.json']['goal']
```

    {'attraction': {},
     'hospital': {},
     'hotel': {},
     'message': ["You want to book a taxi. The taxi should go to pizza hut fen ditton and should depart from saint john's college",
      "The taxi should leave after 17:15",
      "Make sure you get car type and contact number"],
     'police': {},
     'restaurant': {},
     'taxi': {'fail_info': {},
      'info': {'departure': "saint john's college",
       'destination': 'pizza hut fen ditton',
       'leaveAt': '17:15'},
      'reqt': ['car type', 'phone']},
     'train': {}}

The `log` on the other hand contains the dialog. It is a list of dictionaries and each element of this list contains several descriptions as well. Let's look at an example:

```python
# get first element of the log list
DIALOGUE_DB['SNG0073.json']['log'][0]
```

    {'dialog_act': {'Taxi-Inform': [['Dest', 'pizza hut fen ditton'],
       ['Depart', "saint john 's college"]]},
     'metadata': {},
     'span_info': [['Taxi-Inform', 'Dest', 'pizza hut fen ditton', 11, 14],
      ['Taxi-Inform', 'Depart', "saint john 's college", 6, 9]],
     'text': "I would like a taxi from Saint John's college to Pizza Hut Fen Ditton."}

For this assignment, we are only interested in the conversation which is in the `text` field.
The conversation goes back and forth between two persons. Let's call them 'Person 1' and 'Person 2'. This implies that
data['SNG0073.json']['log'][0]['text'] is 'Person 1' and
data['SNG0073.json']['log'][1]['text'] is 'Person 2' and so on. The even offsets are 'Person 1' and the odd offsets are 'Person 2'.

```python
print(' Person 1: ', DIALOGUE_DB['SNG0073.json']['log'][0]['text'])
print(' Person 2: ',DIALOGUE_DB['SNG0073.json']['log'][1]['text'])
```

     Person 1:  I would like a taxi from Saint John's college to Pizza Hut Fen Ditton.
     Person 2:  What time do you want to leave and what time do you want to arrive by?

```python

def get_conversation(file, data_db):
    '''
    Args:
        file (string): filename of the dialogue file saved as json
        data_db (dict): dialogue database

    Returns:
        string: A string containing the 'text' fields of  data[file]['log'][x]
    '''

    # initialize empty string
    result = ''

    # get length of file's log list
    len_msg_log = len(data_db[file]['log'])

    # set the delimiter strings
    delimiter_1 = ' Person 1: '
    delimiter_2 = ' Person 2: '

    # loop over the file's log list
    for i in range(len_msg_log):

        # get i'th element of file log list
        cur_log = data_db[file]['log'][i]

        # check if i is even
        if i%2 == 0:
            # append the 1st delimiter string
            result += delimiter_1
        else:
            # append the 2nd delimiter string
            result += delimiter_2

        # append the message text from the log
        result += cur_log['text']

    return result

```

```python
file = 'SNG01856.json'
conversation = get_conversation(file, DIALOGUE_DB)

# print raw output
print(conversation)
```

     Person 1: am looking for a place to to stay that has cheap price range it should be in a type of hotel Person 2: Okay, do you have a specific area you want to stay in? Person 1: no, i just need to make sure it's cheap. oh, and i need parking Person 2: I found 1 cheap hotel for you that includes parking. Do you like me to book it? Person 1: Yes, please. 6 people 3 nights starting on tuesday. Person 2: I am sorry but I wasn't able to book that for you for Tuesday. Is there another day you would like to stay or perhaps a shorter stay? Person 1: how about only 2 nights. Person 2: Booking was successful.
    Reference number is : 7GAWK763. Anything else I can do for you? Person 1: No, that will be all. Good bye. Person 2: Thank you for using our services.

We can have a utility pretty print function just so we can visually follow the conversation more easily.

```python
def print_conversation(conversation):

    delimiter_1 = 'Person 1: '
    delimiter_2 = 'Person 2: '

    split_list_d1 = conversation.split(delimiter_1)

    for sublist in split_list_d1[1:]:
        split_list_d2 = sublist.split(delimiter_2)
        print(colored(f'Person 1: {split_list_d2[0]}', 'red'))

        if len(split_list_d2) > 1:
            print(colored(f'Person 2: {split_list_d2[1]}', 'green'))

print_conversation(conversation)
```

    Person 1: am looking for a place to to stay that has cheap price range it should be in a type of hotel
    Person 2: Okay, do you have a specific area you want to stay in?
    Person 1: no, i just need to make sure it's cheap. oh, and i need parking
    Person 2: I found 1 cheap hotel for you that includes parking. Do you like me to book it?
    Person 1: Yes, please. 6 people 3 nights starting on tuesday.
    Person 2: I am sorry but I wasn't able to book that for you for Tuesday. Is there another day you would like to stay or perhaps a shorter stay?
    Person 1: how about only 2 nights.
    Person 2: Booking was successful.
    Reference number is : 7GAWK763. Anything else I can do for you?
    Person 1: No, that will be all. Good bye.
    Person 2: Thank you for using our services.

```python

DIALOGUE_DB['SNG01856.json']['log'][0]
```

    {'dialog_act': {'Hotel-Inform': [['Type', 'hotel'], ['Price', 'cheap']]},
     'metadata': {},
     'span_info': [['Hotel-Inform', 'Type', 'hotel', 20, 20],
      ['Hotel-Inform', 'Price', 'cheap', 10, 10]],
     'text': 'am looking for a place to to stay that has cheap price range it should be in a type of hotel'}

The dataset also comes with hotel, hospital, taxi, train, police, and restaurant databases. For example, in case you need to call a doctor, or a hotel, or a taxi, this will allow you to automate the entire conversation. Take a look at the files accompanying the data set.

```python
# this is an example of the attractions file
attraction_file = open('data/attraction_db.json')
attractions = json.load(attraction_file)
print(attractions[0])
```

    {'address': 'pool way, whitehill road, off newmarket road', 'area': 'east', 'entrance fee': '?', 'id': '1', 'location': [52.208789, 0.154883], 'name': 'abbey pool and astroturf pitch', 'openhours': '?', 'phone': '01223902088', 'postcode': 'cb58nt', 'pricerange': '?', 'type': 'swimmingpool'}

```python
# this is an example of the hospital file
hospital_file = open('data/hospital_db.json')
hospitals = json.load(hospital_file)
print(hospitals[0]) # feel free to index into other indices
```

    {'department': 'neurosciences critical care unit', 'id': 0, 'phone': '01223216297'}

```python
# this is an example of the hotel file
hotel_file = open('data/hotel_db.json')
hotels = json.load(hotel_file)
print(hotels[0]) # feel free to index into other indices
```

    {'address': '124 tenison road', 'area': 'east', 'internet': 'yes', 'parking': 'no', 'id': '0', 'location': [52.1963733, 0.1987426], 'name': 'a and b guest house', 'phone': '01223315702', 'postcode': 'cb12dp', 'price': {'double': '70', 'family': '90', 'single': '50'}, 'pricerange': 'moderate', 'stars': '4', 'takesbookings': 'yes', 'type': 'guesthouse'}

```python
# this is an example of the police file
police_file = open('data/police_db.json')
police = json.load(police_file)
print(police[0]) # feel free to index into other indices
```

    {'name': 'Parkside Police Station', 'address': 'Parkside, Cambridge', 'id': 0, 'phone': '01223358966'}

```python
# this is an example of a restuarant file
restaurant_file = open('data/restaurant_db.json')
restaurants = json.load(restaurant_file)
print(restaurants[0]) # feel free to index into other indices
```

    {'address': 'Regent Street City Centre', 'area': 'centre', 'food': 'italian', 'id': '19210', 'introduction': 'Pizza hut is a large chain with restaurants nationwide offering convenience pizzas pasta and salads to eat in or take away', 'location': [52.20103, 0.126023], 'name': 'pizza hut city centre', 'phone': '01223323737', 'postcode': 'cb21ab', 'pricerange': 'cheap', 'type': 'restaurant'}

For more information about the multiwoz 2.1 data set, please run the cell below to read the `ReadMe.txt` file. Feel free to open any other file to explore it.

```python
with open('data/README') as file:
    print(file.read())
```

    #####################################################
    #####################################################
    #  Copyright Cambridge Dialogue Systems Group, 2018 #
    #####################################################
    #####################################################

    Dataset contains the following files:
    1. data.json: the woz dialogue dataset, which contains the conversation  users and wizards, as well as a set of coarse labels for each user turn. This file contains both system and user dialogue acts annotated at the turn level. Files with multi-domain dialogues have "MUL" in their names. Single domain dialogues have either "SNG" or "WOZ" in their names.
    2. restaurant_db.json: the Cambridge restaurant database file, containing restaurants in the Cambridge UK area and a set of attributes.
    3. attraction_db.json: the Cambridge attraction database file, contining attractions in the Cambridge UK area and a set of attributes.
    4. hotel_db.json: the Cambridge hotel database file, containing hotels in the Cambridge UK area and a set of attributes.
    5. train_db.json: the Cambridge train (with artificial connections) database file, containing trains in the Cambridge UK area and a set of attributes.
    6. hospital_db.json: the Cambridge hospital database file, contatining information about departments.
    7. police_db.json: the Cambridge police station information.
    8. taxi_db.json: slot-value list for taxi domain.
    9. valListFile.txt: list of dialogues for validation.
    10. testListFile.txt: list of dialogues for testing.
    11. system_acts.json:
      There are 6 domains ('Booking', 'Restaurant', 'Hotel', 'Attraction', 'Taxi', 'Train') and 1 dummy domain ('general').
      A domain-dependent dialogue act is defined as a domain token followed by a domain-independent dialogue act, e.g. 'Hotel-inform' means it is an 'inform' act in the Hotel domain.
      Dialogue acts which cannot take slots, e.g., 'good bye', are defined under the 'general' domain.
      A slot-value pair defined as a list with two elements. The first element is slot token and the second one is its value.
      If a dialogue act takes no slots, e.g., dialogue act 'offer booking' for an utterance 'would you like to take a reservation?', its slot-value pair is ['none', 'none']
      There are four types of values:
      1) If a slot takes a binary value, e.g., 'has Internet' or 'has park', the value is either 'yes' or 'no'.
      2) If a slot is under the act 'request', e.g., 'request' about 'area', the value is expressed as '?'.
      3) The value that appears in the utterance e.g., the name of a restaurant.
      4) If for some reason the turn does not have an annotation then it is labeled as "No Annotation."
    12. ontology.json: Data-based ontology containing all the values for the different slots in the domains.
    13. slot_descriptions.json: A collection of human-written slot descriptions for each slot in the dataset. Each slot has at least two descriptions.
    14. tokenization.md: A description of the tokenization preprocessing we had to perform to maintain consistency between the dialogue act annotations of DSTC 8 Track 1 and the existing MultiWOZ 2.0 data.

As you can see, there are many other aspects of the MultiWoz dataset. Nonetheless, you'll see that even with just the conversations, your model will still be able to generate useful responses. This concludes our exploration of the dataset. In the next section, we will do some preprocessing before we feed it into our model for training.

# Part 2:   Processing the data for Reformer inputs

You will now use the `get_conversation()` function to process the data. The Reformer expects inputs of this form:

**Person 1: Why am I so happy? Person 2: Because you are learning NLP Person 1: ... Person 2: ...***

And the conversation keeps going with some text. As you can see 'Person 1' and 'Person 2' act as delimiters so the model automatically recognizes the person and who is talking. It can then come up with the corresponding text responses for each person. Let's proceed to process the text in this fashion for the Reformer. First, let's grab all the conversation strings from all dialogue files and put them in a list.

```python
# the keys are the file names
all_files = DIALOGUE_DB.keys()

# initialize empty list
untokenized_data = []

# loop over all files
for file in all_files:
    # this is the graded function you coded
    # returns a string delimited by Person 1 and Person 2
    result = get_conversation(file, DIALOGUE_DB)

    # append to the list
    untokenized_data.append(result)

# print the first element to check if it's the same as the one we got before
print(untokenized_data[0])
```

     Person 1: am looking for a place to to stay that has cheap price range it should be in a type of hotel Person 2: Okay, do you have a specific area you want to stay in? Person 1: no, i just need to make sure it's cheap. oh, and i need parking Person 2: I found 1 cheap hotel for you that includes parking. Do you like me to book it? Person 1: Yes, please. 6 people 3 nights starting on tuesday. Person 2: I am sorry but I wasn't able to book that for you for Tuesday. Is there another day you would like to stay or perhaps a shorter stay? Person 1: how about only 2 nights. Person 2: Booking was successful.
    Reference number is : 7GAWK763. Anything else I can do for you? Person 1: No, that will be all. Good bye. Person 2: Thank you for using our services.

Now let us split the list to a train and eval dataset.

```python
# shuffle the list we generated above
random.shuffle(untokenized_data)

# define a cutoff (5% of the total length for this assignment)
# convert to int because we will use it as a list index
cut_off = int(len(untokenized_data) * .05)

# slice the list. the last elements after the cut_off value will be the eval set. the rest is for training.
train_data, eval_data = untokenized_data[:-cut_off], untokenized_data[-cut_off:]

print(f'number of conversations in the data set: {len(untokenized_data)}')
print(f'number of conversations in train set: {len(train_data)}')
print(f'number of conversations in eval set: {len(eval_data)}')
```

    number of conversations in the data set: 10438
    number of conversations in train set: 9917
    number of conversations in eval set: 521

## 2.1   Tokenizing, batching with bucketing
We can now proceed in generating tokenized batches of our data. Let's first define a utility generator function to yield elements from our data sets:

```python
def stream(data):
    # loop over the entire data
    while True:
        # get a random element
        d = random.choice(data)

        # yield a tuple pair of identical values
        # (i.e. our inputs to the model will also be our targets during training)
        yield (d, d)
```

Now let's define our data pipeline for tokenizing and batching our data. As in the previous assignments, we will bucket by length and also have an upper bound on the token length.

```python
# trax allows us to use combinators to generate our data pipeline
data_pipeline = trax.data.Serial(
    # randomize the stream
    trax.data.Shuffle(),

    # tokenize the data
    trax.data.Tokenize(vocab_dir=VOCAB_DIR,
                       vocab_file=VOCAB_FILE),

    # filter too long sequences
    trax.data.FilterByLength(2048),

    # bucket by length
    trax.data.BucketByLength(boundaries=[128, 256,  512, 1024],
                             batch_sizes=[16,    8,    4,   2, 1]),

    # add loss weights but do not add it to the padding tokens (i.e. 0)
    trax.data.AddLossWeights(id_to_mask=0)
)

# apply the data pipeline to our train and eval sets
train_stream = data_pipeline(stream(train_data))
eval_stream = data_pipeline(stream(eval_data))
```

Peek into the train stream.

```python
# the stream generators will yield (input, target, weights). let's just grab the input for inspection
inp, _, _ = next(train_stream)

# print the shape. format is (batch size, token length)
print("input shape: ", inp.shape)

# detokenize the first element
print(trax.data.detokenize(inp[0], vocab_dir=VOCAB_DIR, vocab_file=VOCAB_FILE))
```

    input shape:  (4, 512)
     Person 1: I'm looking for a moderately priced place to eat in the North please Person 2: There's a restaurant that serves chinese and a restaurant that serves indian food. Will one of those work for you? Person 1: I would prefer the Indian restaurant please. Person 2: Okay. Would you like for me to make a reservation there? Person 1: Yes. Can you book for 8 people at 19:15 on Monday? Person 2: Ok, you're booked, the table will be reserved for 15 minutes. Reference number is : 7Y3YQNNR.

     Person 1: I also need to get a train.  Person 2: Alright when would you like to leave by? Person 1: I want the train from birmingham new street to arrive in cambridge by 09:30. I plan to travel on monday. Please provide me with the travel time for this trip Person 2: There are 4 options. The travel time for each trip is 163 minutes. Person 1: Get the earliest one, please.  Person 2: The earliest train leaves at 05:40 and arrives 08:23. The price is 75.10 pounds and travel time 163 minutes. Do you want to book, and if so for how many people? Person 1: Yes I would like to book, for 8 people. Person 2: Okay, it is booked. Reference number is VSTQEJ0S Person 1: Thank you for all of your help! Person 2: Thanks, have a great trip!

# Part 3:   Reversible layers

When running large deep models, you will often run out of memory as each layer allocates memory to store activations for use in backpropagation. To save this resource, you need to be able to recompute these activations during the backward pass without storing them during the forward pass. Take a look first at the leftmost diagram below.

This is how the residual networks are implemented in the standard Transformer. It follows that, given `F()` is Attention and `G()` is Feed-forward(FF).
:

\begin{align}
\mathrm{y}_\mathrm{a} &= \mathrm{x} + \mathrm{F}\left(\mathrm{x}\right)\tag{1} \\
\mathrm{y}_{b}&=\mathrm{y}_{a}+\mathrm{G}\left(\mathrm{y}_{a}\right)\tag{2}\\
\end{align}

As you can see, it requires that $\mathrm{x}$ and $\mathrm{y}_{a}$ be saved so it can be used during backpropagation. We want to avoid this to conserve memory and this is where reversible residual connections come in. They are shown in the middle and rightmost diagrams above. The key idea is that we will start with two copies of the input to the model and at each layer we will only update one of them. The activations that we *don’t* update are the ones that will be used to compute the residuals.

Now in this reversible set up you get the following instead:

\begin{align}
\mathrm{y}_{1}&=\mathrm{x}_{1}+\mathrm{F}\left(\mathrm{x}_{2}\right)\tag{3}\\
\mathrm{y}_{2}&=\mathrm{x}_{2}+\mathrm{G}\left(\mathrm{y}_{1}\right)\tag{4}\\
\end{align}
To recover $\mathrm{(x_1,x_2)}$ from $\mathrm{(y_1, y_2)}$

\begin{align}
\mathrm{x}_{2}&=\mathrm{y}_{2}-\mathrm{G}\left(\mathrm{y}_{1}\right)\tag{5}\\
\mathrm{x}_{1}&=\mathrm{y}_{1}-\mathrm{F}\left(\mathrm{x}_{2}\right)\tag{6}\\
\end{align}

With this configuration, we’re now able to run the network fully in reverse. You'll notice that during the backward pass, $\mathrm{x2}$ and $\mathrm{x1}$ can be recomputed based solely on the values of $\mathrm{y2}$ and $\mathrm{y1}$. No need to save it during the forward pass.

```python

def reversible_layer_forward(x, f, g):
    """
    Args:
        x (np.array): an input vector or matrix
        f (function): a function which operates on a vector/matrix
        g (function): a function which operates on a vector/matrix
    Returns:
        y (np.array): an output vector or matrix whose form is determined by 'x', f and g
    """
    # split the input vector into two (* along the last axis because it is the depth dimension)
    x1, x2 = np.split(x, 2, axis=-1)

    # get y1 using equation 3
    y1 = x1 + f(x2)

    # get y2 using equation 4
    y2 = x2 + g(y1)

    # concatenate y1 and y2 along the depth dimension. be sure output is of type np.ndarray
    y = np.concatenate([y1, y2], axis=-1)

    return y
```

```python

def reversible_layer_reverse(y, f, g):
    """
    Args:
        y (np.array): an input vector or matrix
        f (function): a function which operates on a vector/matrix of the form of 'y'
        g (function): a function which operates on a vector/matrix of the form of 'y'
    Returns:
        y (np.array): an output vector or matrix whose form is determined by 'y', f and g
    """

    # split the input vector into two (* along the last axis because it is the depth dimension)
    y1, y2 = np.split(y, 2, axis=-1)

    ### START CODE HERE (REPLACE INSTANCES OF 'None' WITH YOUR CODE) ###

    # compute x2 using equation 5
    x2 = y2 - g(y1)

    # compute x1 using equation 6
    x1 = y1 - f(x2)

    # concatenate x1 and x2 along the depth dimension
    x = np.concatenate([x1, x2], axis=-1)

    return x

```

```python
# UNIT TEST COMMENT: assert at the end can be used in grading as well
f = lambda x: x + 2
g = lambda x: x * 3
input_vector = np.random.uniform(size=(32,))

output_vector = reversible_layer_forward(input_vector, f, g)
reversed_vector = reversible_layer_reverse(output_vector, f, g)

assert np.allclose(reversed_vector, input_vector)
```

## 3.1   Reversible layers and randomness

This is why we were learning about fastmath's random functions and keys in Course 3 Week 1. Utilizing the same key, `trax.fastmath.random.uniform()` will return the same values. This is required for the backward pass to return the correct layer inputs when random noise is introduced in the layer.

```python
# Layers like dropout have noise, so let's simulate it here:
f = lambda x: x + np.random.uniform(size=x.shape)

# See that the above doesn't work any more:
output_vector = reversible_layer_forward(input_vector, f, g)
reversed_vector = reversible_layer_reverse(output_vector, f, g)

assert not np.allclose(reversed_vector, input_vector)  # Fails!!

# It failed because the noise when reversing used a different random seed.

random_seed = 27686
rng = trax.fastmath.random.get_prng(random_seed)
f = lambda x: x + trax.fastmath.random.uniform(key=rng, shape=x.shape)

# See that it works now as the same rng is used on forward and reverse.
output_vector = reversible_layer_forward(input_vector, f, g)
reversed_vector = reversible_layer_reverse(output_vector, f, g)

assert np.allclose(reversed_vector, input_vector,  atol=1e-07)
```

    WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

# Part 4:   ReformerLM Training

We will now proceed to training our model. Since we have already know the two main components that differentiates it from the standard Transformer, LSH in Course 1 and reversible layers above, we can just use the pre-built model already implemented in Trax. It will have this architecture:

Similar to the Transformer we learned earlier, we want to apply an attention and feed forward layer to your inputs. For the Reformer, we improve the memory efficiency by using **reversible decoder blocks** and youwecan picture its implementation in Trax like below:

### Exercise 04
**Instructions:** Implement a wrapper function that returns a Reformer Language Model. You can use Trax's [ReformerLM](https://trax-ml.readthedocs.io/en/latest/trax.models.html#trax.models.reformer.reformer.ReformerLM) to do this quickly. It will have the same architecture as shown above.

```python

def ReformerLM(vocab_size=33000, n_layers=2, mode='train', attention_type=tl.SelfAttention):
    """
    Args:
        vocab_size (int): size of the vocabulary
        n_layers (int): number of decoder layers
        mode (string): setting of the model which can be 'train', 'eval', or 'predict'
        attention_type(class): attention class to use
    Returns:
        model (ReformerLM): a reformer language model implemented in Trax
    """

    # initialize an instance of Trax's ReformerLM class
    model = trax.models.reformer.ReformerLM(
        # set vocab size
        vocab_size=vocab_size,
        # set number of layers
        n_layers=n_layers,
        # set mode
        mode=mode,
        # set attention type
        attention_type=attention_type
    )

    return model
```

```python
# display the model
temp_model = ReformerLM('train')
print(str(temp_model))

# free memory
del temp_model
```

    Serial[
      Serial[
        AssertShape
        ShiftRight(1)
        AssertShape
      ]
      Embedding_train_512
      Dropout
      PositionalEncoding
      Dup_out2
      ReversibleSerial_in2_out2[
        ReversibleHalfResidual_in2_out2[
          Serial[
            LayerNorm
          ]
          SelfAttention
        ]
        ReversibleSwap_in2_out2
        ReversibleHalfResidual_in2_out2[
          Serial[
            LayerNorm
            Dense_2048
            Dropout
            FastGelu
            Dense_512
            Dropout
          ]
        ]
        ReversibleSwap_in2_out2
        ReversibleHalfResidual_in2_out2[
          Serial[
            LayerNorm
          ]
          SelfAttention
        ]
        ReversibleSwap_in2_out2
        ReversibleHalfResidual_in2_out2[
          Serial[
            LayerNorm
            Dense_2048
            Dropout
            FastGelu
            Dense_512
            Dropout
          ]
        ]
        ReversibleSwap_in2_out2
      ]
      Concatenate_in2
      LayerNorm
      Dropout
      Dense_train
      LogSoftmax
    ]

- Create `TrainTask` and `EvalTask`
- Create the training loop `trax.supervised.training.Loop`
- Pass in the following depending to train_task :
    - `labeled_data=train_gen`
    - `loss_layer=tl.CrossEntropyLoss()`
    - `optimizer=trax.optimizers.Adam(0.01)`
    - `lr_schedule=lr_schedule`
    - `n_steps_per_checkpoint=10`

```python

def training_loop(ReformerLM, train_gen, eval_gen, output_dir = "./model/"):
    """
    Args:
        ReformerLM:  the Reformer language model you are building
        train_gen (generator): train data generator.
        eval_gen (generator): Validation generator.
        output_dir (string): Path to save the model output. Defaults to './model/'.

    Returns:
        trax.supervised.training.Loop: Training loop for the model.
    """

    # use the warmup_and_rsqrt_decay learning rate schedule
    lr_schedule = trax.lr.warmup_and_rsqrt_decay(
        n_warmup_steps=1000, max_value=0.01)

    # define the train task
    train_task = training.TrainTask(
        # labeled data
        labeled_data=train_gen,
        # loss layer
        loss_layer=tl.CrossEntropyLoss(),
        # optimizer
        optimizer=trax.optimizers.Adam(0.01),
        # lr_schedule
        lr_schedule=lr_schedule,
        # n_steps
        n_steps_per_checkpoint=10
    )

    # define the eval task
    eval_task = training.EvalTask(
        # labeled data
        labeled_data=eval_gen,
        # metrics
        metrics=[tl.CrossEntropyLoss(), tl.Accuracy()]
    )

    loop = training.Loop(ReformerLM(mode='train'),
                         train_task,
                         eval_tasks=[eval_task],
                         output_dir=output_dir)
    return loop
```

```python
# we will now test your function
!rm -f model/model.pkl.gz
loop = training_loop(ReformerLM, train_stream, eval_stream)
loop.run(10)
```

    Step      1: Total number of trainable weights: 58072296
    Step      1: Ran 1 train steps in 71.82 secs
    Step      1: train CrossEntropyLoss |  10.42980957
    Step      1: eval  CrossEntropyLoss |  10.39183235
    Step      1: eval          Accuracy |  0.00000000

    Step     10: Ran 9 train steps in 213.35 secs
    Step     10: train CrossEntropyLoss |  10.19278431
    Step     10: eval  CrossEntropyLoss |  9.78180218
    Step     10: eval          Accuracy |  0.05691554

# Part 5:   Decode from a pretrained model

We will now proceed on decoding using the model architecture you just implemented. As in the previous weeks, we will be giving you a pretrained model so you can observe meaningful output during inference. You will be using the [autoregressive_sample_stream()](https://trax-ml.readthedocs.io/en/latest/trax.supervised.html#trax.supervised.decoding.autoregressive_sample_stream) decoding method from Trax to do fast inference. Let's define a few parameters to initialize our model.

```python
# define the `predict_mem_len` and `predict_drop_len` of tl.SelfAttention
def attention(*args, **kwargs):
    # number of input positions to remember in a cache when doing fast inference.
    kwargs['predict_mem_len'] = 120
    # number of input elements to drop once the fast inference input cache fills up.
    kwargs['predict_drop_len'] = 120
    # return the attention layer with the parameters defined above
    return tl.SelfAttention(*args, **kwargs)

# define the model using the ReformerLM function you implemented earlier.
model = ReformerLM(
    vocab_size=33000,
    n_layers=6,
    mode='predict',
    attention_type=attention,
)

# define an input signature so we can initialize our model. shape will be (1, 1) and the data type is int32.
shape11 = trax.shapes.ShapeDtype((1, 1), dtype=np.int32)
```

We can now initialize our model from a file containing the pretrained weights. We will save this starting state so we can reset the model state when we generate a new conversation. This will become clearer in the `generate_dialogue()` function later.

Let's define a few utility functions as well to help us tokenize and detokenize. We can use the [tokenize()](https://trax-ml.readthedocs.io/en/latest/trax.data.html#trax.data.tf_inputs.tokenize) and [detokenize()](https://trax-ml.readthedocs.io/en/latest/trax.data.html#trax.data.tf_inputs.detokenize) from `trax.data.tf_inputs` to do this.

```python
def tokenize(sentence, vocab_file, vocab_dir):
    return list(trax.data.tokenize(iter([sentence]), vocab_file=vocab_file, vocab_dir=vocab_dir))[0]

def detokenize(tokens, vocab_file, vocab_dir):
    return trax.data.detokenize(tokens, vocab_file=vocab_file, vocab_dir=vocab_dir)
```

We are now ready to define our decoding function. This will return a generator that yields that next symbol output by the model. It will be able to predict the next words by just feeding it a starting sentence.

```python

def ReformerLM_output_gen(ReformerLM, start_sentence, vocab_file, vocab_dir, temperature):
    """
    Args:
        ReformerLM:  the Reformer language model you just trained
        start_sentence (string): starting sentence of the conversation
        vocab_file (string): vocabulary filename
        vocab_dir (string): directory of the vocabulary file
        temperature (float): parameter for sampling ranging from 0.0 to 1.0.
            0.0: same as argmax, always pick the most probable token
            1.0: sampling from the distribution (can sometimes say random things)

    Returns:
        generator: yields the next symbol generated by the model
    """

    # Create input tokens using the the tokenize function
    input_tokens = tokenize(start_sentence, vocab_file=vocab_file, vocab_dir=vocab_dir)

    # Add batch dimension to array. Convert from (n,) to (x, n) where
    # x is the batch size. Default is 1. (hint: you can use np.expand_dims() with axis=0)
    input_tokens_with_batch = np.array(input_tokens)[None, :]

    # call the autoregressive_sample_stream function from trax
    output_gen = trax.supervised.decoding.autoregressive_sample_stream(
        # model
        ReformerLM,
        # inputs will be the tokens with batch dimension
        inputs=input_tokens_with_batch,
        # temperature
        temperature=temperature
    )

    return output_gen
```

```python
shape11 = trax.shapes.ShapeDtype((1, 1), dtype=np.int32)

def attention(*args, **kwargs):
    kwargs['predict_mem_len'] = 120  # max length for predictions
    kwargs['predict_drop_len'] = 120  # never drop old stuff
    return tl.SelfAttention(*args, **kwargs)

model = ReformerLM(
    vocab_size=33000,
    n_layers=6,
    mode='predict',
    attention_type=attention,
)
```

```python
model.init_from_file('chatbot_model1.pkl.gz',
                     weights_only=True, input_signature=shape11)

STARTING_STATE = model.state
```

```python
def generate_dialogue(ReformerLM, model_state, start_sentence, vocab_file, vocab_dir, max_len, temperature):
    """
    Args:
        ReformerLM:  the Reformer language model you just trained
        model_state (np.array): initial state of the model before decoding
        start_sentence (string): starting sentence of the conversation
        vocab_file (string): vocabulary filename
        vocab_dir (string): directory of the vocabulary file
        max_len (int): maximum number of tokens to generate
        temperature (float): parameter for sampling ranging from 0.0 to 1.0.
            0.0: same as argmax, always pick the most probable token
            1.0: sampling from the distribution (can sometimes say random things)

    Returns:
        generator: yields the next symbol generated by the model
    """

    # define the delimiters we used during training
    delimiter_1 = 'Person 1: '
    delimiter_2 = 'Person 2: '

    # initialize detokenized output
    sentence = ''

    # token counter
    counter = 0

    # output tokens. we insert a ': ' for formatting
    result = [tokenize(': ', vocab_file=vocab_file, vocab_dir=vocab_dir)]

    # reset the model state when starting a new dialogue
    ReformerLM.state = model_state

    # calls the output generator implemented earlier
    output = ReformerLM_output_gen(ReformerLM, start_sentence, vocab_file=VOCAB_FILE, vocab_dir=VOCAB_DIR, temperature=temperature)

    # print the starting sentence
    print(start_sentence.split(delimiter_2)[0].strip())

    # loop below yields the next tokens until max_len is reached. the if-elif is just for prettifying the output.
    for o in output:

        result.append(o)

        sentence = detokenize(np.concatenate(result, axis=0), vocab_file=VOCAB_FILE, vocab_dir=VOCAB_DIR)

        if sentence.endswith(delimiter_1):
            sentence = sentence.split(delimiter_1)[0]
            print(f'{delimiter_2}{sentence}')
            sentence = ''
            result.clear()

        elif sentence.endswith(delimiter_2):
            sentence = sentence.split(delimiter_2)[0]
            print(f'{delimiter_1}{sentence}')
            sentence = ''
            result.clear()

        counter += 1

        if counter > max_len:
            break

```

We can now feed in different starting sentences and see how the model generates the dialogue. You can even input your own starting sentence. Just remember to ask a question that covers the topics in the Multiwoz dataset so you can generate a meaningful conversation.

```python
sample_sentence = ' Person 1: Are there theatres in town? Person 2: '
generate_dialogue(ReformerLM=model, model_state=STARTING_STATE, start_sentence=sample_sentence, vocab_file=VOCAB_FILE, vocab_dir=VOCAB_DIR, max_len=120, temperature=0.2)
```

    Person 1: Are there theatres in town?
    Person 2: : There are 4 theatres in town. Do you have a preference?
    Person 1: Not really, but I would like the one in the south.
    Person 2: I have one theatre, the Junction, and the other is Tenpin.
    Person 1: Could I get the address and postcode?
    Person 1: The postcode is cb17gx, and the address is Clifton Way, Cambridge Leisure Park, Clifton Way. Is there anything else I can i give for?

```python
sample_sentence = ' Person 1: Is there a hospital nearby? Person 2: '
generate_dialogue(ReformerLM=model, model_state=STARTING_STATE, start_sentence=sample_sentence, vocab_file=VOCAB_FILE, vocab_dir=VOCAB_DIR, max_len=120, temperature=0.2)
```

    Person 1: Is there a hospital nearby?
    Person 2: : Addensbrookes Hospital is located at Hills Rd, Cambridge, postcode CB20QQ. Do you need a particular department?
    Person 1: No, I just need the phone number, please.
    Person 2: The phone number is 01223245151.
    Person 1: Thank you. That's all I need.
    Person 2: Thank you for using our services.Goodbye.
    Person 1: Goodbye.

```python
sample_sentence = ' Person 1: Can you book a taxi? Person 2: '
generate_dialogue(ReformerLM=model, model_state=STARTING_STATE, start_sentence=sample_sentence, vocab_file=VOCAB_FILE, vocab_dir=VOCAB_DIR, max_len=120, temperature=0.2)
```

    Person 1: Can you book a taxi?
    Person 2: : I sure can. When would you like to leave?
    Person 1: I need to leave after 13:00.
    Person 2: I'd be happy to help with your request, first I will need to know your destination.
    Person 1: I'm going to be going to be from the city stop restaurant.
    Person 2: Booking completed! Booked car type	:	grey volkswagen
    Contact number	:	07262372

    Person 2: Thank bybybybyby