tf-idf.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Source:\n",
    "https://towardsdatascience.com/how-i-used-machine-learning-to-classify-emails-and-turn-them-into-insights-efed37c1e66"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import gzip\n",
    "import os\n",
    "import re\n",
    "import shutil\n",
    "import sqlite3\n",
    "from email import policy\n",
    "from email.parser import HeaderParser\n",
    "import numpy as np\n",
    "from pathlib import Path\n",
    "from typing import Dict, List\n",
    "\n",
    "import pandas as pd\n",
    "from matplotlib import pyplot\n",
    "from sklearn.cluster import KMeans, MiniBatchKMeans\n",
    "from sklearn.decomposition import PCA\n",
    "from sklearn.feature_extraction.text import TfidfVectorizer, ENGLISH_STOP_WORDS\n",
    "from tqdm.notebook import tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Things that won't change"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# To be removed from message bodies\n",
    "LICENSE = '''EDRM Enron Email Data Set has been produced in EML, PST and NSF format by ZL Technologies, Inc. This Data Set is licensed under a Creative Commons Attribution 3.0 United States License <http://creativecommons.org/licenses/by/3.0/us/> . To provide attribution, please cite to \"ZL Technologies, Inc. (http://www.zlti.com).\"'''"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Things that can change"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Study sample size\n",
    "TOTAL_MESSAGES = 10_000\n",
    "\n",
    "# Batch size when reading messages from the database\n",
    "MESSAGE_BATCH_SIZE = 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Input file for this study\n",
    "# Generated from PST files, e.g: `ratom report -pvm /tmp/libratom/test_data/RevisedEDRMv1_Complete/ -o edrm_subset.sqlite3`\n",
    "DB_FILE = Path(\"data/edrm_subset.sqlite3.gz\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Utility functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cleanup_message_body(body: str) -> str:\n",
    "    \"\"\"\n",
    "    Removes non content from message body\n",
    "    \"\"\"\n",
    "    \n",
    "    license = re.escape(LICENSE)\n",
    "    \n",
    "    # Also remove separators around license\n",
    "    separator = re.escape('***********')\n",
    "\n",
    "    body = re.sub(f\"{separator}[\\s]*{license}[\\s]*{separator}\", \"\", body, flags=re.MULTILINE)\n",
    "\n",
    "    return body.strip()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_messages_from_db(db_file: str) -> Dict[str, List]:\n",
    "    \"\"\"\n",
    "    Returns a dict of lists to feed to a dataframe\n",
    "    \"\"\"\n",
    "\n",
    "    header_parser = HeaderParser(policy=policy.default)\n",
    "\n",
    "    # The keys here will become dataframe column names\n",
    "    # Use bracket notation with 'from' and other python reserved keywords, e.g: df['from']\n",
    "    messages = {\n",
    "        'from': [],\n",
    "        'to': [],\n",
    "        'subject': [],\n",
    "        'body': [],\n",
    "    }\n",
    "\n",
    "    with sqlite3.connect(db_file) as conn:\n",
    "        cursor = conn.cursor()\n",
    "\n",
    "        # Our base DB query\n",
    "        cursor.execute(f\"select headers, body from message\")\n",
    "\n",
    "        with tqdm(desc='messages', unit='msg', total=TOTAL_MESSAGES) as pbar:\n",
    "\n",
    "            # Fetch batches of messages\n",
    "            for batch in iter(lambda: cursor.fetchmany(MESSAGE_BATCH_SIZE), []):\n",
    "            \n",
    "                # Iterate over messages until we have the right amount of useful messages\n",
    "                for raw_headers, body in batch:\n",
    "                    try:\n",
    "                        # Parse the headers\n",
    "                        headers = header_parser.parsestr(raw_headers)\n",
    "                        sender = headers['from'] or headers['sender'] or headers['return-path'] or headers['reply-to']\n",
    "                        recipient = headers['to']\n",
    "                        \n",
    "                        # Skip this message if we don't have both a sender and a recipient\n",
    "                        if not (sender and recipient):\n",
    "                            continue\n",
    "\n",
    "                        # Sanitize the body\n",
    "                        body = cleanup_message_body(body)\n",
    "\n",
    "                        # Skip message if body is empty\n",
    "                        if not body:\n",
    "                            continue\n",
    "\n",
    "                        # Add message contents to partial results\n",
    "                        messages['from'].append(sender)\n",
    "                        messages['to'].append(recipient)\n",
    "                        messages['subject'].append(headers['subject'] or '')  # Possibly blank\n",
    "                        messages['body'].append(body)\n",
    "\n",
    "                        # Update progress\n",
    "                        pbar.update()\n",
    "                        \n",
    "                        # Have we reached our sample size?\n",
    "                        if pbar.n >= pbar.total:\n",
    "                            return messages\n",
    "\n",
    "                    except Exception as exc:\n",
    "                        print(exc)\n",
    "                    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Source: https://github.com/anthdm/ml-email-clustering/blob/master/helpers.py\n",
    "\n",
    "# row represents a document in the DTM\n",
    "def top_tfidf_feats(row, features, top_n=25):\n",
    "    topn_ids = np.argsort(row)[-top_n:][::-1]\n",
    "    top_feats = [(features[i], row[i]) for i in topn_ids]\n",
    "    df = pd.DataFrame(top_feats, columns=['features', 'score'])\n",
    "    return df\n",
    "\n",
    "def top_feats_in_doc(X, features, row_id, top_n=25):\n",
    "    row = np.squeeze(X[row_id].toarray())\n",
    "    return top_tfidf_feats(row, features, top_n)\n",
    "\n",
    "def top_mean_feats(X, features, grp_ids=None, min_tfidf=0.1, top_n=25):\n",
    "    if grp_ids:\n",
    "        D = X[grp_ids].toarray()\n",
    "    else:\n",
    "        D = X.toarray()\n",
    "\n",
    "    D[D < min_tfidf] = 0\n",
    "    tfidf_means = np.mean(D, axis=0)\n",
    "    return top_tfidf_feats(tfidf_means, features, top_n)\n",
    "\n",
    "def top_feats_per_cluster(X, y, features, min_tfidf=0.1, top_n=25):\n",
    "    dfs = []\n",
    "\n",
    "    labels = np.unique(y)\n",
    "    for label in labels:\n",
    "        ids = np.where(y==label) \n",
    "        feats_df = top_mean_feats(X, features, ids, min_tfidf=min_tfidf, top_n=top_n)\n",
    "        feats_df.label = label\n",
    "        dfs.append(feats_df)\n",
    "    return dfs\n",
    "\n",
    "\n",
    "def plot_tfidf_classfeats_h(dfs):\n",
    "    fig = pyplot.figure(figsize=(12, 9), facecolor=\"w\")\n",
    "    x = np.arange(len(dfs[0]))\n",
    "    for i, df in enumerate(dfs):\n",
    "        ax = fig.add_subplot(1, len(dfs), i+1)\n",
    "        ax.spines[\"top\"].set_visible(False)\n",
    "        ax.spines[\"right\"].set_visible(False)\n",
    "        ax.set_frame_on(False)\n",
    "        ax.get_xaxis().tick_bottom()\n",
    "        ax.get_yaxis().tick_left()\n",
    "        ax.set_xlabel(\"Tf-Idf Score\", labelpad=16, fontsize=14)\n",
    "        ax.set_title(\"cluster = \" + str(df.label), fontsize=16)\n",
    "        ax.ticklabel_format(axis='x', style='sci', scilimits=(-2,2))\n",
    "        ax.barh(x, df.score, align='center', color='#7530FF')\n",
    "        ax.set_yticks(x)\n",
    "        ax.set_ylim([-1, x[-1]+1])\n",
    "        yticks = ax.set_yticklabels(df.features)\n",
    "        pyplot.subplots_adjust(bottom=0.09, right=0.97, left=0.15, top=0.95, wspace=0.52)\n",
    "    pyplot.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Study"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Build and visualize the document-term matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Decompress DB file\n",
    "if DB_FILE.suffix == '.gz':\n",
    "    out_path = DB_FILE.parent / DB_FILE.stem\n",
    "\n",
    "    with gzip.open(str(DB_FILE), 'rb') as f_in, open(out_path, 'wb') as f_out:\n",
    "        shutil.copyfileobj(f_in, f_out)\n",
    "\n",
    "    unzipped_db_file = str(out_path)\n",
    "else:\n",
    "    unzipped_db_file = str(DB_FILE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract messages from DB file\n",
    "messages = read_messages_from_db(unzipped_db_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load messages into dataframe\n",
    "df = pd.DataFrame(messages)\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Quick visual check on the message bodies\n",
    "for body in df.body[:100]:\n",
    "    print('=' * 80)\n",
    "    print(body)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here max_df is 50% of total documents, min_df is 2 documents\n",
    "vect = TfidfVectorizer(stop_words='english', max_df=0.50, min_df=2)\n",
    "\n",
    "# Get document-term matrix\n",
    "X = vect.fit_transform(df.body)\n",
    "\n",
    "# Number of unique terms\n",
    "# print(len(vect.get_feature_names()))\n",
    "\n",
    "# Visualize document-term matrix\n",
    "X_dense = X.toarray()\n",
    "coords = PCA(n_components=2).fit_transform(X_dense)\n",
    "pyplot.scatter(coords[:, 0], coords[:, 1], c='m')\n",
    "pyplot.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Print the top terms by tf-idf value, aka the most \"original\" terms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pick an arbitrary message\n",
    "msg_index = 50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Print top terms for that message\n",
    "features = vect.get_feature_names_out()\n",
    "top_feats_in_doc(X, features, msg_index, 25)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Look at that message's body\n",
    "print(df.body[msg_index])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Print the top terms for our entire message collection\n",
    "top_mean_feats(X, features, None, 0.1, 25)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Clustering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use KMeans with 3 clusters\n",
    "n_clusters = 3\n",
    "clf = KMeans(n_clusters=n_clusters, \n",
    "            max_iter=100, \n",
    "            init='k-means++', \n",
    "            n_init=1)\n",
    "labels = clf.fit_predict(X)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's plot this with matplotlib to visualize it.\n",
    "# First we need to make 2D coordinates from the sparse matrix.\n",
    "X_dense = X.toarray()\n",
    "pca = PCA(n_components=2).fit(X_dense)\n",
    "coords = pca.transform(X_dense)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Lets plot it again, but this time we add some color to it.\n",
    "# This array needs to be at least the length of the n_clusters.\n",
    "label_colors = [\"#2AB0E9\", \"#2BAF74\", \"#D7665E\", \"#CCCCCC\", \n",
    "                \"#D2CA0D\", \"#522A64\", \"#A3DB05\", \"#FC6514\"]\n",
    "colors = [label_colors[i] for i in labels]\n",
    "\n",
    "pyplot.scatter(coords[:, 0], coords[:, 1], c=colors)\n",
    "\n",
    "# Plot the cluster centers\n",
    "centroids = clf.cluster_centers_\n",
    "centroid_coords = pca.transform(centroids)\n",
    "pyplot.scatter(centroid_coords[:, 0], centroid_coords[:, 1], marker='X', s=200, linewidths=2, c='#444d60')\n",
    "\n",
    "\n",
    "pyplot.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# View top terms per cluster\n",
    "plot_tfidf_classfeats_h(top_feats_per_cluster(X, labels, features, 0.1, 25))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": 4
}