bert_semantic_search_example.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/nlongcn/examples/blob/master/bert_semantic_search_example.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install sentence-transformers==2.2.2"
      ],
      "metadata": {
        "id": "2Nd9BxKRXn78",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "60c87580-5705-4d98-ad62-5c2c71408ed0"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: sentence-transformers==2.2.2 in /usr/local/lib/python3.10/dist-packages (2.2.2)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.10/dist-packages (from sentence-transformers==2.2.2) (1.10.1)\n",
            "Requirement already satisfied: torch>=1.6.0 in /usr/local/lib/python3.10/dist-packages (from sentence-transformers==2.2.2) (2.0.0+cu118)\n",
            "Requirement already satisfied: torchvision in /usr/local/lib/python3.10/dist-packages (from sentence-transformers==2.2.2) (0.15.1+cu118)\n",
            "Requirement already satisfied: nltk in /usr/local/lib/python3.10/dist-packages (from sentence-transformers==2.2.2) (3.8.1)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.10/dist-packages (from sentence-transformers==2.2.2) (4.65.0)\n",
            "Requirement already satisfied: transformers<5.0.0,>=4.6.0 in /usr/local/lib/python3.10/dist-packages (from sentence-transformers==2.2.2) (4.28.1)\n",
            "Requirement already satisfied: huggingface-hub>=0.4.0 in /usr/local/lib/python3.10/dist-packages (from sentence-transformers==2.2.2) (0.14.1)\n",
            "Requirement already satisfied: scikit-learn in /usr/local/lib/python3.10/dist-packages (from sentence-transformers==2.2.2) (1.2.2)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from sentence-transformers==2.2.2) (1.22.4)\n",
            "Requirement already satisfied: sentencepiece in /usr/local/lib/python3.10/dist-packages (from sentence-transformers==2.2.2) (0.1.99)\n",
            "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.4.0->sentence-transformers==2.2.2) (4.5.0)\n",
            "Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.4.0->sentence-transformers==2.2.2) (2023.4.0)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.4.0->sentence-transformers==2.2.2) (2.27.1)\n",
            "Requirement already satisfied: packaging>=20.9 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.4.0->sentence-transformers==2.2.2) (23.1)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.4.0->sentence-transformers==2.2.2) (3.12.0)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.4.0->sentence-transformers==2.2.2) (6.0)\n",
            "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch>=1.6.0->sentence-transformers==2.2.2) (1.11.1)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch>=1.6.0->sentence-transformers==2.2.2) (3.1)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch>=1.6.0->sentence-transformers==2.2.2) (3.1.2)\n",
            "Requirement already satisfied: triton==2.0.0 in /usr/local/lib/python3.10/dist-packages (from torch>=1.6.0->sentence-transformers==2.2.2) (2.0.0)\n",
            "Requirement already satisfied: lit in /usr/local/lib/python3.10/dist-packages (from triton==2.0.0->torch>=1.6.0->sentence-transformers==2.2.2) (16.0.2)\n",
            "Requirement already satisfied: cmake in /usr/local/lib/python3.10/dist-packages (from triton==2.0.0->torch>=1.6.0->sentence-transformers==2.2.2) (3.25.2)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.10/dist-packages (from transformers<5.0.0,>=4.6.0->sentence-transformers==2.2.2) (2022.10.31)\n",
            "Requirement already satisfied: tokenizers!=0.11.3,<0.14,>=0.11.1 in /usr/local/lib/python3.10/dist-packages (from transformers<5.0.0,>=4.6.0->sentence-transformers==2.2.2) (0.13.3)\n",
            "Requirement already satisfied: click in /usr/local/lib/python3.10/dist-packages (from nltk->sentence-transformers==2.2.2) (8.1.3)\n",
            "Requirement already satisfied: joblib in /usr/local/lib/python3.10/dist-packages (from nltk->sentence-transformers==2.2.2) (1.2.0)\n",
            "Requirement already satisfied: threadpoolctl>=2.0.0 in /usr/local/lib/python3.10/dist-packages (from scikit-learn->sentence-transformers==2.2.2) (3.1.0)\n",
            "Requirement already satisfied: pillow!=8.3.*,>=5.3.0 in /usr/local/lib/python3.10/dist-packages (from torchvision->sentence-transformers==2.2.2) (8.4.0)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch>=1.6.0->sentence-transformers==2.2.2) (2.1.2)\n",
            "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->huggingface-hub>=0.4.0->sentence-transformers==2.2.2) (1.26.15)\n",
            "Requirement already satisfied: charset-normalizer~=2.0.0 in /usr/local/lib/python3.10/dist-packages (from requests->huggingface-hub>=0.4.0->sentence-transformers==2.2.2) (2.0.12)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->huggingface-hub>=0.4.0->sentence-transformers==2.2.2) (3.4)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->huggingface-hub>=0.4.0->sentence-transformers==2.2.2) (2022.12.7)\n",
            "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.10/dist-packages (from sympy->torch>=1.6.0->sentence-transformers==2.2.2) (1.3.0)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Connect to Gmail API"
      ],
      "metadata": {
        "id": "ts31ntiaY31I"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# already installed by default i believe\n",
        "! pip install --upgrade google-api-python-client==2.83.0 google-auth==2.16.3 google-auth-httplib2==0.1.0 google-auth-oauthlib==0.4.1"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HpgjbtdXZT9Q",
        "outputId": "a3185a42-7855-4c88-d56c-eacfb1df3b02"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: google-api-python-client==2.83.0 in /usr/local/lib/python3.10/dist-packages (2.83.0)\n",
            "Requirement already satisfied: google-auth==2.16.3 in /usr/local/lib/python3.10/dist-packages (2.16.3)\n",
            "Requirement already satisfied: google-auth-httplib2==0.1.0 in /usr/local/lib/python3.10/dist-packages (0.1.0)\n",
            "Requirement already satisfied: google-auth-oauthlib==0.4.1 in /usr/local/lib/python3.10/dist-packages (0.4.1)\n",
            "Requirement already satisfied: google-api-core!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.0,<3.0.0dev,>=1.31.5 in /usr/local/lib/python3.10/dist-packages (from google-api-python-client==2.83.0) (2.11.0)\n",
            "Requirement already satisfied: httplib2<1dev,>=0.15.0 in /usr/local/lib/python3.10/dist-packages (from google-api-python-client==2.83.0) (0.21.0)\n",
            "Requirement already satisfied: uritemplate<5,>=3.0.1 in /usr/local/lib/python3.10/dist-packages (from google-api-python-client==2.83.0) (4.1.1)\n",
            "Requirement already satisfied: six>=1.9.0 in /usr/local/lib/python3.10/dist-packages (from google-auth==2.16.3) (1.16.0)\n",
            "Requirement already satisfied: cachetools<6.0,>=2.0.0 in /usr/local/lib/python3.10/dist-packages (from google-auth==2.16.3) (5.3.0)\n",
            "Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.10/dist-packages (from google-auth==2.16.3) (0.3.0)\n",
            "Requirement already satisfied: rsa<5,>=3.1.4 in /usr/local/lib/python3.10/dist-packages (from google-auth==2.16.3) (4.9)\n",
            "Requirement already satisfied: requests-oauthlib>=0.7.0 in /usr/local/lib/python3.10/dist-packages (from google-auth-oauthlib==0.4.1) (1.3.1)\n",
            "Requirement already satisfied: googleapis-common-protos<2.0dev,>=1.56.2 in /usr/local/lib/python3.10/dist-packages (from google-api-core!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.0,<3.0.0dev,>=1.31.5->google-api-python-client==2.83.0) (1.59.0)\n",
            "Requirement already satisfied: protobuf!=3.20.0,!=3.20.1,!=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5,<5.0.0dev,>=3.19.5 in /usr/local/lib/python3.10/dist-packages (from google-api-core!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.0,<3.0.0dev,>=1.31.5->google-api-python-client==2.83.0) (3.20.3)\n",
            "Requirement already satisfied: requests<3.0.0dev,>=2.18.0 in /usr/local/lib/python3.10/dist-packages (from google-api-core!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.0,<3.0.0dev,>=1.31.5->google-api-python-client==2.83.0) (2.27.1)\n",
            "Requirement already satisfied: pyparsing!=3.0.0,!=3.0.1,!=3.0.2,!=3.0.3,<4,>=2.4.2 in /usr/local/lib/python3.10/dist-packages (from httplib2<1dev,>=0.15.0->google-api-python-client==2.83.0) (3.0.9)\n",
            "Requirement already satisfied: pyasn1<0.6.0,>=0.4.6 in /usr/local/lib/python3.10/dist-packages (from pyasn1-modules>=0.2.1->google-auth==2.16.3) (0.5.0)\n",
            "Requirement already satisfied: oauthlib>=3.0.0 in /usr/local/lib/python3.10/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib==0.4.1) (3.2.2)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests<3.0.0dev,>=2.18.0->google-api-core!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.0,<3.0.0dev,>=1.31.5->google-api-python-client==2.83.0) (2022.12.7)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests<3.0.0dev,>=2.18.0->google-api-core!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.0,<3.0.0dev,>=1.31.5->google-api-python-client==2.83.0) (3.4)\n",
            "Requirement already satisfied: charset-normalizer~=2.0.0 in /usr/local/lib/python3.10/dist-packages (from requests<3.0.0dev,>=2.18.0->google-api-core!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.0,<3.0.0dev,>=1.31.5->google-api-python-client==2.83.0) (2.0.12)\n",
            "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests<3.0.0dev,>=2.18.0->google-api-core!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.0,<3.0.0dev,>=1.31.5->google-api-python-client==2.83.0) (1.26.15)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "from google.oauth2.credentials import Credentials\n",
        "from google_auth_oauthlib.flow import InstalledAppFlow\n",
        "from google.auth.transport.requests import Request\n",
        "from googleapiclient.errors import HttpError\n",
        "import google.auth.exceptions\n",
        "from googleapiclient.discovery import build\n",
        "\n",
        "# If modifying these scopes, delete the file token.json.\n",
        "SCOPES = ['https://www.googleapis.com/auth/gmail.readonly']\n",
        "\n",
        "creds = None\n",
        "if os.path.exists('token.json'):\n",
        "    creds = Credentials.from_authorized_user_file('token.json')\n",
        "if not creds or not creds.valid:\n",
        "    if creds and creds.expired and creds.refresh_token:\n",
        "        creds.refresh(Request())\n",
        "    else:\n",
        "        flow = InstalledAppFlow.from_client_secrets_file('credentials.json', ['https://www.googleapis.com/auth/gmail.readonly'])\n",
        "        creds = flow.run_local_server(port=0)\n",
        "    with open('token.json', 'w') as token:\n",
        "        token.write(creds.to_json())\n",
        "return creds"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_HpiVDsxY2in",
        "outputId": "e2fcea6d-4d0f-47da-fc93-e6a11c8a4427"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Please visit this URL to authorize this application: https://accounts.google.com/o/oauth2/auth?response_type=code&client_id=661640239003-5ncgrgactghtf48on2jnm9br9oqdug56.apps.googleusercontent.com&redirect_uri=http%3A%2F%2Flocalhost%3A34021%2F&scope=https%3A%2F%2Fwww.googleapis.com%2Fauth%2Fgmail.readonly&state=1OZJVtH73gwxH80uO30YydPAeez6I9&access_type=offline\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "emails = [] # indvidiual raw emails stored\n",
        "CUT_OFF = 10000 # maximum number of emails to be downloaded\n",
        "try:\n",
        "  service = build(\"gmail\", \"v1\", credentials=creds)\n",
        "  gmail_messages = service.users().messages()\n",
        "  \n",
        "  has_next_token = True\n",
        "  next_token = None\n",
        "  count = 0\n",
        "  while(has_next_token):\n",
        "    results = gmail_messages.list(userId='me', pageToken=next_token).execute()\n",
        "    messages = results[\"messages\"]\n",
        "    if \"nextPageToken\" in results:\n",
        "      next_token = results[\"nextPageToken\"]\n",
        "    else:\n",
        "      next_token = None\n",
        "      has_next_token = False\n",
        "    size_estimate = results[\"resultSizeEstimate\"]\n",
        "    print(f\"next_token {next_token}, size estimate: {size_estimate}\")\n",
        "    for msg in messages:\n",
        "      msg_dict = gmail_messages.get(userId='me', id=msg['id'], format='raw').execute()\n",
        "      emails.append(msg_dict)\n",
        "    count += len(emails)\n",
        "    if len(emails) > CUT_OFF:\n",
        "      has_next_token = False\n",
        "      next_token = None\n",
        "except HttpError as error:\n",
        "  print(f'An error occurred: {error}')"
      ],
      "metadata": {
        "id": "WTzjpEcNflzd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [],
      "metadata": {
        "id": "gCaqZXIjlK2c"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "! pip install bs4 "
      ],
      "metadata": {
        "id": "-L91aRUpkYyk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from bs4 import BeautifulSoup\n",
        "import base64\n",
        "\n",
        "# Function to remove tags\n",
        "def remove_tags(html):\n",
        " \n",
        "    # parse html content\n",
        "    soup = BeautifulSoup(html, \"html.parser\")\n",
        " \n",
        "    for data in soup(['style', 'script']):\n",
        "        # Remove tags\n",
        "        data.decompose()\n",
        " \n",
        "    # return data by retrieving the tag content\n",
        "    return ' '.join(soup.stripped_strings)"
      ],
      "metadata": {
        "id": "flszF5hukhES"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def extract_header(headers, header_name:str) -> str:\n",
        "  for h in headers:\n",
        "    if h[\"name\"] == header_name:\n",
        "      return h[\"value\"]"
      ],
      "metadata": {
        "id": "xs8G8jV9r4Or"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "len(emails)"
      ],
      "metadata": {
        "id": "d7vBWsqjNGoS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import email\n",
        "from email import policy\n",
        "from email.parser import Parser\n",
        "from io import BytesIO\n",
        "import io\n",
        "\n",
        "parser = Parser(policy=policy.default)\n",
        "\n",
        "# reference content dictionary\n",
        "content_dictionary = {\n",
        "}\n",
        "\n",
        "counter = 0\n",
        "for email in emails:\n",
        "  msg_str = base64.urlsafe_b64decode(email['raw'].encode('ASCII'))\n",
        "  \n",
        "  subject = from_email = to_email = None\n",
        "  msg = parser.parsestr(msg_str.decode('utf-8'))\n",
        "  for key in msg.keys():\n",
        "    if key == \"Subject\":\n",
        "      subject = msg.get_all(\"Subject\")\n",
        "    if key == \"From\":\n",
        "      from_email = msg.get_all(\"From\")\n",
        "    if key == \"To\":\n",
        "      to_email = msg.get_all(\"To\")\n",
        "  print(f\"from: {from_email}, subject: {subject}\")\n",
        "  for part in msg.walk():\n",
        "     if part.get_content_type() == \"text/html\":\n",
        "       content = part.get_content()\n",
        "       clear_content = remove_tags(content)\n",
        "       content_dictionary[counter] = {\n",
        "           \"id\": email[\"id\"],\n",
        "           \"subject\": subject,\n",
        "           \"from\": from_email,\n",
        "           \"to\": to_email,\n",
        "           \"body\": clear_content,\n",
        "           \"snippet\": email[\"snippet\"]\n",
        "       }\n",
        "       counter += 1"
      ],
      "metadata": {
        "id": "vJx6twL6kpFW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from sentence_transformers import SentenceTransformer, util\n",
        "import torch\n",
        "\n",
        "embedder = SentenceTransformer('all-MiniLM-L6-v2')"
      ],
      "metadata": {
        "id": "bohvFPIOB6dE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "len(content_dictionary)"
      ],
      "metadata": {
        "id": "gKhNLzgZODjW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "corpus = [c[\"body\"] for c in content_dictionary.values()]\n",
        "print(f\"embedding {len(corpus)} sentences\")\n",
        "corpus_embeddings = embedder.encode(corpus, convert_to_tensor=True)"
      ],
      "metadata": {
        "id": "N01ovkL8X2r3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Query sentences:\n",
        "queries = ['where to hire a wall fixing guy in Salt Lake City?']"
      ],
      "metadata": {
        "id": "51tXNIIaX8rW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Find the closest 5 sentences of the corpus for each query sentence based on cosine similarity\n",
        "top_k = min(1, len(corpus))\n",
        "for query in queries:\n",
        "    query_embedding = embedder.encode(query, convert_to_tensor=True)\n",
        "\n",
        "    # We use cosine-similarity and torch.topk to find the highest 5 scores\n",
        "    # cos_scores = util.cos_sim(query_embedding, corpus_embeddings)[0]\n",
        "    # top_results = torch.topk(cos_scores, k=top_k)\n",
        "\n",
        "    # print(\"\\n\\n======================\\n\\n\")\n",
        "    # print(\"Query:\", query)\n",
        "    # print(\"\\nTop 5 most similar sentences in corpus:\")\n",
        "\n",
        "    # for score, idx in zip(top_results[0], top_results[1]):\n",
        "    #     print(corpus[idx], \"(Score: {:.4f})\".format(score))\n",
        "    # Alternatively, we can also use util.semantic_search to perform cosine similarty + topk\n",
        "    hits = util.semantic_search(query_embedding, corpus_embeddings, top_k=2)\n",
        "    hits = hits[0]      #Get the hits for the first query\n",
        "    for hit in hits:\n",
        "        id = hit['corpus_id']\n",
        "        item = content_dictionary[id]\n",
        "        print(f\"query: {query}, subject: {item['subject']}, {item['snippet']} https://mail.google.com/mail/u/2/#all/{item['id']}\")\n",
        "        # print(corpus[hit['corpus_id']], \"(Score: {:.4f})\".format(hit['score']))\n",
        "\n",
        "    \"\"\"\n",
        "    # Alternatively, we can also use util.semantic_search to perform cosine similarty + topk\n",
        "    hits = util.semantic_search(query_embedding, corpus_embeddings, top_k=5)\n",
        "    hits = hits[0]      #Get the hits for the first query\n",
        "    for hit in hits:\n",
        "        print(corpus[hit['corpus_id']], \"(Score: {:.4f})\".format(hit['score']))\n",
        "    \"\"\""
      ],
      "metadata": {
        "id": "2OyFFCRKX9q-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "Gx2OS2S1M7wN"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}