create_percentage_multiwoz_dataset.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import json\n",
    "import re\n",
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "percentage = \"20\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# with open(\"/dccstor/gpandey11/gaurav/data/multiwoz21/data.json\") as f:\n",
    "#     data = json.load(f)\n",
    "with open(\"createData/multiwoz21/delex_query_results.json\") as f:\n",
    "    data = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"createData/multiwoz21/valListFile.json\", 'r') as f:\n",
    "    valList = f.read().split('\\n')[:-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"createData/multiwoz21/testListFile.json\", 'r') as f:\n",
    "    testList = f.read().split('\\n')[:-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print(valList[1000])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_goal(message):\n",
    "#     message = goals['message']\n",
    "    if isinstance(message, list):  \n",
    "        message = \". \".join(message)\n",
    "    message = \" . \".join(message.split(\". \"))\n",
    "\n",
    "    span_start = \"class='emphasis'>\"\n",
    "    span_end = \"</span>\"\n",
    "    span_start_len = len(span_start)\n",
    "    span_end_len = len(span_end)\n",
    "    \n",
    "    words = message.split()\n",
    "    msg = \"\"\n",
    "    keys = []\n",
    "    key_words = []\n",
    "    i = 0\n",
    "    j = 0\n",
    "    while i < len(words):\n",
    "        if words[i] == '<span':\n",
    "            i += 1\n",
    "            continue\n",
    "        elif words[i].startswith(span_start):\n",
    "            if span_end in words[i]:\n",
    "                sind = len(msg) + 6\n",
    "                end_idx = words[i].find(span_end)\n",
    "                msg += '[key] ' + words[i][span_start_len:end_idx] + \" [key] \"\n",
    "                keys.extend([sind, len(msg)-7])\n",
    "                key_words.append(msg[sind:len(msg)-1])\n",
    "                i += 1\n",
    "            else:\n",
    "                sind = len(msg) + 6\n",
    "                msg += '[key] ' + words[i][span_start_len:] + \" \"\n",
    "#                 msg.append(words[i][span_start_len:])\n",
    "                keys.extend([sind, len(msg)-1])\n",
    "                wk = [msg[sind:len(msg)-1]]\n",
    "                i += 1\n",
    "                while span_end not in words[i]:\n",
    "                    sind = len(msg)\n",
    "                    msg += words[i] + \" \"\n",
    "                    keys.extend([sind, len(msg)-1])\n",
    "                    wk.append(msg[sind:len(msg)-1])\n",
    "                    i += 1\n",
    "                sind = len(msg)  \n",
    "                end_idx = words[i].find(span_end)\n",
    "                msg += words[i][:end_idx] + \" \"\n",
    "                keys.extend([sind, len(msg)-1])\n",
    "                wk.append(msg[sind:len(msg)-1])\n",
    "                key_words.append(\" \".join(wk))\n",
    "                msg += \"[key] \"\n",
    "                i += 1\n",
    "        else:\n",
    "            sind = len(msg)\n",
    "            msg += words[i] + \" \"\n",
    "            i += 1   \n",
    "    \n",
    "    return msg, keys, key_words                             \n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def getMetaData(data):\n",
    "    if isinstance(data, dict): \n",
    "        md = []\n",
    "        for key in data:\n",
    "            md.extend(getMetaData(data[key]))\n",
    "        \n",
    "        return md\n",
    "    \n",
    "    elif isinstance(data, list):\n",
    "        md = []\n",
    "        for d in data:\n",
    "            md.extend(getMetaData(d))\n",
    "        \n",
    "        return md\n",
    "            \n",
    "    else:\n",
    "        if data != \"\":\n",
    "            return data.split()\n",
    "        else:\n",
    "            return []\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def getTopicsFilled(metadata):\n",
    "    topic_lst = []\n",
    "    for m in metadata:\n",
    "        if len(getMetaData(metadata[m])) != 0:\n",
    "            topic_lst.append(m)\n",
    "    \n",
    "    return topic_lst"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "delexKeys = {}\n",
    "delexKeys['train'] = ['leaveAt', 'destination', 'departure', 'arriveBy']\n",
    "delexKeys['hotel'] = ['name']\n",
    "delexKeys['restaurant'] = ['food', 'name', 'time']\n",
    "delexKeys['attraction'] = ['type', 'name']\n",
    "delexKeys['taxi'] = ['leaveAt', 'destination', 'departure', 'arriveBy']\n",
    "delexKeys['hospital'] = ['department']\n",
    "\n",
    "def delexicaliseUser(response, state, query):\n",
    "#     print(response)\n",
    "    topic  = query.split('|')[0]\n",
    "    topic = topic.replace('[Q]', '').strip()\n",
    "#     print(state[topic])\n",
    "    for key in delexKeys[topic]:\n",
    "        if 'info' in state[topic] and key in state[topic]['info']:\n",
    "#             print(key, state[topic]['info'][key])\n",
    "            response = response.replace(state[topic]['info'][key], '['+topic+'_'+key.lower()+']')\n",
    "        \n",
    "        if 'semi' in state[topic] and key in state[topic]['semi']:\n",
    "#             print(key, state[topic]['info'][key])\n",
    "            response = response.replace(state[topic]['semi'][key], '['+topic+'_'+key.lower()+']')\n",
    "    \n",
    "        if 'book' in state[topic] and key == 'time':\n",
    "#             print(key, state[topic]['info'][key])\n",
    "            response = response.replace(state[topic]['book'][key], '['+topic+'_'+key.lower()+']')\n",
    "    \n",
    "        if 'fail_info' in state[topic] and key in state[topic]['fail_info']:\n",
    "            response = response.replace(state[topic]['fail_info'][key], '['+topic+'_'+key.lower()+']')\n",
    "        \n",
    "        if 'fail_book' in state[topic] and key in state[topic]['fail_book']:\n",
    "            response = response.replace(state[topic]['fail_book'][key], '['+topic+'_'+key.lower()+']')\n",
    "\n",
    "#     print(response, '\\n', '\\n')\n",
    "    return response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "importantKeys = {}\n",
    "importantKeys['train'] = ['leaveAt', 'destination', 'departure', 'arriveBy', 'day', 'people']\n",
    "importantKeys['hospital'] = ['department']\n",
    "importantKeys['hotel'] = ['name', 'area', 'parking', 'pricerange', 'stars', 'internet', 'type', 'stay', 'day', 'people']\n",
    "importantKeys['restaurant'] = ['food', 'pricerange', 'name', 'area', 'people', 'day', 'time']\n",
    "importantKeys['attraction'] = ['type', 'name', 'area']\n",
    "importantKeys['police'] = []\n",
    "importantKeys['taxi'] = ['leaveAt', 'destination', 'departure', 'arriveBy']\n",
    "def importantKey(key, topic):\n",
    "    if key in importantKeys[topic]:\n",
    "        return True\n",
    "    return False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def getTopicsFromMsg(msg):\n",
    "    msg = msg.lower()\n",
    "    if msg.find('train') != -1:\n",
    "        return 'train'\n",
    "    elif msg.find('restaurant') != -1 or msg.find('food') != -1:\n",
    "        return 'restaurant'\n",
    "    elif msg.find('hotel') != -1: \n",
    "        return 'hotel'\n",
    "    elif msg.find('attraction') != -1:\n",
    "        return 'attraction'\n",
    "    elif msg.find('hospital') != -1 or msg.find('accident') != -1:\n",
    "        return 'hospital'\n",
    "    elif msg.find('taxi') != -1:\n",
    "        return 'taxi'\n",
    "    elif msg.find('police') != -1 or msg.find('robbed') != -1:\n",
    "        return 'police'\n",
    "    else:\n",
    "        return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def delexQuery(query):\n",
    "    topic = query[0]\n",
    "    query = query[1]\n",
    "    for key in delexKeys[topic]:\n",
    "        if key in query and query[key] != '' and query[key] != 'not mentioned':\n",
    "            query[key] = '['+topic+'_'+key.lower()+']'\n",
    "    \n",
    "#     print(query)\n",
    "    return [topic,query]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def formatResults(results, query, fail_book, txt, failed):\n",
    "    if len(query)!= 0:\n",
    "            topic = query[0]\n",
    "    else:\n",
    "        print('Empty')\n",
    "        return '[KB] Total = 0 [KB]'\n",
    "    \n",
    "    if topic == 'police':\n",
    "        return '[KB] Total = 1 [KB]'\n",
    "    elif topic == 'taxi':\n",
    "        return '[KB] Total = 1 [KB]'\n",
    "    elif len(results) == 0:\n",
    "        return '[KB] Total = 0 [KB]'\n",
    "    else:\n",
    "        if topic not in fail_book:\n",
    "            return \"[KB] \" + \" Total = \" + str(len(results)) + ' [KB]'\n",
    "        for k in fail_book[topic]:\n",
    "            if k in query[1] and fail_book[topic][k] == query[1][k] and failed[topic] == False:\n",
    "                return '[KB] Total = 0 [KB]'\n",
    "        final_str = \"[KB] \" + \" Total = \" + str(len(results)) + ' [KB]'\n",
    "        return final_str\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def formatQueries(query, context):\n",
    "    topic = query[0]\n",
    "    if topic == 'police':\n",
    "        return '[Q] police [Q]'\n",
    "    elif len(query) == 1:\n",
    "        final_str = '[Q] ' + topic + ' | '\n",
    "        for q in importantKeys[topic]:\n",
    "            final_str += q + ' = ' + '*' + ' | '\n",
    "        \n",
    "        final_str += ' [Q]'\n",
    "        return final_str    \n",
    "    else:    \n",
    "        query = query[1]\n",
    "        final_str = '[Q] ' + topic + ' | '\n",
    "        for q in importantKeys[topic]:\n",
    "            if q in query and query[q] != '' and query[q] != 'not mentioned':\n",
    "                final_str += q + ' = ' + str(query[q]) + ' | '\n",
    "            else:\n",
    "                final_str += q + ' = ' + '*' + ' | ' \n",
    "        final_str = final_str[:-3] + ' [Q]'\n",
    "        return final_str    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'prev_delexquery' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-17-efb4bf4e1839>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m    149\u001b[0m             \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    150\u001b[0m                 \u001b[0mdialogue_query\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mprev_query\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 151\u001b[0;31m                 \u001b[0mdialogue_delexquery\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mprev_delexquery\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    152\u001b[0m                 \u001b[0mdialogue_kb\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mprev_kb\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    153\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'prev_delexquery' is not defined"
     ]
    }
   ],
   "source": [
    "user_context_train = []\n",
    "user_context_valid = []\n",
    "user_context_test = []\n",
    "user_context_test_single = []\n",
    "user_context_test_multiple = []\n",
    "user_response_train = []\n",
    "user_response_valid = []\n",
    "user_response_test = []\n",
    "user_response_test_single = []\n",
    "user_response_test_multiple = []\n",
    "agent_context_train = []\n",
    "agent_context_valid = []\n",
    "agent_context_test = []\n",
    "agent_context_test_single = []\n",
    "agent_context_test_multiple = []\n",
    "agent_response_train = []\n",
    "agent_response_valid = []\n",
    "agent_response_test = []\n",
    "agent_response_test_single = []\n",
    "agent_response_test_multiple = []\n",
    "goals_train = []\n",
    "goals_valid = []\n",
    "goals_test = []\n",
    "goals_test_single = []\n",
    "goals_test_multiple = []\n",
    "\n",
    "agent_kb_train = []\n",
    "agent_kb_valid = []\n",
    "agent_kb_test = []\n",
    "agent_kb_test_single = []\n",
    "agent_kb_test_multiple = []\n",
    "\n",
    "agent_query_train = []\n",
    "agent_query_valid = []\n",
    "agent_query_test = []\n",
    "agent_query_test_single = []\n",
    "agent_query_test_multiple = []\n",
    "\n",
    "agent_delexquery_train = []\n",
    "agent_delexquery_valid = []\n",
    "agent_delexquery_test = []\n",
    "agent_delexquery_test_single = []\n",
    "agent_delexquery_test_multiple = []\n",
    "\n",
    "state_train = []\n",
    "state_valid = []\n",
    "state_test = []\n",
    "state_test_single = []\n",
    "state_test_multiple = []\n",
    "\n",
    "dialogue_names_train = []\n",
    "dialogue_names_valid = []\n",
    "dialogue_names_test = []\n",
    "dialogue_names_test_single = []\n",
    "dialogue_names_test_multiple = []\n",
    "\n",
    "u=0\n",
    "multi_goal = []\n",
    "\n",
    "data_keys = list(data.keys())\n",
    "random.shuffle(data_keys, random = lambda: 0.5)\n",
    "\n",
    "for dd in data_keys:\n",
    "#     print('dialogue : ', u)\n",
    "    u+=1\n",
    "    dialogue = data[dd]\n",
    "    \n",
    "    goal = dialogue['goal']['message']\n",
    "    # indicates if the agent has already stated the failure\n",
    "    failed = {}\n",
    "    # from user goal stating that failure will happen\n",
    "    fail_book = {}\n",
    "    topic_count = 0\n",
    "    v_state = {}\n",
    "    te_state = {}\n",
    "    tr_state = {}\n",
    "    for topic in ['train', 'restaurant', 'taxi', 'hotel', 'hospital', 'police', 'attraction']:  \n",
    "        if dialogue['goal'][topic]:\n",
    "#             print(dd, topic)\n",
    "            topic_count += 1\n",
    "            if dd in valList:\n",
    "                v_state[topic] = dialogue['goal'][topic]\n",
    "            elif dd in testList:\n",
    "                te_state[topic] = dialogue['goal'][topic]\n",
    "            else:\n",
    "                tr_state[topic] = dialogue['goal'][topic]            \n",
    "            if 'fail_book' in dialogue['goal'][topic]:\n",
    "                fail_book[topic] = dialogue['goal'][topic]['fail_book']\n",
    "            else:\n",
    "                fail_book[topic] = {}\n",
    "                \n",
    "            failed[topic] = False\n",
    "    \n",
    "    if dd in valList:\n",
    "        state_valid.append(v_state)\n",
    "    elif dd in testList:\n",
    "        state_test.append(te_state) \n",
    "    else:\n",
    "        state_train.append(tr_state) \n",
    "    \n",
    "    if topic_count > 1:\n",
    "        multi_goal.append(dd)\n",
    "    \n",
    "    msg, key, _ = get_goal(goal)\n",
    "\n",
    "    utterances = ['[st@rt]']\n",
    "    speakers = ['Agent']\n",
    "    prev_query = None\n",
    "    prev_kb = None\n",
    "    \n",
    "    dialogue_query = []\n",
    "    dialogue_delexquery = []\n",
    "    dialogue_kb = []\n",
    "    \n",
    "    prev_query = 'no'\n",
    "    prev_kb = 'no'\n",
    "    domain_topic = None\n",
    "    \n",
    "    for log in dialogue['log']:\n",
    "        if bool(log['metadata']):\n",
    "            speakers.append('Agent')\n",
    "            utterances.append('[Agent] ' + log['text'].strip().lower().replace(\"\\n\", \" \"))\n",
    "            if 'queries' in log:\n",
    "                if len(log['queries']) < 1:\n",
    "                    dialogue_query.append(prev_query)\n",
    "                    dialogue_delexquery.append(prev_delexquery)\n",
    "                    dialogue_kb.append(prev_kb)\n",
    "#                     agent_modified_query.append(agent_modified_query[-1])\n",
    "                else:\n",
    "                    domain_topic = log['queries'][0]\n",
    "                    query = formatQueries(log['queries'], utterances)\n",
    "                    delex_query = formatQueries(delexQuery(log['queries']), utterances)\n",
    "                    kb = formatResults(log['results'], log['queries'], fail_book, log['text'], failed)\n",
    "                    dialogue_query.append(query)\n",
    "                    dialogue_delexquery.append(delex_query)\n",
    "                    dialogue_kb.append(kb)\n",
    "#                     agent_modified_query.append(formatQueries(log['modified_queries']))\n",
    "                    prev_query = query\n",
    "                    prev_delexquery = delex_query\n",
    "                    prev_kb = kb\n",
    "                    t = log['queries'][0]\n",
    "                    for l in range(len(dialogue_query)-1,-1,-1):\n",
    "                        if dialogue_query[l-1] == 'no':\n",
    "                            dialogue_query[l-1] = formatQueries([t], utterances)\n",
    "                            dialogue_delexquery[l-1] = formatQueries([t], utterances)\n",
    "                            dialogue_kb[l-1] = formatResults([], [t], fail_book, log['text'], failed)\n",
    "                        else:\n",
    "                            break\n",
    "            else:\n",
    "                dialogue_query.append(prev_query)\n",
    "                dialogue_delexquery.append(prev_delexquery)\n",
    "                dialogue_kb.append(prev_kb)\n",
    "            \n",
    "            if any(word in log['text'].strip().lower().replace(\"\\n\", \" \") for word in [' no ', ' not ', 'unavailable', 'unfortunate', 'sorry', 'unable', 'unsuccessful']):\n",
    "                    if domain_topic is not None:\n",
    "                        failed[domain_topic] = True\n",
    "            \n",
    "        else:\n",
    "            speakers.append('User')\n",
    "            utterances.append('[User] ' + log['text'].strip().replace(\"\\n\", \" \"))\n",
    "    \n",
    "    if 'no' in dialogue_query:\n",
    "        topic = getTopicsFromMsg(msg)\n",
    "        if topic != None:\n",
    "            dialogue_query = [formatQueries([topic], utterances)]*len(dialogue_query)\n",
    "            dialogue_delexquery = [formatQueries([topic], utterances)]*len(dialogue_query)\n",
    "            dialogue_kb = [formatResults([], [topic], fail_book, log['text'], failed)]*len(dialogue_kb)\n",
    "     \n",
    "    if 'no' in dialogue_query:\n",
    "        print(dialogue_query)\n",
    "        print(dialogue_kb)\n",
    "        print(msg)\n",
    "        \n",
    "    if len(dialogue_query) != len(dialogue_delexquery):\n",
    "        raise Exception('Length  Not equal')\n",
    "    \n",
    "#     print(dialogue_query, '\\n')\n",
    "#     print(dialogue_delexquery, '\\n\\n')\n",
    "    \n",
    "    if dd in valList:\n",
    "        agent_query_valid.append(dialogue_query)\n",
    "        agent_delexquery_valid.append(dialogue_delexquery)\n",
    "        agent_kb_valid.append(dialogue_kb)\n",
    "    elif dd in testList:\n",
    "        agent_query_test.append(dialogue_query)\n",
    "        agent_delexquery_test.append(dialogue_delexquery)\n",
    "        agent_kb_test.append(dialogue_kb)\n",
    "        if dd in multi_goal:\n",
    "            agent_query_test_multiple.append(dialogue_query)\n",
    "            agent_kb_test_multiple.append(dialogue_kb)\n",
    "        else:\n",
    "            agent_query_test_single.append(dialogue_query)\n",
    "            agent_kb_test_single.append(dialogue_kb)\n",
    "            state_test_single.append(dialogue['goal'])\n",
    "    else:\n",
    "        agent_query_train.append(dialogue_query)\n",
    "        agent_delexquery_train.append(dialogue_delexquery)        \n",
    "        agent_kb_train.append(dialogue_kb)\n",
    "    \n",
    "    utterances.extend(['[User] [e*d]']) \n",
    "    speakers.extend(['User']) \n",
    "    \n",
    "    dialogue_user_context = []\n",
    "    dialogue_user_response = []\n",
    "    dialogue_agent_context = []\n",
    "    dialogue_agent_response = []\n",
    "    dialogue_goals = []\n",
    "\n",
    "    counts=0\n",
    "    for i in range(1,len(utterances)):\n",
    "        context = utterances[:i]\n",
    "        response = utterances[i]\n",
    "\n",
    "        if speakers[i] == 'User':\n",
    "            if response.find('[e*d]') == -1 and topic != 'police':\n",
    "                response = delexicaliseUser(response, dialogue['goal'], dialogue_query[counts])\n",
    "#                 print(response, '\\n', '\\n')\n",
    "            dialogue_user_context.append(context)\n",
    "            dialogue_user_response.append(response)\n",
    "        \n",
    "            dialogue_goals.append(msg)\n",
    "            counts += 1\n",
    "            \n",
    "        else:\n",
    "            dialogue_agent_context.append(context)\n",
    "            dialogue_agent_response.append(response) \n",
    "#             agent_keys.append(' '.join(map(str,key)))\n",
    "#             agent_goal.append(msg)\n",
    "    \n",
    "    if dd in valList:\n",
    "        user_context_valid.append(dialogue_user_context)\n",
    "        user_response_valid.append(dialogue_user_response)\n",
    "        agent_context_valid.append(dialogue_agent_context)\n",
    "        agent_response_valid.append(dialogue_agent_response)\n",
    "        goals_valid.append(dialogue_goals)\n",
    "        dialogue_names_valid.append(dd)\n",
    "            \n",
    "    elif dd in testList:\n",
    "        user_context_test.append(dialogue_user_context)\n",
    "        user_response_test.append(dialogue_user_response)\n",
    "        agent_context_test.append(dialogue_agent_context)\n",
    "        agent_response_test.append(dialogue_agent_response)\n",
    "        goals_test.append(dialogue_goals)\n",
    "        dialogue_names_test.append(dd)\n",
    "        if dd in multi_goal:\n",
    "            user_context_test_multiple.append(dialogue_user_context)\n",
    "            user_response_test_multiple.append(dialogue_user_response)\n",
    "            agent_context_test_multiple.append(dialogue_agent_context)\n",
    "            agent_response_test_multiple.append(dialogue_agent_response)\n",
    "            goals_test_multiple.append(dialogue_goals)\n",
    "            dialogue_names_test_multiple.append(dd)\n",
    "        else:\n",
    "            user_context_test_single.append(dialogue_user_context)\n",
    "            user_response_test_single.append(dialogue_user_response)\n",
    "            agent_context_test_single.append(dialogue_agent_context)\n",
    "            agent_response_test_single.append(dialogue_agent_response)\n",
    "            goals_test_single.append(dialogue_goals)\n",
    "            dialogue_names_test_single.append(dd)\n",
    "    else:\n",
    "        user_context_train.append(dialogue_user_context)\n",
    "        user_response_train.append(dialogue_user_response)\n",
    "        agent_context_train.append(dialogue_agent_context)\n",
    "        agent_response_train.append(dialogue_agent_response)\n",
    "        goals_train.append(dialogue_goals)\n",
    "        dialogue_names_train.append(dd)\n",
    "#     if u>5:\n",
    "#         break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print(state_train[0][1])\n",
    "# def get_state(state):\n",
    "#     for s in state:\n",
    "#         new_state = {}\n",
    "#         if len(state[s]) > 1:\n",
    "#             for k in state[s]:\n",
    "#                 new_state[s+'_'+k] = state[s][k]\n",
    "            \n",
    "#             new_state[s] = get_state(new_state[s])\n",
    "        \n",
    "# get_state(state_train[0][1])\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "save_dir = \"data/multiwiz/user/\"+percentage+\"p/\"\n",
    "train_length = int(float(percentage)*len(user_context_train)/100)\n",
    "valid_length = int(float(percentage)*len(user_context_valid)/100)\n",
    "\n",
    "with open(save_dir + \"train_input.json\", \"w\") as f:\n",
    "    json.dump(user_context_train[:train_length],f,indent=4)\n",
    "    \n",
    "with open(save_dir + \"rest_input.json\", \"w\") as f:\n",
    "    json.dump(user_context_train[train_length:],f,indent=4)\n",
    "    \n",
    "with open(save_dir + \"valid_input.json\", \"w\") as f:\n",
    "    json.dump(user_context_valid[:valid_length],f,indent=4)\n",
    "    \n",
    "with open(save_dir + \"test_input.json\", \"w\") as f:\n",
    "    json.dump(user_context_test,f,indent=4) \n",
    "    \n",
    "with open(save_dir + \"test_input_single.json\", \"w\") as f:\n",
    "    json.dump(user_context_test_single,f,indent=4) \n",
    "    \n",
    "with open(save_dir + \"test_input_multiple.json\", \"w\") as f:\n",
    "    json.dump(user_context_test_multiple,f,indent=4) \n",
    "\n",
    "\n",
    "with open(save_dir + \"train_tgt.json\", \"w\") as f:\n",
    "    json.dump(user_response_train[:train_length],f,indent=4)\n",
    "    \n",
    "with open(save_dir + \"rest_tgt.json\", \"w\") as f:\n",
    "    json.dump(user_response_train[train_length:],f,indent=4)\n",
    "\n",
    "with open(save_dir + \"valid_tgt.json\", \"w\") as f:\n",
    "    json.dump(user_response_valid[:valid_length],f,indent=4)\n",
    "    \n",
    "with open(save_dir + \"test_tgt.json\", \"w\") as f:\n",
    "    json.dump(user_response_test,f,indent=4)\n",
    "    \n",
    "with open(save_dir + \"test_tgt_single.json\", \"w\") as f:\n",
    "    json.dump(user_response_test_single,f,indent=4)\n",
    "    \n",
    "with open(save_dir + \"test_tgt_multiple.json\", \"w\") as f:\n",
    "    json.dump(user_response_test_multiple,f,indent=4)\n",
    "    \n",
    "with open(save_dir + \"train_goal.json\", \"w\") as f:\n",
    "    json.dump(goals_train[:train_length],f,indent=4)\n",
    "    \n",
    "with open(save_dir + \"rest_goal.json\", \"w\") as f:\n",
    "    json.dump(goals_train[train_length:],f,indent=4)\n",
    "    \n",
    "with open(save_dir + \"valid_goal.json\", \"w\") as f:\n",
    "    json.dump(goals_valid[:valid_length],f,indent=4)\n",
    "    \n",
    "with open(save_dir + \"test_goal.json\", \"w\") as f:\n",
    "    json.dump(goals_test,f,indent=4)\n",
    "    \n",
    "with open(save_dir + \"test_goal_single.json\", \"w\") as f:\n",
    "    json.dump(goals_test_single,f,indent=4)\n",
    "    \n",
    "with open(save_dir + \"test_goal_multiple.json\", \"w\") as f:\n",
    "    json.dump(goals_test_multiple,f,indent=4)\n",
    "    \n",
    "    \n",
    "# with open(save_dir + \"train_adverse.json\", \"w\") as f:\n",
    "#     json.dump(user_adverse[:train_len],f,indent=4)\n",
    "    \n",
    "# with open(save_dir + \"valid_adverse.json\", \"w\") as f:\n",
    "#     json.dump(user_adverse[train_len:valid_len],f,indent=4)  \n",
    "\n",
    "# with open(save_dir + \"test_adverse.json\", \"w\") as f:\n",
    "#     json.dump(user_adverse[valid_len:],f,indent=4)   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "save_dir = \"data/multiwiz/agent/\"+percentage+\"p/\"\n",
    "train_length = int(float(percentage)*len(agent_context_train)/100)\n",
    "valid_length = int(float(percentage)*len(agent_context_valid)/100)\n",
    "\n",
    "with open(save_dir + \"train_input.json\", \"w\") as f:\n",
    "    json.dump(agent_context_train[:train_length],f,indent=4)\n",
    "    \n",
    "with open(save_dir + \"rest_input.json\", \"w\") as f:\n",
    "    json.dump(agent_context_train[train_length:],f,indent=4)\n",
    "    \n",
    "with open(save_dir + \"valid_input.json\", \"w\") as f:\n",
    "    json.dump(agent_context_valid[:valid_length],f,indent=4)\n",
    "    \n",
    "with open(save_dir + \"test_input.json\", \"w\") as f:\n",
    "    json.dump(agent_context_test,f,indent=4) \n",
    "    \n",
    "with open(save_dir + \"test_input_single.json\", \"w\") as f:\n",
    "    json.dump(agent_context_test_single,f,indent=4) \n",
    "    \n",
    "with open(save_dir + \"test_input_multiple.json\", \"w\") as f:\n",
    "    json.dump(agent_context_test_multiple,f,indent=4) \n",
    "    \n",
    "with open(save_dir + \"train_tgt.json\", \"w\") as f:\n",
    "    json.dump(agent_response_train[:train_length],f,indent=4)\n",
    "    \n",
    "with open(save_dir + \"rest_tgt.json\", \"w\") as f:\n",
    "    json.dump(agent_response_train[train_length:],f,indent=4)\n",
    "    \n",
    "with open(save_dir + \"valid_tgt.json\", \"w\") as f:\n",
    "    json.dump(agent_response_valid[:valid_length],f,indent=4)\n",
    "\n",
    "with open(save_dir + \"test_tgt.json\", \"w\") as f:\n",
    "    json.dump(agent_response_test,f,indent=4) \n",
    "    \n",
    "with open(save_dir + \"test_tgt_single.json\", \"w\") as f:\n",
    "    json.dump(agent_response_test_single,f,indent=4) \n",
    "    \n",
    "with open(save_dir + \"test_tgt_multiple.json\", \"w\") as f:\n",
    "    json.dump(agent_response_test_multiple,f,indent=4) \n",
    "    \n",
    "with open(save_dir + \"train_kb.json\", \"w\") as f:\n",
    "    json.dump(agent_kb_train[:train_length],f,indent=4)\n",
    "    \n",
    "with open(save_dir + \"rest_kb.json\", \"w\") as f:\n",
    "    json.dump(agent_kb_train[train_length:],f,indent=4)\n",
    "    \n",
    "with open(save_dir + \"valid_kb.json\", \"w\") as f:\n",
    "    json.dump(agent_kb_valid[:valid_length],f,indent=4)\n",
    "    \n",
    "with open(save_dir + \"test_kb.json\", \"w\") as f:\n",
    "    json.dump(agent_kb_test,f,indent=4)\n",
    "    \n",
    "with open(save_dir + \"test_kb_single.json\", \"w\") as f:\n",
    "    json.dump(agent_kb_test_single,f,indent=4)\n",
    "    \n",
    "with open(save_dir + \"test_kb_multiple.json\", \"w\") as f:\n",
    "    json.dump(agent_kb_test_multiple,f,indent=4)\n",
    "\n",
    "with open(save_dir + \"train_query.json\", \"w\") as f:\n",
    "    json.dump(agent_query_train[:train_length],f,indent=4)\n",
    "    \n",
    "with open(save_dir + \"rest_query.json\", \"w\") as f:\n",
    "    json.dump(agent_query_train[train_length:],f,indent=4)\n",
    "    \n",
    "with open(save_dir + \"valid_query.json\", \"w\") as f:\n",
    "    json.dump(agent_query_valid[:valid_length],f,indent=4) \n",
    "    \n",
    "with open(save_dir + \"test_query.json\", \"w\") as f:\n",
    "    json.dump(agent_query_test,f,indent=4) \n",
    "    \n",
    "with open(save_dir + \"train_delex_query.json\", \"w\") as f:\n",
    "    json.dump(agent_delexquery_train[:train_length],f,indent=4)\n",
    "    \n",
    "with open(save_dir + \"rest_delex_query.json\", \"w\") as f:\n",
    "    json.dump(agent_delexquery_train[train_length:],f,indent=4)\n",
    "    \n",
    "with open(save_dir + \"valid_delex_query.json\", \"w\") as f:\n",
    "    json.dump(agent_delexquery_valid[:valid_length],f,indent=4) \n",
    "    \n",
    "with open(save_dir + \"test_delex_query.json\", \"w\") as f:\n",
    "    json.dump(agent_delexquery_test,f,indent=4) \n",
    "    \n",
    "with open(save_dir + \"test_query_single.json\", \"w\") as f:\n",
    "    json.dump(agent_query_test_single,f,indent=4)\n",
    "    \n",
    "with open(save_dir + \"test_query_multiple.json\", \"w\") as f:\n",
    "    json.dump(agent_query_test_multiple,f,indent=4) \n",
    "    \n",
    "with open(save_dir + \"train_state.json\", \"w\") as f:\n",
    "    json.dump(state_train[:train_length],f,indent=4)\n",
    "    \n",
    "with open(save_dir + \"rest_state.json\", \"w\") as f:\n",
    "    json.dump(state_train[train_length:],f,indent=4)\n",
    "    \n",
    "with open(save_dir + \"valid_state.json\", \"w\") as f:\n",
    "    json.dump(state_valid[:valid_length],f,indent=4) \n",
    "    \n",
    "with open(save_dir + \"test_state.json\", \"w\") as f:\n",
    "    json.dump(state_test,f,indent=4) \n",
    "    \n",
    "with open(save_dir + \"test_state_single.json\", \"w\") as f:\n",
    "    json.dump(state_test_single,f,indent=4)\n",
    "    \n",
    "with open(save_dir + \"test_state_multiple.json\", \"w\") as f:\n",
    "    json.dump(state_test_multiple,f,indent=4) \n",
    "\n",
    "with open(save_dir + \"train_dialogue_names.json\", \"w\") as f:\n",
    "    json.dump(dialogue_names_train[:train_length],f,indent=4)\n",
    "    \n",
    "with open(save_dir + \"rest_dialogue_names.json\", \"w\") as f:\n",
    "    json.dump(dialogue_names_train[train_length:],f,indent=4)\n",
    "    \n",
    "with open(save_dir + \"valid_dialogue_names.json\", \"w\") as f:\n",
    "    json.dump(dialogue_names_valid[:valid_length],f,indent=4)\n",
    "    \n",
    "with open(save_dir + \"test_dialogue_names.json\", \"w\") as f:\n",
    "    json.dump(dialogue_names_test,f,indent=4)\n",
    "    \n",
    "with open(save_dir + \"test_dialogue_names_single.json\", \"w\") as f:\n",
    "    json.dump(dialogue_names_test_single,f,indent=4)\n",
    "    \n",
    "with open(save_dir + \"test_dialogue_names_multiple.json\", \"w\") as f:\n",
    "    json.dump(dialogue_names_test_multiple,f,indent=4)\n",
    "    \n",
    "# with open(save_dir + \"train_mod_query.txt\", \"w\") as fq:\n",
    "#     for c in agent_modified_query[:train_len]:\n",
    "#         fq.write(c + \"\\n\")\n",
    "    \n",
    "# with open(save_dir + \"valid_mod_query.txt\", \"w\") as fq:\n",
    "#     for c in agent_modified_query[train_len:valid_len]:\n",
    "#         fq.write(c + \"\\n\") \n",
    "    \n",
    "# with open(save_dir + \"test_mod_query.txt\", \"w\") as fq:\n",
    "#     for c in agent_modified_query[valid_len:]:\n",
    "#         fq.write(c + \"\\n\")     "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "226"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(state_test_single)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "210"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with open(save_dir + \"train_delex_query.json\", \"r\") as f:\n",
    "    s =json.load(f)\n",
    "\n",
    "len(s)"
   ]
  },
  {
   "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.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}