Hybrid_Search_Workshop.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "collapsed_sections": [
        "A514i6Fbqx28"
      ],
      "authorship_tag": "ABX9TyMl8wuoJ61Cth3nzl0u5Qka",
      "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/trancethehuman/ai-workshop-code/blob/main/Hybrid_Search_Workshop.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Introduction"
      ],
      "metadata": {
        "id": "lrM-ldk7rQDV"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Setup documents for knowledge base"
      ],
      "metadata": {
        "id": "j__l8tgHh7gF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!wget https://raw.githubusercontent.com/trancethehuman/ai-workshop-code/main/datasets/legal_text_classification_first_1000.csv -q"
      ],
      "metadata": {
        "id": "YQroPABw7OaS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "pip install tabulate -q"
      ],
      "metadata": {
        "id": "eUWUcuRDPz0E"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import pandas as pd\n",
        "from tabulate import tabulate\n",
        "import textwrap\n",
        "\n",
        "# Read the CSV file\n",
        "df = pd.read_csv('legal_text_classification_first_1000.csv')\n",
        "\n",
        "# Function to wrap text\n",
        "def wrap_text(text, width):\n",
        "    return \"\\n\".join(textwrap.wrap(text, width))\n",
        "\n",
        "# Set the width for wrapping\n",
        "wrap_width = 60\n",
        "\n",
        "# Create a copy of the DataFrame for display so we don't modify the original data\n",
        "df_display = df.head(3).copy()\n",
        "\n",
        "# Apply the wrap_text function to each column in the copied DataFrame\n",
        "df_display = df_display.applymap(lambda x: wrap_text(str(x), wrap_width) if isinstance(x, str) else x)\n",
        "\n",
        "# Display the modified DataFrame using tabulate\n",
        "print(tabulate(df_display, headers='keys', tablefmt='pretty'))"
      ],
      "metadata": {
        "id": "a453vmErPEe2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "row_count = df.shape[0]\n",
        "\n",
        "print(f\"The DataFrame contains {row_count} rows.\")"
      ],
      "metadata": {
        "id": "BmPFfIz4WOxQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "pip install openai -q"
      ],
      "metadata": {
        "id": "gERW3o7x9zHW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import getpass\n",
        "\n",
        "OPENAI_API_KEY = \"\""
      ],
      "metadata": {
        "id": "I39KvgUw95Cy",
        "collapsed": true
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "pip install langchain-openai -q"
      ],
      "metadata": {
        "id": "l7SG-4DJ_MSv",
        "collapsed": true
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from langchain_openai import OpenAIEmbeddings\n",
        "\n",
        "EMBEDDINGS_DIMENSIONS = 512\n",
        "\n",
        "embedding_client = OpenAIEmbeddings(api_key=OPENAI_API_KEY,\n",
        "    model=\"text-embedding-3-large\", dimensions=EMBEDDINGS_DIMENSIONS)"
      ],
      "metadata": {
        "id": "y180Q-FL7Xmk",
        "collapsed": true
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "pip install tiktoken -q"
      ],
      "metadata": {
        "id": "VakX9GCZz4RH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import tiktoken\n",
        "\n",
        "def num_tokens_from_string(string: str, encoding_name: str) -> int:\n",
        "    \"\"\"Returns the number of tokens in a text string.\"\"\"\n",
        "    encoding = tiktoken.get_encoding(encoding_name)\n",
        "    num_tokens = len(encoding.encode(string))\n",
        "    return num_tokens"
      ],
      "metadata": {
        "id": "5nMr7fZwyHsX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import pandas as pd\n",
        "import time\n",
        "from typing import List, Dict\n",
        "\n",
        "def get_embeddings(\n",
        "    df: pd.DataFrame,\n",
        "    num_rows: int = 2,\n",
        "    max_tokens: int = 8191,\n",
        "    encoding_name: str = \"cl100k_base\",\n",
        "    price_per_token: float = 0.13 / 1000000\n",
        ") -> List[Dict]:\n",
        "    \"\"\"\n",
        "    Process a DataFrame, adding embeddings to a specified number of rows,\n",
        "    ensuring the total number of tokens per request does not exceed max_tokens.\n",
        "    \"\"\"\n",
        "\n",
        "    result_list = []\n",
        "    batch_case_text = []\n",
        "    batch_ids = []\n",
        "    current_tokens = 0\n",
        "    total_tokens = 0\n",
        "    total_batches = 0\n",
        "\n",
        "    # Start the timer to see how long it'd take to get all the embeddings\n",
        "    start_time = time.time()\n",
        "\n",
        "    for index, row in df.head(num_rows).iterrows():\n",
        "        case_text = row['case_text']\n",
        "\n",
        "        if not isinstance(case_text, str) or not case_text.strip():\n",
        "            continue\n",
        "\n",
        "        tokens = num_tokens_from_string(case_text, encoding_name)\n",
        "\n",
        "        # Check if adding this text would exceed the token limit\n",
        "        if current_tokens + tokens > max_tokens:\n",
        "            if batch_case_text:\n",
        "                # Process the current batch\n",
        "                embeddings = embedding_client.embed_documents(batch_case_text)\n",
        "                for i, text in enumerate(batch_case_text):\n",
        "                    result_list.append({\n",
        "                        'case_id': batch_ids[i],\n",
        "                        'case_title': df.loc[df['case_id'] == batch_ids[i], 'case_title'].values[0],\n",
        "                        'case_text': text,\n",
        "                        'embeddings': embeddings[i]\n",
        "                    })\n",
        "\n",
        "            # Reset batch and start a new one with the current text\n",
        "            batch_case_text = [case_text]\n",
        "            batch_ids = [row['case_id']]\n",
        "            current_tokens = tokens\n",
        "            total_batches += 1  # Increment batch count\n",
        "        else:\n",
        "            # Add the text to the current batch\n",
        "            batch_case_text.append(case_text)\n",
        "            batch_ids.append(row['case_id'])\n",
        "            current_tokens += tokens\n",
        "\n",
        "        total_tokens += tokens  # Increment total token count\n",
        "\n",
        "    # Process the final batch\n",
        "    if batch_case_text:\n",
        "        embeddings = embedding_client.embed_documents(batch_case_text)\n",
        "        for i, text in enumerate(batch_case_text):\n",
        "            result_list.append({\n",
        "                'case_id': batch_ids[i],\n",
        "                'case_title': df.loc[df['case_id'] == batch_ids[i], 'case_title'].values[0],\n",
        "                'case_text': text,\n",
        "                'embeddings': embeddings[i]\n",
        "            })\n",
        "        total_batches += 1  # Increment batch count for the final batch\n",
        "\n",
        "    end_time = time.time()\n",
        "    duration = end_time - start_time\n",
        "\n",
        "    # Print the statistics\n",
        "    print(f\"Completed in {duration:.2f} seconds.\")\n",
        "    print(f\"Total number of tokens: {total_tokens:,}\")\n",
        "    print(f\"Total number of batches: {total_batches:,}\")\n",
        "    print(f\"Money burned: ${total_tokens * price_per_token:.6f} - Thanks, Invest Ottawa ❤️\")\n",
        "\n",
        "    return result_list"
      ],
      "metadata": {
        "id": "oOqyXq5SO_Hs"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "test_cases_with_embeddings = get_embeddings(df, num_rows=3)\n",
        "\n",
        "# For displaying our test\n",
        "MAX_TABLE_WIDTH = 60\n",
        "MAX_CHARACTERS_LENGTH = 150\n",
        "\n",
        "# Utility function to wrap and truncate text\n",
        "def wrap_and_truncate_text(text, width=MAX_TABLE_WIDTH, max_length=MAX_CHARACTERS_LENGTH):\n",
        "    if isinstance(text, list):\n",
        "        text = str(text)\n",
        "    if len(text) > max_length:\n",
        "        text = text[:max_length] + '...'\n",
        "    return \"\\n\".join(textwrap.wrap(text, width))\n",
        "\n",
        "# Apply wrapping and truncation to all items in the test data before we display them so we get nice table\n",
        "wrapped_test_cases_with_embeddings = [\n",
        "    {key: wrap_and_truncate_text(value) for key, value in item.items()}\n",
        "    for item in test_cases_with_embeddings\n",
        "]\n",
        "\n",
        "# Define the alignment for each column (quite funny)\n",
        "colalign = (\"left\", \"left\", \"left\", \"left\")\n",
        "\n",
        "print(tabulate(wrapped_test_cases_with_embeddings, headers='keys', tablefmt='pretty', colalign=colalign))"
      ],
      "metadata": {
        "id": "n-JOTMv_UFa4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "all_cases_with_embeddings = get_embeddings(df, num_rows=row_count)"
      ],
      "metadata": {
        "id": "rwV-gUroVg9v"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Setup vector databases' clients"
      ],
      "metadata": {
        "id": "cRbM-44gpa8F"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Pinecone"
      ],
      "metadata": {
        "id": "JFXriKeVqKji"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "pip install pinecone-client pinecone-notebooks pinecone-text -q"
      ],
      "metadata": {
        "id": "03Up66Nq-FdD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from pinecone_notebooks.colab import Authenticate\n",
        "\n",
        "Authenticate()"
      ],
      "metadata": {
        "id": "si2S_h_H-JxX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "from pinecone import Pinecone, ServerlessSpec\n",
        "\n",
        "pc = Pinecone(api_key=os.environ.get('PINECONE_API_KEY'))"
      ],
      "metadata": {
        "id": "l9Z3ivBx-d1F"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "index_name = \"hybridhearchexperiment\"\n",
        "\n",
        "# Just checking to see if this index already exist\n",
        "existing_indexes = pc.list_indexes().names()\n",
        "\n",
        "# If index doesn't exist yet, then delete it and create one (we're starting from scratch)\n",
        "if index_name in existing_indexes:\n",
        "  pc.delete_index(index_name)\n",
        "  print(\"Deleted old index.\")\n",
        "\n",
        "pc.create_index(\n",
        "    name=index_name,\n",
        "    dimension=EMBEDDINGS_DIMENSIONS,\n",
        "    metric=\"dotproduct\", # to use sparse-dense index (aka hybrid search) in Pinecone, they require us to use dotproduct.\n",
        "    spec=ServerlessSpec(\n",
        "        cloud='aws',\n",
        "        region='us-east-1'\n",
        "    )\n",
        ")\n",
        "# wait for index to be initialized\n",
        "while not pc.describe_index(index_name).status['ready']:\n",
        "    time.sleep(1)\n",
        "\n",
        "# connect to index\n",
        "pinecone_index = pc.Index(index_name)\n",
        "# view index stats\n",
        "pinecone_index.describe_index_stats()"
      ],
      "metadata": {
        "id": "ZqaE9zzs-qQT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from pinecone_text.sparse import BM25Encoder\n",
        "\n",
        "bm25 = BM25Encoder()"
      ],
      "metadata": {
        "id": "SGrlhwS9KK-u"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!wget https://raw.githubusercontent.com/trancethehuman/ai-workshop-code/main/datasets/Legal_Text_Classification_Data_500_train.csv -q"
      ],
      "metadata": {
        "id": "Jn8uf_vwL_ec"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Load our training set into memory as a DataFrame\n",
        "train_df = pd.read_csv('Legal_Text_Classification_Data_500_train.csv')\n",
        "\n",
        "# Train baby\n",
        "bm25.fit(train_df.get(\"case_text\").astype(str).tolist())"
      ],
      "metadata": {
        "id": "yZjkyrE0MS8y"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from tqdm import tqdm\n",
        "\n",
        "def group_embeddings_and_generate_sparse_vectors(cases, sparse_vector_model):\n",
        "    all_cases_embeddings_and_sparse_vectors = []\n",
        "\n",
        "    for case in tqdm(cases, desc=\"Processing cases\"):\n",
        "        case_id = case['case_id']\n",
        "        case_title = case['case_title']\n",
        "        case_text = case['case_text']\n",
        "        embeddings = case['embeddings']\n",
        "\n",
        "        # Encode the case text using a sparse vector model\n",
        "        sparse_values = sparse_vector_model.encode_documents(case_text)\n",
        "\n",
        "        # Create the new dictionary with the required structure\n",
        "        new_case_dict = {\n",
        "            'id': case_id,\n",
        "            'sparse_values': sparse_values,\n",
        "            'values': embeddings,\n",
        "            'metadata': {\n",
        "                'case_title': case_title,\n",
        "                'case_text': case_text\n",
        "            }\n",
        "        }\n",
        "\n",
        "        # Add the new dictionary to the list\n",
        "        all_cases_embeddings_and_sparse_vectors.append(new_case_dict)\n",
        "\n",
        "    return all_cases_embeddings_and_sparse_vectors\n"
      ],
      "metadata": {
        "id": "oljzXCrXNwwo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "all_cases_embeddings_and_sparse_vectors = group_embeddings_and_generate_sparse_vectors(all_cases_with_embeddings, bm25)"
      ],
      "metadata": {
        "id": "CkBzD8CsTO30"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Prepare the data for tabulate\n",
        "table_data = []\n",
        "for case in all_cases_embeddings_and_sparse_vectors:\n",
        "    table_data.append([\n",
        "        case['id'],\n",
        "        case['metadata']['case_title'][:30],\n",
        "        case['metadata']['case_text'][:30],\n",
        "        str(case['sparse_values'])[:30],\n",
        "        str(case['values'])[:30],\n",
        "    ])\n",
        "\n",
        "# Limit the table_data to the first 10 rows\n",
        "table_data = table_data[:10]\n",
        "\n",
        "# Define the headers based on the keys\n",
        "headers = [\"ID\", \"Case Title\", \"Case Text\", \"Sparse Values\", \"Embeddings\"]\n",
        "\n",
        "# Print the table\n",
        "print(tabulate(table_data, headers=headers, tablefmt=\"grid\"))"
      ],
      "metadata": {
        "id": "W8t8MZQlUYhU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Upsert our hard work\n",
        "pinecone_index.upsert(vectors=all_cases_embeddings_and_sparse_vectors, batch_size=100) # batch size is very important here because our request is a bit large.\n",
        "\n",
        "# See if our index's stats changed\n",
        "print(pinecone_index.describe_index_stats())"
      ],
      "metadata": {
        "id": "ETgB-9TDVTjN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from typing import Any\n",
        "\n",
        "def convert_string_query_to_vectors(query: str) -> Dict[str, Any]:\n",
        "    dense_vector = embedding_client.embed_query(query)\n",
        "    sparse_vector = bm25.encode_queries(query)\n",
        "\n",
        "    return {\n",
        "        \"dense\": dense_vector,\n",
        "        \"sparse\": sparse_vector\n",
        "    }"
      ],
      "metadata": {
        "id": "P-5s05-JtdFf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def hybrid_scale(dense, sparse, alpha: float):\n",
        "    \"\"\"Hybrid vector scaling using a convex combination\n",
        "\n",
        "    alpha * dense + (1 - alpha) * sparse\n",
        "\n",
        "    Args:\n",
        "        dense: Array of floats representing\n",
        "        sparse: a dict of `indices` and `values`\n",
        "        alpha: float between 0 and 1 where 0 == sparse only\n",
        "               and 1 == dense only\n",
        "    \"\"\"\n",
        "    if alpha < 0 or alpha > 1:\n",
        "        raise ValueError(\"Alpha must be between 0 and 1\")\n",
        "    # scale sparse and dense vectors to create hybrid search vecs\n",
        "    hsparse = {\n",
        "        'indices': sparse['indices'],\n",
        "        'values':  [v * (1 - alpha) for v in sparse['values']]\n",
        "    }\n",
        "    hdense = [v * alpha for v in dense]\n",
        "    return hdense, hsparse"
      ],
      "metadata": {
        "id": "rxnZyBcbubhy"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Weaviate"
      ],
      "metadata": {
        "id": "vRs3Uriaq00Z"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "pip install weaviate-client -q"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "KP6uSNMPjDgj",
        "outputId": "26aa1965-a15a-42d6-fa0c-3e91f3405706"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[?25l     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/325.7 kB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r\u001b[2K     \u001b[91m━━━━━━━━━━━━\u001b[0m\u001b[91m╸\u001b[0m\u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m102.4/325.7 kB\u001b[0m \u001b[31m2.9 MB/s\u001b[0m eta \u001b[36m0:00:01\u001b[0m\r\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m325.7/325.7 kB\u001b[0m \u001b[31m5.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m40.0/40.0 kB\u001b[0m \u001b[31m3.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m223.8/223.8 kB\u001b[0m \u001b[31m20.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.3/2.3 MB\u001b[0m \u001b[31m37.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m309.3/309.3 kB\u001b[0m \u001b[31m26.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "cudf-cu12 24.4.1 requires protobuf<5,>=3.20, but you have protobuf 5.27.2 which is incompatible.\n",
            "google-ai-generativelanguage 0.6.4 requires 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, but you have protobuf 5.27.2 which is incompatible.\n",
            "google-api-core 2.16.2 requires 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.0.dev0,>=3.19.5, but you have protobuf 5.27.2 which is incompatible.\n",
            "google-cloud-aiplatform 1.59.0 requires 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, but you have protobuf 5.27.2 which is incompatible.\n",
            "google-cloud-bigquery-connection 1.12.1 requires 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, but you have protobuf 5.27.2 which is incompatible.\n",
            "google-cloud-bigquery-storage 2.25.0 requires 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, but you have protobuf 5.27.2 which is incompatible.\n",
            "google-cloud-bigtable 2.24.0 requires 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, but you have protobuf 5.27.2 which is incompatible.\n",
            "google-cloud-datastore 2.15.2 requires 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, but you have protobuf 5.27.2 which is incompatible.\n",
            "google-cloud-firestore 2.11.1 requires 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, but you have protobuf 5.27.2 which is incompatible.\n",
            "google-cloud-functions 1.13.3 requires 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, but you have protobuf 5.27.2 which is incompatible.\n",
            "google-cloud-translate 3.11.3 requires 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, but you have protobuf 5.27.2 which is incompatible.\n",
            "tensorflow 2.15.0 requires protobuf!=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5,<5.0.0dev,>=3.20.3, but you have protobuf 5.27.2 which is incompatible.\n",
            "tensorflow-metadata 1.15.0 requires protobuf<4.21,>=3.20.3; python_version < \"3.11\", but you have protobuf 5.27.2 which is incompatible.\u001b[0m\u001b[31m\n",
            "\u001b[0m"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Then, we setup our Weaviate client."
      ],
      "metadata": {
        "id": "CXd43-H9qra8"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "WCD_URL = \"\"\n",
        "WCD_API_KEY = \"\""
      ],
      "metadata": {
        "id": "qkKNKc_srBQT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import weaviate\n",
        "import weaviate.classes as wvc\n",
        "import os\n",
        "import requests\n",
        "import json\n",
        "\n",
        "weaviate_client = weaviate.connect_to_weaviate_cloud(\n",
        "    cluster_url=WCD_URL,\n",
        "    auth_credentials=weaviate.auth.AuthApiKey(WCD_API_KEY)\n",
        ")"
      ],
      "metadata": {
        "id": "0Zs85GcAjAPA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "weaviate_index = None\n",
        "\n",
        "if weaviate_client.collections.exists(index_name):\n",
        "  print(\"Collection (index) already exists. Deleting and making a new one..\")\n",
        "  weaviate_client.collections.delete(index_name)\n",
        "\n",
        "  print(\"Creating a new collection (index) with some metadata as properties\")\n",
        "  weaviate_client.collections.create(name=index_name, properties=[\n",
        "            wvc.config.Property(\n",
        "                name=\"case_id\",\n",
        "                data_type=wvc.config.DataType.TEXT,\n",
        "                vectorize_property_name=False,\n",
        "                ),\n",
        "            wvc.config.Property(\n",
        "                name=\"case_title\",\n",
        "                data_type=wvc.config.DataType.TEXT,\n",
        "                vectorize_property_name=False,\n",
        "            ),\n",
        "            wvc.config.Property(\n",
        "                name=\"case_text\",\n",
        "                data_type=wvc.config.DataType.TEXT,\n",
        "                vectorize_property_name=False,\n",
        "            )\n",
        "        ])\n",
        "  print(\"Weaviate collection (index) created.\")"
      ],
      "metadata": {
        "id": "0kjFNGCvsmMu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "collection_weaviate = weaviate_client.collections.get(index_name)\n",
        "\n",
        "# Get the total number of items to process\n",
        "total_items = len(all_cases_embeddings_and_sparse_vectors)\n",
        "\n",
        "with collection_weaviate.batch.dynamic() as batch:\n",
        "    # Wrap the loop with tqdm for nice progress bar lol\n",
        "    for i, d in tqdm(enumerate(all_cases_embeddings_and_sparse_vectors), total=total_items, desc=\"Uploading to Weaviate\"):\n",
        "        properties = {\n",
        "            \"case_id\": d[\"id\"],\n",
        "            \"case_title\": d[\"metadata\"][\"case_title\"],\n",
        "            \"case_text\": d[\"metadata\"][\"case_text\"],\n",
        "        }\n",
        "\n",
        "        custom_vector = d[\"values\"]\n",
        "\n",
        "        batch.add_object(\n",
        "            properties=properties,\n",
        "            vector=custom_vector\n",
        "        )\n"
      ],
      "metadata": {
        "id": "DcPUYzy66wKx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## The search (setup the search query and determine what we're looking for)"
      ],
      "metadata": {
        "id": "pRtV2qtDqh5B"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def display_pinecone_result_in_nice_table(data):\n",
        "    \"\"\"\n",
        "    Displays Pinecone result data in a nicely formatted table.\n",
        "\n",
        "    Args:\n",
        "    data (dict): The data to be displayed, in the format provided.\n",
        "\n",
        "    Returns:\n",
        "    str: A string representation of the data in a nicely formatted table.\n",
        "    \"\"\"\n",
        "    # Extract relevant data for the table\n",
        "    table_data = []\n",
        "    for match in data['matches']:\n",
        "        row = {\n",
        "            'ID': match['id'],\n",
        "            'Case Title': match['metadata']['case_title'],\n",
        "            'Relevance Score': match['score'],\n",
        "            'Case Text': match['metadata']['case_text'][:200] + '...' if len(match['metadata']['case_text']) > 200 else match['metadata']['case_text'],\n",
        "        }\n",
        "        table_data.append(row)\n",
        "\n",
        "    df = pd.DataFrame(table_data)\n",
        "\n",
        "    # Display the DataFrame in a nice table format using tabulate\n",
        "    table = tabulate(df, headers='keys', tablefmt='grid')\n",
        "\n",
        "    print(table)"
      ],
      "metadata": {
        "id": "Uexc8TLC0lad"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "search_query = \"Whats the verdict from Palmer J in Macleay Nominees Pty\""
      ],
      "metadata": {
        "id": "OUQgfXOcvVni"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "the_case_we_need_to_find = df[df['case_id'] == \"Case269\"]\n",
        "\n",
        "print(tabulate(the_case_we_need_to_find, headers='keys', tablefmt='pretty'))"
      ],
      "metadata": {
        "id": "TTYTiZNA4xaI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "search_query_as_vectors = convert_string_query_to_vectors(search_query)"
      ],
      "metadata": {
        "id": "7hySZbzQvtWz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Let's test similarity search alone"
      ],
      "metadata": {
        "id": "aeyvM8LcoD0Y"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Pinecone"
      ],
      "metadata": {
        "id": "A514i6Fbqx28"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "hdense, hsparse = hybrid_scale(search_query_as_vectors.get(\"dense\"), search_query_as_vectors.get(\"sparse\"), alpha=1)"
      ],
      "metadata": {
        "id": "VGMDFkTRxKTf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "pinecone_result = pinecone_index.query(\n",
        "    top_k=3,\n",
        "    vector=hdense,\n",
        "    sparse_vector=hsparse,\n",
        "    include_metadata=True\n",
        ")\n",
        "\n",
        "display_pinecone_result_in_nice_table(pinecone_result)"
      ],
      "metadata": {
        "collapsed": true,
        "id": "GM95OtMCxtbv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Weaviate"
      ],
      "metadata": {
        "id": "l1DwLU0SKZk7"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "We're going to pass our query as dense vectors in as search query, and only do similarity search with Weaviate."
      ],
      "metadata": {
        "id": "VzJiTr7qKYeK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from pprint import pprint\n",
        "\n",
        "response = collection_weaviate.query.near_vector(\n",
        "    near_vector=search_query_as_vectors[\"dense\"],\n",
        "    limit=2,\n",
        "    return_metadata=wvc.query.MetadataQuery(certainty=True)\n",
        ")\n",
        "\n",
        "pprint(response)"
      ],
      "metadata": {
        "id": "8M2yDP3QFk3Z"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Hybrid Search (Pinecone and Weaviate)"
      ],
      "metadata": {
        "id": "8PakPltapk4B"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Pinecone"
      ],
      "metadata": {
        "id": "5tbAP7L_4gHN"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# First, we re-define the hybrid scale to weigh in more on keyword search (sparse vectors)\n",
        "hdense, hsparse = hybrid_scale(search_query_as_vectors.get(\"dense\"), search_query_as_vectors.get(\"sparse\"), alpha=0.6)\n",
        "\n",
        "# Then, we let it go ham\n",
        "pinecone_result = pinecone_index.query(\n",
        "    top_k=3,\n",
        "    vector=hdense,\n",
        "    sparse_vector=hsparse,\n",
        "    include_metadata=True\n",
        ")\n",
        "\n",
        "display_pinecone_result_in_nice_table(pinecone_result)\n"
      ],
      "metadata": {
        "collapsed": true,
        "id": "_X8TO5sS0B3J"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Weaviate"
      ],
      "metadata": {
        "id": "QQ_VlRY2MIQF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "alpha_weaviate = 0.5"
      ],
      "metadata": {
        "id": "fJ_WiWvnOMVk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "response = collection_weaviate.query.hybrid(\n",
        "    query=search_query,\n",
        "    vector=search_query_as_vectors[\"dense\"],\n",
        "    alpha=alpha_weaviate,\n",
        "    limit=3,\n",
        ")\n",
        "\n",
        "pprint(response)"
      ],
      "metadata": {
        "id": "ri8XVoHqMNzR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Hybrid search part II: Full-text search (Postgres) + Reranker (Jina AI)"
      ],
      "metadata": {
        "id": "oeh8zTwLptc_"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Setup Postgres & pgvector with Supabase and vecs"
      ],
      "metadata": {
        "id": "PxVpvkOaqU41"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "pip install vecs supabase -q"
      ],
      "metadata": {
        "id": "uXW3xTrvQEWd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from vecs import IndexArgsHNSW, IndexMeasure, create_client, IndexMethod\n",
        "import vecs\n",
        "\n",
        "DB_STRING = f\"postgresql://postgres.gyiucazdpikoqpigfhvs:ntmWCuLtAfTeX8mk@aws-0-ca-central-1.pooler.supabase.com:6543/postgres\"\n",
        "\n",
        "# create vector store client\n",
        "vx = vecs.create_client(DB_STRING)"
      ],
      "metadata": {
        "id": "Akg88kS8WZCT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "list_of_vecs_collections = vx.list_collections()\n",
        "\n",
        "# Find our collection and delete it if it already exists\n",
        "matching_collection = next((col for col in list_of_vecs_collections if col.name == index_name), None)\n",
        "\n",
        "if matching_collection:\n",
        "    print(\"Collection already exists. Deleting now..\")\n",
        "    vx.delete_collection(matching_collection.name)\n",
        "\n",
        "vecs_collection = vx.get_or_create_collection(name=index_name, dimension=EMBEDDINGS_DIMENSIONS)\n",
        "\n",
        "# Create an index over our new collection (aka a Postgres table under the hood) for faster querying\n",
        "vecs_collection.create_index(method=IndexMethod.hnsw, measure=IndexMeasure.cosine_distance,\n",
        "                                        index_arguments=IndexArgsHNSW(m=8))\n",
        "\n",
        "print(\"Collection in vecs created.\")"
      ],
      "metadata": {
        "id": "RUgPuA_xXpYR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "records_to_upsert_to_vecs = [\n",
        "    (\n",
        "        item[\"id\"],          # the vector's identifier\n",
        "        item[\"values\"],  # the vector\n",
        "        item[\"metadata\"]     # associated metadata\n",
        "    )\n",
        "    for item in all_cases_embeddings_and_sparse_vectors\n",
        "]"
      ],
      "metadata": {
        "id": "25_FgPuERHF-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "vecs_collection.upsert(records=records_to_upsert_to_vecs)"
      ],
      "metadata": {
        "id": "w4VLXuEqY9Ne"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "vecs_similarity_search = vecs_collection.query(\n",
        "    data=search_query_as_vectors[\"dense\"], # Again, we use dense vector to search similarity here.\n",
        "    limit=2,\n",
        "    filters={},\n",
        "    measure=\"cosine_distance\",\n",
        "    include_value=False,\n",
        "    include_metadata=True,\n",
        ")\n",
        "\n",
        "pprint(vecs_similarity_search)\n",
        "\n"
      ],
      "metadata": {
        "id": "yr0T80cFfWAV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from supabase import create_client, Client\n",
        "from supabase.client import ClientOptions\n",
        "\n",
        "url: str = \"\"\n",
        "key: str = \"\"\n",
        "\n",
        "supabase: Client = create_client(url, key,\n",
        "  options=ClientOptions(\n",
        "    postgrest_client_timeout=10,\n",
        "    storage_client_timeout=10,\n",
        "    schema=\"public\",\n",
        "  ))"
      ],
      "metadata": {
        "id": "a1OxuMffeYp-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "vecs_full_text_search_results = (\n",
        "    supabase.table(\"legal_cases\")\n",
        "    .select(\"case_id, case_text\") # The columns we want back in our response\n",
        "    .text_search(\n",
        "        \"case_text\", # The column we want to search over\n",
        "        f\"'{search_query}'\", # Be careful here; it's sensitive to special characters and will error out\n",
        "        options={\"type\": \"websearch\", \"config\": \"english\"},\n",
        "    )\n",
        "    .execute()\n",
        ")\n",
        "\n",
        "pprint(vecs_full_text_search_results)"
      ],
      "metadata": {
        "id": "ouqMlxZFf8tw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "vecs_full_text_search_results = (\n",
        "    supabase.table(\"legal_cases\")\n",
        "    .select(\"case_id, case_text\") # The columns we want back in our response\n",
        "    .text_search(\n",
        "        \"case_text\", # The column we want to search over\n",
        "        \"'Palmer J' & 'Macleay Nominees Pty'\", # Be careful here; it's sensitive to special characters and will error out\n",
        "        options={\"type\": \"websearch\", \"config\": \"english\"},\n",
        "    )\n",
        "    .execute()\n",
        ")\n",
        "\n",
        "pprint(vecs_full_text_search_results)"
      ],
      "metadata": {
        "id": "3NOWkrWOkedm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Setup Jina Reranker"
      ],
      "metadata": {
        "id": "wPkuSeU9qfvs"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "JINA_RERANKER_URL = \"https://api.jina.ai/v1/rerank\"\n",
        "\n",
        "def jina_rerank(query: str, text_list: List[str]):\n",
        "    headers = {\"Content-Type\": \"application/json\", \"Authorization\": \"Bearer jina_d68362712b5143188d360eaadef63cf16WjSf5hb686SC-yBocaJLq-2xvo7\"}\n",
        "\n",
        "    json_data = {\n",
        "      \"model\": \"jina-reranker-v2-base-multilingual\",\n",
        "      \"documents\": text_list,\n",
        "      \"query\": query,\n",
        "      \"top_n\": 3,\n",
        "    }\n",
        "\n",
        "    response = requests.post(JINA_RERANKER_URL, headers=headers, data=json.dumps(json_data))\n",
        "    return response.json()"
      ],
      "metadata": {
        "id": "b390WLiHnVYC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "merged_list_for_reranking = []\n",
        "\n",
        "# Process the `vecs_similarity_search` data\n",
        "for case_id, case_data in vecs_similarity_search:\n",
        "    merged_list_for_reranking.append({\n",
        "        'case_id': case_id,\n",
        "        'case_text': case_data['case_text']\n",
        "    })\n",
        "\n",
        "# Process the `vecs_full_text_search_results` data\n",
        "for case_data in vecs_full_text_search_results.data:\n",
        "    merged_list_for_reranking.append({\n",
        "        'case_id': case_data['case_id'],\n",
        "        'case_text': case_data['case_text']\n",
        "    })\n",
        "\n",
        "pprint(merged_list_for_reranking)"
      ],
      "metadata": {
        "id": "osOg42_Dp3R7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "just_case_text = [item['case_text'] for item in merged_list_for_reranking]"
      ],
      "metadata": {
        "id": "zxc3fy-SylWx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "reranked_results = jina_rerank(search_query, just_case_text)\n",
        "\n",
        "pprint(reranked_results)"
      ],
      "metadata": {
        "id": "iWlvArpfpecQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## The verdict"
      ],
      "metadata": {
        "id": "HPwhP5GW0azJ"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "1.   Pinecone's developer experience was awesome. Honestly no complaints here other than costs (their pods pricing), but their serverless solution is quite affordable.\n",
        "2.   Weaviate's docs are ok-ish and took me longer than expected to setup. But they do try to abstract a lot of the vectorization and and sparse vectors away so you can get started quickly without having to know what they are.\n",
        "3. Whip up your own hybrid search with Postgres pgvector and full text search is fine as long as you extract the right keywords for full text search and pick the right components (especially reranker). This should be the cheapest option. I use Cohere's Reranker in production. Be careful of JinaAI's API downtime.\n",
        "\n",
        "BONUS: JinaAI's API is down a lot lol.\n"
      ],
      "metadata": {
        "id": "YoGfbtQq0cZO"
      }
    }
  ]
}