imdb_preprocess.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %mkdir ./data\n",
    "# !wget -O ./data/aclImdb_v1.tar.gz http://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz\n",
    "# !tar -zxf ./data/aclImdb_v1.tar.gz -C ./data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "\n",
    "def read_imdb_data(data_dir='./data/aclImdb'):\n",
    "    data = {}\n",
    "    labels = {}\n",
    "    \n",
    "    for data_type in ['train', 'test']:\n",
    "        data[data_type] = {}\n",
    "        labels[data_type] = {}\n",
    "        \n",
    "        for sentiment in ['pos', 'neg']:\n",
    "            data[data_type][sentiment] = []\n",
    "            labels[data_type][sentiment] = []\n",
    "            \n",
    "            path = os.path.join(data_dir, data_type, sentiment, '*.txt')\n",
    "            files = glob.glob(path)\n",
    "            \n",
    "            for f in files:\n",
    "                with open(f) as review:\n",
    "                    data[data_type][sentiment].append(review.read())\n",
    "                    labels[data_type][sentiment].append(1 if sentiment == 'pos' else 0)\n",
    "                    \n",
    "            assert len(data[data_type][sentiment]) == len(labels[data_type][sentiment]), \\\n",
    "                    \"{}/{} data size does not match labels size\".format(data_type, sentiment)\n",
    "                \n",
    "    return data, labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data, labels = read_imdb_data()\n",
    "print(\"IMDB reviews: train = {} pos / {} neg, test = {} pos / {} neg\".format(\n",
    "            len(data['train']['pos']), len(data['train']['neg']),\n",
    "            len(data['test']['pos']), len(data['test']['neg'])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.utils import shuffle\n",
    "\n",
    "def prepare_imdb_data(data, labels):\n",
    "    \"\"\"Prepare training and test sets from IMDb movie reviews.\"\"\"\n",
    "    \n",
    "    #Combine positive and negative reviews and labels\n",
    "    data_train = data['train']['pos'] + data['train']['neg']\n",
    "    data_test = data['test']['pos'] + data['test']['neg']\n",
    "    labels_train = labels['train']['pos'] + labels['train']['neg']\n",
    "    labels_test = labels['test']['pos'] + labels['test']['neg']\n",
    "    \n",
    "    #Shuffle reviews and corresponding labels within training and test sets\n",
    "    data_train, labels_train = shuffle(data_train, labels_train)\n",
    "    data_test, labels_test = shuffle(data_test, labels_test)\n",
    "    \n",
    "    # Return a unified training data, test data, training labels, test labets\n",
    "    return data_train, data_test, labels_train, labels_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_X, test_X, train_y, test_y = prepare_imdb_data(data, labels)\n",
    "print(\"IMDb reviews (combined): train = {}, test = {}\".format(len(train_X), len(test_X)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nltk\n",
    "from nltk.corpus import stopwords\n",
    "from nltk.stem.porter import *\n",
    "\n",
    "import re\n",
    "from bs4 import BeautifulSoup\n",
    "\n",
    "def review_to_words(review):\n",
    "    nltk.download(\"stopwords\", quiet=True)\n",
    "    stemmer = PorterStemmer()\n",
    "    # TODO: try out SnowBallStemmer\n",
    "    text = BeautifulSoup(review, \"html.parser\").get_text() # Remove HTML tags\n",
    "    text = re.sub(r\"[^a-zA-Z0-9]\", \" \", text.lower()) # Convert to lower case\n",
    "    words = text.split() # Split string into words\n",
    "    words = [w for w in words if w not in stopwords.words(\"english\")] # Remove stopwords\n",
    "    words = [PorterStemmer().stem(w) for w in words] # stem\n",
    "    \n",
    "    return words"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "\n",
    "cache_dir = os.path.join(\"./cache\", \"sentiment_analysis\")  # where to store cache files\n",
    "os.makedirs(cache_dir, exist_ok=True)  # ensure cache directory exists\n",
    "\n",
    "def preprocess_data(data_train, data_test, labels_train, labels_test,\n",
    "                    cache_dir=cache_dir, cache_file=\"preprocessed_data.pkl\"):\n",
    "    \"\"\"Convert each review to words; read from cache if available.\"\"\"\n",
    "\n",
    "    # If cache_file is not None, try to read from it first\n",
    "    cache_data = None\n",
    "    if cache_file is not None:\n",
    "        try:\n",
    "            with open(os.path.join(cache_dir, cache_file), \"rb\") as f:\n",
    "                cache_data = pickle.load(f)\n",
    "            print(\"Read preprocessed data from cache file:\", cache_file)\n",
    "        except:\n",
    "            pass  # unable to read from cache, but that's okay\n",
    "    \n",
    "    # If cache is missing, then do the heavy lifting\n",
    "    if cache_data is None:\n",
    "        words_train = [review_to_words(review) for review in data_train]\n",
    "        words_test = [review_to_words(review) for review in data_test]\n",
    "        \n",
    "        # Write to cache file for future runs\n",
    "        if cache_file is not None:\n",
    "            cache_data = dict(words_train=words_train, words_test=words_test,\n",
    "                              labels_train=labels_train, labels_test=labels_test)\n",
    "            with open(os.path.join(cache_dir, cache_file), \"wb\") as f:\n",
    "                pickle.dump(cache_data, f)\n",
    "            print(\"Wrote preprocessed data to cache file:\", cache_file)\n",
    "    else:\n",
    "        # Unpack data loaded from cache file\n",
    "        words_train, words_test, labels_train, labels_test = (cache_data['words_train'],\n",
    "                cache_data['words_test'], cache_data['labels_train'], cache_data['labels_test'])\n",
    "    \n",
    "    return words_train, words_test, labels_train, labels_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_X, test_X, train_y, test_y = preprocess_data(train_X, test_X, train_y, test_y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "def build_dict(data, vocab_size = 5000):\n",
    "    word_count = {} # A dict storing the words that appear in the reviews along with how often they occur\n",
    "    for sentence in data:\n",
    "        for word in sentence:\n",
    "            if word in word_count.keys():\n",
    "                word_count[word] = word_count[word] + 1\n",
    "            else:\n",
    "                word_count[word] = 1\n",
    "   \n",
    "    sorted_words = {k: v for k, v in sorted(word_count.items(), key=lambda item: item[1], reverse=True)}\n",
    "    sorted_words = [word for word in sorted_words.keys()]\n",
    "    word_dict = {} # This is what we are building, a dictionary that translates words into integers\n",
    "\n",
    "    for idx, word in enumerate(sorted_words[:vocab_size - 2]): # The -2 is so that we save room for the 'no word'\n",
    "        word_dict[word] = idx + 2                              # 'infrequent' labels\n",
    "        \n",
    "    return word_dict\n",
    "\n",
    "def convert_and_pad(word_dict, sentence, pad=500):\n",
    "    NOWORD = 0 # We will use 0 to represent the 'no word' category\n",
    "    INFREQ = 1 # and we use 1 to represent the infrequent words, i.e., words not appearing in word_dict\n",
    "    \n",
    "    working_sentence = [NOWORD] * pad\n",
    "    \n",
    "    for word_index, word in enumerate(sentence[:pad]):\n",
    "        if word in word_dict:\n",
    "            working_sentence[word_index] = word_dict[word]\n",
    "        else:\n",
    "            working_sentence[word_index] = INFREQ\n",
    "            \n",
    "    return working_sentence, min(len(sentence), pad)\n",
    "\n",
    "def convert_and_pad_data(word_dict, data, pad=500):\n",
    "    result = []\n",
    "    lengths = []\n",
    "    \n",
    "    for sentence in data:\n",
    "        converted, leng = convert_and_pad(word_dict, sentence, pad)\n",
    "        result.append(converted)\n",
    "        lengths.append(leng)\n",
    "        \n",
    "    return np.array(result), np.array(lengths)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_X, test_X, train_y, test_y = preprocess_data()\n",
    "type(train_X), type(train_y)\n",
    "train_X.extend(test_X)\n",
    "train_y.extend(test_y)\n",
    "len(train_X), len(train_y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "type(train_X), type(train_y)\n",
    "temp = list(zip(train_X, train_y))\n",
    "random.shuffle(temp)\n",
    "train_X, train_y = zip(*temp)\n",
    "# res1 and res2 come out as tuples, and so must be converted to lists.\n",
    "train_X, train_y = list(train_X), list(train_y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "type(train_X), type(train_y)\n",
    "print(len(train_X))\n",
    "train_list = train_X[:30000]\n",
    "train_labels_list = train_y[:30000]\n",
    "test_list = train_X[30000:]\n",
    "test_labels_list = train_y[30000:]\n",
    "print(len(train_list), len(train_labels_list), len(test_list), len(test_labels_list))\n",
    "print(train_list[0], test_list[0], train_labels_list[0], test_labels_list[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "word_dict = build_dict(train_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dir = './data/pytorch' # The folder we will use for storing data\n",
    "if not os.path.exists(data_dir): # Make sure that the folder exists\n",
    "    os.makedirs(data_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join(data_dir, 'word_dict_qrnn.pkl'), \"wb\") as f:\n",
    "    pickle.dump(word_dict, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_X, train_X_len = convert_and_pad_data(word_dict, train_list)\n",
    "test_X, test_X_len = convert_and_pad_data(word_dict, test_list)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(train_X.shape, len(train_X_len))\n",
    "print(test_X.shape, len(test_X_len))\n",
    "print(train_y[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "pd.concat([pd.DataFrame(train_labels_list), pd.DataFrame(train_X_len), pd.DataFrame(train_X)], axis=1).to_csv(os.path.join(data_dir, 'train_qrnn.csv'), header=False, index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.concat([pd.DataFrame(test_labels_list), pd.DataFrame(test_X_len), pd.DataFrame(test_X)], axis=1).to_csv(os.path.join(data_dir, 'test_qrnn.csv'), header=False, index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "my_env2_gpu",
   "language": "python",
   "name": "my_env2_gpu"
  },
  "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": 4
}