Redial_Data_Preprocessing.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "8494f226",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy \n",
    "import torch\n",
    "import json\n",
    "import pandas\n",
    "import re\n",
    "\n",
    "from tmdbv3api import TMDb\n",
    "from tmdbv3api import Movie\n",
    "\n",
    "tmdb = TMDb()\n",
    "tmdb.api_key = 'ecbe5b50079424c4372da15d2999da46'\n",
    "tmdb.language = 'en'\n",
    "tmdb.debug = True\n",
    "\n",
    "movie = Movie()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "a208fe40",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'his script returns four list:\\n1. all_movie_db: ID to descripiton (plot) using all movies in the IMDb\\n2. metioned_movie_db: ID to descripiton using the movies mentioned in redial (include both train and test)\\n3. all_data_train: train data\\n4. all_data_test: test data\\n\\n5. not_exist_metioned_movie_db\\n'"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"his script returns four list:\n",
    "1. all_movie_db: ID to descripiton (plot) using all movies in the IMDb\n",
    "2. metioned_movie_db: ID to descripiton using the movies mentioned in redial (include both train and test)\n",
    "3. all_data_train: train data\n",
    "4. all_data_test: test data\n",
    "\n",
    "5. not_exist_metioned_movie_db\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "570e880e",
   "metadata": {},
   "source": [
    "## movie database"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "5bf42bfc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# IMDb dataset\n",
    "\n",
    "df = pandas.read_csv(\"IMDb.csv\")\n",
    "\n",
    "# our ID to description\n",
    "all_movie_db = {} # ID to descripiton\n",
    "imdb_name_with_year_to_id = {}\n",
    "imdb_name_without_year_to_id = {}\n",
    "\n",
    "# unique imdb id set\n",
    "imdb_id_set = {}\n",
    "\n",
    "# unique duplicate name set\n",
    "duplicate_name_set = {}\n",
    "\n",
    "key = 0\n",
    "for i in range(len(df)):\n",
    "    mid = df['imdb_title_id'][i]\n",
    "    if mid not in imdb_id_set.keys():\n",
    "        imdb_id_set[mid]=1\n",
    "        name = df['original_title'][i].lower()\n",
    "        \n",
    "        top3_actors = \"nan\" if str(df['actors'][i]) == \"nan\" else ','.join(df['actors'][i].split(',')[:3])\n",
    "        director  = df['director'][i]\n",
    "        genres = df['genre'][i]\n",
    "        \n",
    "        year = df['year'][i]\n",
    "        full_name = name + ' (' + str(year) + ')'\n",
    "        desc = df['description'][i]\n",
    "        \n",
    "        concated_meta = str(name) + ' [SEP] ' + str(top3_actors) + ' [SEP] ' + \\\n",
    "            str(director) + ' [SEP] ' + str(genres) + ' [SEP] ' + str(desc)\n",
    "    \n",
    "        if name in imdb_name_without_year_to_id.keys():\n",
    "            duplicate_name_set[name] = 1\n",
    "        \n",
    "        # for movie match in redial dataset\n",
    "        all_movie_db[key] = concated_meta\n",
    "        imdb_name_with_year_to_id[full_name] = key\n",
    "        imdb_name_without_year_to_id[name] = key\n",
    "        key +=1 \n",
    "\n",
    "# Movies_metadata dataset\n",
    "new_count = 0\n",
    "df = pandas.read_csv(\"movies_metadata.csv\")\n",
    "for i in range(len(df)):\n",
    "    mid = df['imdb_id'][i]\n",
    "    if mid not in imdb_id_set.keys():\n",
    "        new_count += 1\n",
    "        imdb_id_set[mid]=1\n",
    "        name = df['original_title'][i].lower()\n",
    "        \n",
    "        genres = json.loads(df[\"genres\"][i].replace('\\'', '\\\"'))\n",
    "        genres = ', '.join([ g['name'] for g in genres])\n",
    "        \n",
    "        year = df['release_date'][i]\n",
    "        year = re.findall(r\"\\d\\d\\d\\d\", str(year))\n",
    "        if len(year) > 0:\n",
    "            year = year[0]\n",
    "            full_name = name + ' (' + str(year) + ')'\n",
    "        else:\n",
    "            full_name = name\n",
    "        desc = df['overview'][i]\n",
    "        \n",
    "        concated_meta = str(name) + ' [SEP] [SEP] [SEP] ' + str(genres) + ' [SEP] ' + str(desc)\n",
    "        \n",
    "        if name in imdb_name_without_year_to_id.keys():\n",
    "            duplicate_name_set[name] = 1\n",
    "        \n",
    "        # for movie match in redial dataset\n",
    "        all_movie_db[key] = concated_meta\n",
    "        imdb_name_with_year_to_id[full_name] = key\n",
    "        imdb_name_without_year_to_id[name] = key\n",
    "        key +=1 \n",
    "\n",
    "# Inspired dataset\n",
    "new_count2 = 0       \n",
    "df = pandas.read_csv(\"movie_database.tsv\", sep = '\\t')\n",
    "for i in range(len(df)):\n",
    "    mid = df['imdb_id'][i]\n",
    "    if mid not in imdb_id_set.keys():\n",
    "        new_count2+=1\n",
    "        imdb_id_set[mid]=1\n",
    "        name = df['title'][i].lower()\n",
    "        year = df['year'][i]\n",
    "        desc = df['short_plot'][i]\n",
    "        full_name = name + ' (' + str(year) + ')'\n",
    "        if name in imdb_name_without_year_to_id.keys():\n",
    "            duplicate_name_set[name] = 1\n",
    "            \n",
    "        concated_meta = str(name) + \" [SEP] \" + \\\n",
    "              str(df['actors'][i]) + \" [SEP] \" + \\\n",
    "              str(df['director'][i]) + \" [SEP] \" + \\\n",
    "              str(df['genre'][i]) + \" [SEP] \" + \\\n",
    "              str(desc)\n",
    "\n",
    "        # for movie match in redial dataset\n",
    "        all_movie_db[key] = concated_meta\n",
    "        imdb_name_with_year_to_id[full_name] = key\n",
    "        imdb_name_without_year_to_id[name] = key\n",
    "        key +=1\n",
    "        \n",
    "# sparql responses\n",
    "sparql_dic = torch.load(\"/local-scratch1/data/by2299/sparql_processed_dic\")\n",
    "\n",
    "for k, v in sparql_dic.items():\n",
    "    all_movie_db[key] = v\n",
    "    if '(' in k:\n",
    "        imdb_name_with_year_to_id[k.lower()] = key\n",
    "        imdb_name_without_year_to_id[k.split('(')[0].strip().lower()] = key\n",
    "    else:\n",
    "        imdb_name_without_year_to_id[k.lower()] = key\n",
    "    key += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "e34da93e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "97694\n",
      "91503\n",
      "4600\n",
      "11345 672\n"
     ]
    }
   ],
   "source": [
    "print(len(imdb_name_with_year_to_id.keys()))\n",
    "print(len(imdb_name_without_year_to_id.keys()))\n",
    "print(len(duplicate_name_set.keys()))\n",
    "print(new_count, new_count2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "86e8d6dc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded 1342 test conversations\n"
     ]
    }
   ],
   "source": [
    "test_data = []\n",
    "for line in open(\"./website-data/redial_dataset/test_data.jsonl\", \"r\"):\n",
    "    test_data.append(json.loads(line))\n",
    "print(\"Loaded {} test conversations\".format(len(test_data)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "eda6c18e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'movieMentions': {'111776': 'Super Troopers (2001)',\n",
       "  '91481': 'Beverly Hills Cop (1984)',\n",
       "  '151656': 'Police Academy  (1984)',\n",
       "  '134643': 'American Pie  (1999)',\n",
       "  '192131': 'American Pie ',\n",
       "  '124771': '48 Hrs. (1982)',\n",
       "  '94688': 'Police Academy 2: Their First Assignment (1985)',\n",
       "  '101794': 'Lethal Weapon (1987)'},\n",
       " 'respondentQuestions': {'111776': {'suggested': 0, 'seen': 1, 'liked': 1},\n",
       "  '91481': {'suggested': 1, 'seen': 2, 'liked': 2},\n",
       "  '151656': {'suggested': 1, 'seen': 0, 'liked': 1},\n",
       "  '134643': {'suggested': 0, 'seen': 1, 'liked': 1},\n",
       "  '192131': {'suggested': 0, 'seen': 1, 'liked': 1},\n",
       "  '124771': {'suggested': 1, 'seen': 2, 'liked': 2},\n",
       "  '94688': {'suggested': 1, 'seen': 0, 'liked': 1},\n",
       "  '101794': {'suggested': 1, 'seen': 0, 'liked': 2}},\n",
       " 'messages': [{'timeOffset': 0,\n",
       "   'text': 'Hi I am looking for a movie like @111776',\n",
       "   'senderWorkerId': 956,\n",
       "   'messageId': 204171},\n",
       "  {'timeOffset': 48,\n",
       "   'text': 'You should watch @151656',\n",
       "   'senderWorkerId': 957,\n",
       "   'messageId': 204172},\n",
       "  {'timeOffset': 90,\n",
       "   'text': 'Is that a great one? I have never seen it. I have seen @192131',\n",
       "   'senderWorkerId': 956,\n",
       "   'messageId': 204173},\n",
       "  {'timeOffset': 122,\n",
       "   'text': 'I mean @134643',\n",
       "   'senderWorkerId': 956,\n",
       "   'messageId': 204174},\n",
       "  {'timeOffset': 180,\n",
       "   'text': 'Yes @151656 is very funny and so is @94688',\n",
       "   'senderWorkerId': 957,\n",
       "   'messageId': 204175},\n",
       "  {'timeOffset': 199,\n",
       "   'text': 'It sounds like I need to check them out',\n",
       "   'senderWorkerId': 956,\n",
       "   'messageId': 204176},\n",
       "  {'timeOffset': 219,\n",
       "   'text': 'yes you will enjoy them',\n",
       "   'senderWorkerId': 957,\n",
       "   'messageId': 204177},\n",
       "  {'timeOffset': 253,\n",
       "   'text': 'I appreciate your time. I will need to check those out. Are there any others you would recommend?',\n",
       "   'senderWorkerId': 956,\n",
       "   'messageId': 204178},\n",
       "  {'timeOffset': 297,\n",
       "   'text': 'yes @101794',\n",
       "   'senderWorkerId': 957,\n",
       "   'messageId': 204179},\n",
       "  {'timeOffset': 311,\n",
       "   'text': 'Thank you i will watch that too',\n",
       "   'senderWorkerId': 956,\n",
       "   'messageId': 204180},\n",
       "  {'timeOffset': 312,\n",
       "   'text': 'and also @91481',\n",
       "   'senderWorkerId': 957,\n",
       "   'messageId': 204181},\n",
       "  {'timeOffset': 326,\n",
       "   'text': 'Thanks for the suggestions.',\n",
       "   'senderWorkerId': 956,\n",
       "   'messageId': 204182},\n",
       "  {'timeOffset': 341,\n",
       "   'text': 'you are welcome',\n",
       "   'senderWorkerId': 957,\n",
       "   'messageId': 204183},\n",
       "  {'timeOffset': 408,\n",
       "   'text': 'and also @124771',\n",
       "   'senderWorkerId': 957,\n",
       "   'messageId': 204184},\n",
       "  {'timeOffset': 518,\n",
       "   'text': 'thanks goodbye',\n",
       "   'senderWorkerId': 956,\n",
       "   'messageId': 204185}],\n",
       " 'conversationId': '20001',\n",
       " 'respondentWorkerId': 957,\n",
       " 'initiatorWorkerId': 956,\n",
       " 'initiatorQuestions': {'111776': {'suggested': 0, 'seen': 1, 'liked': 1},\n",
       "  '91481': {'suggested': 1, 'seen': 2, 'liked': 2},\n",
       "  '151656': {'suggested': 1, 'seen': 0, 'liked': 1},\n",
       "  '134643': {'suggested': 0, 'seen': 1, 'liked': 1},\n",
       "  '192131': {'suggested': 0, 'seen': 1, 'liked': 1},\n",
       "  '124771': {'suggested': 1, 'seen': 2, 'liked': 2},\n",
       "  '94688': {'suggested': 1, 'seen': 0, 'liked': 1},\n",
       "  '101794': {'suggested': 0, 'seen': 2, 'liked': 2}}}"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_data[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "8e3efce6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def refine_movie_name(old_name):\n",
    "    \"\"\"\n",
    "    regularize move name to match the database. default is name + space + (year), if without year, return name only\n",
    "    return (new name, flag)\n",
    "    flag 0 means without year\n",
    "    flag 1 means with year\n",
    "    \"\"\"\n",
    "    year = re.findall(r\"\\(\\d\\d\\d\\d\\)\", old_name)\n",
    "    if len(year) > 0:\n",
    "        year = year[0]\n",
    "        name = old_name.split(year)[0]\n",
    "        movie_name = name.strip() + ' ' + year\n",
    "        return (movie_name, 1)\n",
    "    else:\n",
    "        movie_name = old_name.strip()\n",
    "        return (movie_name, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "ddda26c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# def movie_search(mv_name_with_year_to_id, mv_name_without_year_to_id, movie_name):\n",
    "#     movie_name, year_flag = refine_movie_name(movie_name)\n",
    "    \n",
    "#     if year_flag == 1 and movie_name.lower() in mv_name_with_year_to_id.keys():\n",
    "#         our_id = mv_name_with_year_to_id[movie_name.lower()]\n",
    "\n",
    "#     elif year_flag == 0 and movie_name.lower() in imdb_name_without_year_to_id.keys():  \n",
    "#         our_id = imdb_name_without_year_to_id[movie_name.lower()]\n",
    "#         if movie_name.lower() in duplicate_name_set.keys():\n",
    "#             mentioned_is_in_duplicate += 1\n",
    "            \n",
    "#     elif year_flag == 1 and movie_name.split(' (')[0].lower() in mv_name_without_year_to_id.keys():\n",
    "#         our_id = mv_name_without_year_to_id[movie_name.split(' (')[0].lower()]\n",
    "#         if movie_name.split(' (')[0].lower() in duplicate_name_set.keys():\n",
    "#             mentioned_is_in_duplicate += 1\n",
    "#     else:\n",
    "#         our_id = None\n",
    "#         not_exist_count += 1\n",
    "#         not_exist_metioned_movie_db[movie_name] = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "6a770a2e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded 1342 test conversations\n",
      "Loaded 10006 train conversations\n"
     ]
    }
   ],
   "source": [
    "metioned_movie_db = {} # ID to descripiton\n",
    "not_exist_metioned_movie_db = {}\n",
    "all_mentioned_movie = {} # unique movie mentioned\n",
    "mentioned_is_in_duplicate = 0\n",
    "\n",
    "mention_count = 0 # can contain duplicate\n",
    "not_exist_count = 0 # movie not in the database\n",
    "\n",
    "test_data = []\n",
    "for line in open(\"./website-data/redial_dataset/test_data.jsonl\", \"r\"):\n",
    "    test_data.append(json.loads(line))\n",
    "print(\"Loaded {} test conversations\".format(len(test_data)))\n",
    "\n",
    "train_data = []\n",
    "for line in open(\"./website-data/redial_dataset/train_data.jsonl\", \"r\"):\n",
    "    train_data.append(json.loads(line))\n",
    "print(\"Loaded {} train conversations\".format(len(train_data)))\n",
    "\n",
    "all_data = train_data + test_data\n",
    "\n",
    "for i in range(len(all_data)):\n",
    "    temp = all_data[i]\n",
    "    mentioned_movies = temp['movieMentions']\n",
    "    if len(mentioned_movies) == 0:\n",
    "        continue\n",
    "    for key in list(mentioned_movies.keys()):\n",
    "        movie_name = mentioned_movies[key]\n",
    "        if movie_name is None:\n",
    "            continue\n",
    "        movie_name, year_flag = refine_movie_name(movie_name)\n",
    "        all_mentioned_movie[movie_name] = 1\n",
    "        mention_count += 1\n",
    "        \n",
    "        if year_flag == 1 and movie_name.lower() in imdb_name_with_year_to_id.keys():\n",
    "            our_id = imdb_name_with_year_to_id[movie_name.lower()]\n",
    "            metioned_movie_db[our_id] = all_movie_db[our_id]\n",
    "            \n",
    "        elif year_flag == 0 and movie_name.lower() in imdb_name_without_year_to_id.keys():  \n",
    "            our_id = imdb_name_without_year_to_id[movie_name.lower()]\n",
    "            metioned_movie_db[our_id] = all_movie_db[our_id]\n",
    "            if movie_name.lower() in duplicate_name_set.keys():\n",
    "                mentioned_is_in_duplicate += 1\n",
    "                \n",
    "        elif year_flag == 1 and movie_name.split(' (')[0].lower() in imdb_name_without_year_to_id.keys():\n",
    "            our_id = imdb_name_without_year_to_id[movie_name.split(' (')[0].lower()]\n",
    "            metioned_movie_db[our_id] = all_movie_db[our_id]\n",
    "            if movie_name.split(' (')[0].lower() in duplicate_name_set.keys():\n",
    "                mentioned_is_in_duplicate += 1\n",
    "        else:\n",
    "            not_exist_count += 1\n",
    "            not_exist_metioned_movie_db[movie_name] = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "6fc536b1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "6629 6222 180\n",
      "60070 1171 1443\n"
     ]
    }
   ],
   "source": [
    "print(len(all_mentioned_movie.keys()), len(metioned_movie_db.keys()), len(not_exist_metioned_movie_db.keys()))\n",
    "print(mention_count, not_exist_count, mentioned_is_in_duplicate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "8be83f9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def meta_with_id(m_id, movie_name_no_year):\n",
    "    movie_deets = movie.details(m_id)\n",
    "    \n",
    "    #actors\n",
    "    actors = ''\n",
    "    for i, cast in enumerate(movie_deets['casts']['cast']):\n",
    "        if i <= 1:\n",
    "            actors += cast['name'] + ', '\n",
    "        else:\n",
    "            actors += cast['name']\n",
    "            break\n",
    "    \n",
    "    # genre\n",
    "    genres = ''\n",
    "    for i, genre in enumerate(movie_deets[\"genres\"]):\n",
    "        genres += genre['name'] + ', '\n",
    "        \n",
    "    # overview\n",
    "    overview = movie_deets['overview']\n",
    "    \n",
    "    return movie_name_no_year + ' [SEP] ' + actors + ' [SEP] [SEP] ' + genres + ' [SEP] ' + overview"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "d0ed335a",
   "metadata": {},
   "outputs": [],
   "source": [
    "for key in not_exist_metioned_movie_db.keys():\n",
    "    our_id = len(all_movie_db)\n",
    "    movie_name, year_flag = refine_movie_name(key)\n",
    "    movie_name_no_year = None\n",
    "    if year_flag == 1:\n",
    "        movie_name_no_year = movie_name.split(' (')[0].lower()\n",
    "    else:\n",
    "        movie_name_no_year = movie_name\n",
    "    if movie_name_no_year == \"\": continue\n",
    "    results = movie.search(movie_name_no_year)\n",
    "    if len(results) == 1:\n",
    "        all_movie_db[our_id] = meta_with_id(results[0]['id'], movie_name_no_year) #results[0][\"overview\"]\n",
    "        imdb_name_with_year_to_id[movie_name.lower()] = our_id\n",
    "        imdb_name_without_year_to_id[movie_name_no_year.lower()] = our_id\n",
    "        continue\n",
    "    valid = False\n",
    "    for result in results:\n",
    "        o_title = result['original_title']\n",
    "        title = result['title']\n",
    "        # year = result['release_date'].split('-')[0]\n",
    "        if o_title.lower() == movie_name_no_year.lower() or title.lower() == movie_name_no_year.lower():\n",
    "            valid = True\n",
    "            \n",
    "            m_id = result['id']\n",
    "            \n",
    "            all_movie_db[our_id] = meta_with_id(m_id, movie_name_no_year)\n",
    "            imdb_name_with_year_to_id[movie_name.lower()] = our_id\n",
    "            imdb_name_without_year_to_id[movie_name_no_year.lower()] = our_id\n",
    "            break\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "359b7255",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import csv\n",
    "\n",
    "# with open('unfound_movies.csv', mode='w') as unfound_file:\n",
    "#     employee_writer = csv.writer(unfound_file, delimiter=',', quotechar='\"', quoting=csv.QUOTE_MINIMAL)\n",
    "\n",
    "#     for k, v in not_exist_metioned_movie_db.items():\n",
    "#         employee_writer.writerow([str(k)])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "19d8a7ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "multi_rec_count = 0 # more than one rec in on sentence\n",
    "not_label_count = 0 # movie without {suggested} labels\n",
    "suggest_not_exist_count = 0 # movie not in the database\n",
    "suggested_is_in_duplicate = 0\n",
    "not_exist_suggested_movie_db = {}\n",
    "\n",
    "rec_count = 0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c60b08f9",
   "metadata": {},
   "source": [
    "## test data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "3d8b929e",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded 1342 test conversations\n"
     ]
    }
   ],
   "source": [
    "test_data = []\n",
    "for line in open(\"./website-data/redial_dataset/test_data.jsonl\", \"r\"):\n",
    "    test_data.append(json.loads(line))\n",
    "print(\"Loaded {} test conversations\".format(len(test_data)))\n",
    "\n",
    "all_data_test = [] #result\n",
    "\n",
    "for i in range(len(test_data)):\n",
    "    temp = test_data[i]\n",
    "    mentioned_movies = temp['movieMentions']\n",
    "    status_movies = temp['respondentQuestions']\n",
    "    messages = temp['messages']\n",
    "    \n",
    "    # A is recommender B is customer\n",
    "    AB_dict = {}\n",
    "    AB_dict[temp['respondentWorkerId']] = 'A'\n",
    "    AB_dict[temp['initiatorWorkerId']] = 'B'\n",
    "\n",
    "    # dialog temp\n",
    "    this_dialog_tmp = []\n",
    "    for j in range(len(messages)):\n",
    "        message = messages[j]\n",
    "        userID = message['senderWorkerId']\n",
    "        text = message['text']\n",
    "        this_dialog_tmp.append([userID, text])\n",
    "\n",
    "    # merge consecutive sentences from the same user into one sentence\n",
    "    this_dialog = []\n",
    "    j = 0\n",
    "    while j < len(this_dialog_tmp):\n",
    "        userID = this_dialog_tmp[j][0]\n",
    "        text = this_dialog_tmp[j][1]\n",
    "        k = 1\n",
    "        while j + k < len(this_dialog_tmp) and userID == this_dialog_tmp[j+k][0]:\n",
    "            text = text + ' ' + this_dialog_tmp[j+k][1]\n",
    "            k += 1\n",
    "        j = j + k   \n",
    "        this_dialog.append(AB_dict[userID] + \": \" + text) \n",
    "        \n",
    "    # extract @movie ID, check which is recommendaiton\n",
    "    this_dialog_post = []\n",
    "\n",
    "    for item in this_dialog:\n",
    "        this_setence_movie_id = re.findall(r\"@\\d+\", item)\n",
    "        recommended_list = []\n",
    "        for m in this_setence_movie_id:\n",
    "            m_id = m.split('@')[1]\n",
    "            \n",
    "            if len(status_movies) > 0 and m_id in status_movies.keys():# and status_movies[m_id]['suggested']==1:\n",
    "                recommended_list.append(m_id)\n",
    "            \n",
    "            # replace id with name in the dialog\n",
    "            movie_name = mentioned_movies[m_id]\n",
    "            movie_name, year_flag = refine_movie_name(movie_name)\n",
    "#             item = item.replace(str(m), movie_name)\n",
    "            item = item.replace(str(m), \"[MOVIE_ID]\") # use placeholder\n",
    "        our_ids = []\n",
    "        if len(recommended_list) > 0:\n",
    "            for rec_id in range(len(recommended_list)):\n",
    "                movie_name = mentioned_movies[recommended_list[rec_id]]\n",
    "                movie_name, year_flag = refine_movie_name(movie_name)\n",
    "                if year_flag == 1 and movie_name.lower() in imdb_name_with_year_to_id.keys():\n",
    "                    our_ids.append( imdb_name_with_year_to_id[movie_name.lower()] )\n",
    "\n",
    "                elif year_flag == 0 and movie_name.lower() in imdb_name_without_year_to_id.keys():  \n",
    "                    our_ids.append( imdb_name_without_year_to_id[movie_name.lower()] )\n",
    "                    if movie_name.lower() in duplicate_name_set.keys():\n",
    "                        suggested_is_in_duplicate += 1\n",
    "\n",
    "                elif year_flag == 1 and movie_name.split(' (')[0].lower() in imdb_name_without_year_to_id.keys():\n",
    "                    our_ids.append( imdb_name_without_year_to_id[movie_name.split(' (')[0].lower()] )\n",
    "                    if movie_name.split(' (')[0].lower() in duplicate_name_set.keys():\n",
    "                        suggested_is_in_duplicate += 1\n",
    "                else:\n",
    "                    not_exist_suggested_movie_db[movie_name]=1\n",
    "                    suggest_not_exist_count += 1\n",
    "#                     our_ids = None\n",
    "\n",
    "            if len(recommended_list) > 1:\n",
    "                multi_rec_count += 1\n",
    "        else:\n",
    "            our_ids = None\n",
    "        if our_ids == []: our_ids = None\n",
    "        this_dialog_post.append((item, our_ids))\n",
    "    \n",
    "    all_data_test.append(this_dialog_post)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "3ee93e03",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "52 179\n",
      "0 1747\n",
      "154\n"
     ]
    }
   ],
   "source": [
    "print(len(not_exist_suggested_movie_db), suggest_not_exist_count)\n",
    "print(rec_count, multi_rec_count)\n",
    "print(suggested_is_in_duplicate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "dd7df6d5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "52 179\n",
      "0 1747\n",
      "154\n"
     ]
    }
   ],
   "source": [
    "print(len(not_exist_suggested_movie_db), suggest_not_exist_count)\n",
    "print(rec_count, multi_rec_count)\n",
    "print(suggested_is_in_duplicate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "2f4bbbf1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('A: Hello have you seen [MOVIE_ID] ? It is one of my favorite movies [MOVIE_ID]',\n",
       "  [28381, 28381]),\n",
       " ('B: Hi yes, that is a classic! I love movies like that Also, [MOVIE_ID]',\n",
       "  [25300]),\n",
       " ('A: all the mob movies are great like [MOVIE_ID]', [28769]),\n",
       " ('B: another one is [MOVIE_ID] I tend to enjoy gangster movies', [28769]),\n",
       " ('A: nice me too', None),\n",
       " ('B: yes!!', None),\n",
       " ('A: have you seen the [MOVIE_ID] ? Its the ultimate classic', [43116]),\n",
       " ('B: I actually hated [MOVIE_ID]  and never watched any of the others.  It was just too slow lol',\n",
       "  [15528]),\n",
       " ('A: they do take some time to finish, theyre all 3 hours plus', None),\n",
       " ('B: Any other good suggestions? Like the movie [MOVIE_ID] for me that was another classic',\n",
       "  [30671]),\n",
       " (\"A: if you haven't seen it yet, [MOVIE_ID] is great\", [45462]),\n",
       " ('B: Thanks, sounds good.', None),\n",
       " ('A: thoes are all good', None),\n",
       " ('B: Take care!', None),\n",
       " ('A: thanks', None)]"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "all_data_test[108]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08f0681d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51c9393e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "7965b4db",
   "metadata": {},
   "source": [
    "## train data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "a936a873",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded 10006 train conversations\n"
     ]
    }
   ],
   "source": [
    "train_data = []\n",
    "for line in open(\"./website-data/redial_dataset/train_data.jsonl\", \"r\"):\n",
    "    train_data.append(json.loads(line))\n",
    "print(\"Loaded {} train conversations\".format(len(train_data)))\n",
    "\n",
    "all_data_train = [] #result\n",
    "\n",
    "for i in range(len(train_data)):\n",
    "    temp = train_data[i]\n",
    "    mentioned_movies = temp['movieMentions']\n",
    "    status_movies = temp['respondentQuestions']\n",
    "    messages = temp['messages']\n",
    "    \n",
    "    # A is recommender B is customer\n",
    "    AB_dict = {}\n",
    "    AB_dict[temp['respondentWorkerId']] = 'A'\n",
    "    AB_dict[temp['initiatorWorkerId']] = 'B'\n",
    "\n",
    "    # dialog temp\n",
    "    this_dialog_tmp = []\n",
    "    for j in range(len(messages)):\n",
    "        message = messages[j]\n",
    "        userID = message['senderWorkerId']\n",
    "        text = message['text']\n",
    "        this_dialog_tmp.append([userID, text])\n",
    "\n",
    "    # merge consecutive sentences from the same user into one sentence\n",
    "    this_dialog = []\n",
    "    j = 0\n",
    "    while j < len(this_dialog_tmp):\n",
    "        userID = this_dialog_tmp[j][0]\n",
    "        text = this_dialog_tmp[j][1]\n",
    "        k = 1\n",
    "        while j + k < len(this_dialog_tmp) and userID == this_dialog_tmp[j+k][0]:\n",
    "            text = text + ' ' + this_dialog_tmp[j+k][1]\n",
    "            k += 1\n",
    "        j = j + k   \n",
    "        this_dialog.append(AB_dict[userID] + \": \" + text) \n",
    "        \n",
    "    # extract @movie ID, check which is recommendaiton\n",
    "    this_dialog_post = []\n",
    "\n",
    "    for item in this_dialog:\n",
    "        this_setence_movie_id = re.findall(r\"@\\d+\", item)\n",
    "        recommended_list = []\n",
    "        for m in this_setence_movie_id:\n",
    "            m_id = m.split('@')[1]\n",
    "            \n",
    "            if len(status_movies) > 0 and m_id in status_movies.keys():# and status_movies[m_id]['suggested']==1:\n",
    "                recommended_list.append(m_id)\n",
    "            \n",
    "            # replace id with name in the dialog\n",
    "            if m_id in mentioned_movies.keys():\n",
    "                movie_name = mentioned_movies[m_id]\n",
    "                movie_name, year_flag = refine_movie_name(movie_name)\n",
    "#                 item = item.replace(str(m), movie_name)\n",
    "                item = item.replace(str(m), \"[MOVIE_ID]\") # use placeholder\n",
    "        \n",
    "        our_ids = []\n",
    "        if len(recommended_list) > 0:\n",
    "            for rec_id in range(len(recommended_list)):\n",
    "                rec_count += 1\n",
    "                movie_name = mentioned_movies[recommended_list[rec_id]]\n",
    "                movie_name, year_flag = refine_movie_name(movie_name)\n",
    "                if year_flag == 1 and movie_name.lower() in imdb_name_with_year_to_id.keys():\n",
    "                    our_ids.append( imdb_name_with_year_to_id[movie_name.lower()] )\n",
    "\n",
    "                elif year_flag == 0 and movie_name.lower() in imdb_name_without_year_to_id.keys():  \n",
    "                    our_ids.append( imdb_name_without_year_to_id[movie_name.lower()] )\n",
    "                    if movie_name.lower() in duplicate_name_set.keys():\n",
    "                        suggested_is_in_duplicate += 1\n",
    "\n",
    "                elif year_flag == 1 and movie_name.split(' (')[0].lower() in imdb_name_without_year_to_id.keys():\n",
    "                    our_ids.append( imdb_name_without_year_to_id[movie_name.split(' (')[0].lower()] )\n",
    "                    if movie_name.split(' (')[0].lower() in duplicate_name_set.keys():\n",
    "                        suggested_is_in_duplicate += 1\n",
    "                else:\n",
    "                    not_exist_suggested_movie_db[movie_name]=1\n",
    "                    suggest_not_exist_count += 1\n",
    "#                     our_id = None\n",
    "            \n",
    "            if len(recommended_list) > 1:\n",
    "                multi_rec_count += 1\n",
    "        else:\n",
    "            our_ids = None\n",
    "        if our_ids == []: our_ids = None\n",
    "        this_dialog_post.append((item, our_ids))\n",
    "    \n",
    "    all_data_train.append(this_dialog_post)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "bf377bd0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "179 1351\n",
      "63669 14322\n",
      "1692\n"
     ]
    }
   ],
   "source": [
    "print(len(not_exist_suggested_movie_db), suggest_not_exist_count)\n",
    "print(rec_count, multi_rec_count)\n",
    "print(suggested_is_in_duplicate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "1736f153",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "179 1351\n",
      "63669 14322\n",
      "1692\n"
     ]
    }
   ],
   "source": [
    "print(len(not_exist_suggested_movie_db), suggest_not_exist_count)\n",
    "print(rec_count, multi_rec_count)\n",
    "print(suggested_is_in_duplicate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "1d4571d3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('B: HI there~', None),\n",
       " ('A: Hi what kind of movies doe you like to watch?', None),\n",
       " ('B: I like mostly anything. Especially muscials and comedy Im not fond of movies like [MOVIE_ID]',\n",
       "  [89461]),\n",
       " ('A: I love comedy movies have you seen [MOVIE_ID] it is very funny',\n",
       "  [59079]),\n",
       " ('B: I have not seen that yet.Is that with the Rock&quot;', None),\n",
       " ('A: No it has Will Ferrel and Mark Wahlberg', None),\n",
       " ('B: Oh, ok, and they just made a second one too right?', None),\n",
       " ('A: [MOVIE_ID] has the rock and Kevin Hart you may like that Yeah they did it is suppose to be even funnier.',\n",
       "  [58635]),\n",
       " ('B: I did see [MOVIE_ID] and liked that a lot!', [58635]),\n",
       " (\"A: Then you should definitely see daddy's home\", None),\n",
       " ('B: do you have one more suggestion before we go?', None),\n",
       " ('A: Yeah [MOVIE_ID] and [MOVIE_ID] were great Kevin Hart movies',\n",
       "  [68914, 57792]),\n",
       " ('B: Great, I still have yet to see those. makes me laugh when he was on Conan with Ice Cube promoting htat thank you! we can submit now and fill out the movie forms',\n",
       "  None),\n",
       " ('A: Ok great bye.', None)]"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "all_data_train[12]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e07face",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "c6d5b9c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "cleaned_data_train = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "60ad4654",
   "metadata": {},
   "outputs": [],
   "source": [
    "for data_train in all_data_train:\n",
    "    data_point = []\n",
    "    for u, r in data_train:\n",
    "        if u[0] == \"B\":\n",
    "            data_point.append((u, r))\n",
    "        else:\n",
    "            data_point.append((u, r))\n",
    "    cleaned_data_train.append(data_point)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "ea27490c",
   "metadata": {},
   "outputs": [],
   "source": [
    "super_cleaned_data_train = []\n",
    "for data_train in cleaned_data_train:\n",
    "    valid = True\n",
    "    for u, r in data_train: \n",
    "        if r == None: continue\n",
    "        else:\n",
    "            for item in r:\n",
    "                if item not in metioned_movie_db.keys():\n",
    "                    valid = False\n",
    "                    break\n",
    "    if valid:\n",
    "        super_cleaned_data_train.append(data_train)\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "281bebb7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10006"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(super_cleaned_data_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "bd42c5c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "cleaned_data_test = []\n",
    "for data_test in all_data_test:\n",
    "    data_point = []\n",
    "    for u, r in data_test:\n",
    "        if u[0] == \"B\":\n",
    "            data_point.append((u, r))\n",
    "        else:\n",
    "            data_point.append((u, r))\n",
    "    cleaned_data_test.append(data_point)\n",
    "    \n",
    "super_cleaned_data_test = []\n",
    "for data_test in cleaned_data_test:\n",
    "    valid = True\n",
    "    for u, r in data_test: \n",
    "        if r == None: continue\n",
    "        else:\n",
    "            for item in r:\n",
    "                if item not in metioned_movie_db.keys():\n",
    "                    valid = False\n",
    "                    break\n",
    "    if valid:\n",
    "        super_cleaned_data_test.append(data_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "7794a070",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1342"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(super_cleaned_data_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "47bba14f",
   "metadata": {},
   "outputs": [],
   "source": [
    "movie_db_60000 = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "521a7095",
   "metadata": {},
   "outputs": [],
   "source": [
    "for k, v in metioned_movie_db.items():\n",
    "    movie_db_60000[k] = v\n",
    "    all_movie_db.pop(k)\n",
    "\n",
    "for k, v in all_movie_db.items():\n",
    "    movie_db_60000[k] = v\n",
    "    if len(movie_db_60000) >= 60000:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "7965a3e6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "60000"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(movie_db_60000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "09b8e665",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(super_cleaned_data_train, \"/local-scratch1/data/by2299/redial_full_train_placeholder\")\n",
    "torch.save(super_cleaned_data_test, \"/local-scratch1/data/by2299/redial_full_test_placeholder\")\n",
    "torch.save(metioned_movie_db, \"/local-scratch1/data/by2299/redial_full_movie_db_placeholder\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17e47a3e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}