Vectors_v1.ipynb

{
    "cells": [
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Vectors v1\n",
                "\n",
                "### Goal\n",
                "\n",
                "Let users specify notification topics, for example:\n",
                "\n",
                "\"Hey BeepGPT, let me know when there’s an important engineering decision being made about the Fraud Detection Project.\"\n",
                "\n",
                "### Method\n",
                "\n",
                "This method will do the opposite of Vectors v0. We will store the topics in a vector store and do retrieval on each conversation. \n",
                "\n",
                "Also we will explore various different embedding models to try to determine the best for this use case.\n",
                "\n",
                "#### Pros:\n",
                "* Fast response. Only requires an LLM to do embeddings of the conversation and search over the vector space.\n",
                "* Fairly cheap to run.\n",
                "* Can potentially learn user's specific interests over time from via clustering analysis in the vector space.\n",
                "\n",
                "#### Cons:\n",
                "* Requires a vector database\n",
                "\n",
                "### Questions:\n",
                "* Will the topic embeddings be close to the associated conversations in the vector space?\n",
                "* Are there different embedding methods that will work better than others?"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "\n",
                "### The code:"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "#### Install the tools, initiate the things"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "%pip install -q kaskada openai llama-cpp-python ipywidgets"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import openai, getpass\n",
                "\n",
                "# Initialize OpenAI\n",
                "openai.api_key = getpass.getpass('OpenAI: API Key')"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import pandas as pd\n",
                "import kaskada as kd\n",
                "\n",
                "# Initialize Kaskada with a local execution context.\n",
                "kd.init_session()\n",
                "\n",
                "# set pandas to display all floats with 6 decimal places\n",
                "pd.options.display.float_format = '{:.6f}'.format"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "#### Pull in the user list, create a `format_user()` method"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "users_df = pd.read_json(\"slack-generation.users.json\")\n",
                "\n",
                "columns_to_keep = [\"id\", \"team_id\", \"name\", \"deleted\", \"real_name\", \"is_bot\", \"updated\"]\n",
                "\n",
                "users_df.drop(columns=users_df.columns.difference(columns_to_keep), inplace=True)\n",
                "\n",
                "users = {}\n",
                "for user in users_df.to_dict(orient='index').values():\n",
                "    users[user[\"id\"]] = user"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "def get_user(user_id):\n",
                "    return users[user_id] if user_id in users.keys() else None\n",
                "\n",
                "def format_user(user_id):\n",
                "    user = get_user(user_id)\n",
                "    return f\"{user['name']} ({user_id})\" if user else f\"({user_id})\"\n",
                "\n",
                "format_user(\"UBB9D2B01\")"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "#### Load the slack data, clean the message text, format message users"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# Load events from a Parquet file\n",
                "#\n",
                "# if you wan to load in your own slack data, change this to the path of your output file from 1.1 above\n",
                "# otherwise continue with `slack-generation.parquet`, which contains generated slack data for\n",
                "# example purposes. See the `slack-generation/notebook.ipynb` notebook for more info.\n",
                "input_file = \"slack-generation.parquet\"\n",
                "\n",
                "# Use the \"ts\" column as the time associated with each row,\n",
                "# and the \"channel\" column as the entity associated with each row.\n",
                "raw_msgs = await kd.sources.Parquet.create(\n",
                "    input_file,\n",
                "    time_column = \"ts\",\n",
                "    key_column = \"channel\",\n",
                "    time_unit = \"s\"\n",
                ")\n",
                "raw_msgs.preview(5)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import json\n",
                "\n",
                "@kd.udf(\"f<N: any>(x: N) -> string\")\n",
                "def format_users(batch: pd.Series):\n",
                "    # Apply to each row in the batch\n",
                "    return batch.map(format_user)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# Clean Text\n",
                "import re\n",
                "\n",
                "def strip_code_blocks(line):\n",
                "    return re.sub(r\"```.*?```\", '', line)\n",
                "\n",
                "def user_repl(match_obj):\n",
                "    user_id = match_obj.group(1)\n",
                "    return format_user(user_id)\n",
                "\n",
                "def update_users(line):\n",
                "    return re.sub(r\"<@(.*?)>\", user_repl, line)\n",
                "\n",
                "def clean_message(text):\n",
                "        text = strip_code_blocks(update_users(text)).strip()\n",
                "        return None if text == \"\" else text\n",
                "\n",
                "@kd.udf(\"f<N: any>(x: N) -> string\")\n",
                "def clean_text(batch: pd.Series):\n",
                "    # Apply to each row in the batch\n",
                "    return batch.map(clean_message)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "formatted_msgs = raw_msgs.extend({\n",
                "    \"text\": raw_msgs.col(\"text\").pipe(clean_text),\n",
                "    \"user\": raw_msgs.col(\"user\").pipe(format_users)\n",
                "})\n",
                "formatted_msgs.preview(5)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "#### Convert the non-threaded messages into threaded messages. See `FineTuning_v2.ipynb` for more details on this."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "from datetime import timedelta\n",
                "\n",
                "ts = formatted_msgs.col(\"ts\")\n",
                "thread_ts = formatted_msgs.col(\"thread_ts\")\n",
                "\n",
                "# split messages into two subgroups: threads and non-threads\n",
                "threads = formatted_msgs.filter(thread_ts.is_not_null())\n",
                "non_threads = formatted_msgs.filter(thread_ts.is_null())\n",
                "\n",
                "# for non-threads, consider a message a new conversation when\n",
                "# more than 10 mins have elapsed since the previous message\n",
                "is_new = ts.seconds_since_previous() > timedelta(minutes=10)\n",
                "\n",
                "# Eventually this will just be: `thread_ts = ts.first(window=kd.windows.Since(is_new, start=\"inclusive\"))`\n",
                "#\n",
                "# However, `Since()` is currently exclusive on the start of the window, inclusive on the end.\n",
                "# But we need inclusive on the start and exclusive on the end.\n",
                "#\n",
                "# The hack below does what we need until `Since()` provides additional options for inclusivity\n",
                "shifted_non_threads = non_threads.shift_by(timedelta(microseconds=0.001))\n",
                "shifted_ts = shifted_non_threads.lag(1).col(\"ts\").first(window=kd.windows.Since(is_new))\n",
                "thread_ts = ts.if_(is_new).else_(shifted_ts)\n",
                "\n",
                "# create threads_ts column for non-threaded messages\n",
                "non_threads_threads = non_threads.extend({\"thread_ts\": thread_ts}).filter(ts.is_not_null().and_(thread_ts.is_not_null()))\n",
                "\n",
                "# re-join the two message subgroups\n",
                "joined = threads.else_(non_threads_threads)\n",
                "\n",
                "# join non-threads and threads back up, and key by conversations\n",
                "messages = joined.with_key(kd.record({\n",
                "        \"channel\": joined.col(\"channel\"),\n",
                "        \"thread\": joined.col(\"thread_ts\"),\n",
                "    }))\n",
                "messages.preview(5)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "#### Collect up all the messages, reactions, users in each conversation."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "@kd.udf(\"f<N: any>(x: N) -> string\")\n",
                "def format_message(batch: pd.Series):\n",
                "    def formatter(raw):\n",
                "        return f\"{raw['user']} --> {raw['text']}\" # --> {raw['reactions']}\"\n",
                "    return batch.map(formatter)\n",
                "\n",
                "@kd.udf(\"f<N: any>(x: N) -> string\")\n",
                "def format_messages(batch: pd.Series):\n",
                "    def formatter(raw):\n",
                "        return \"\\n---\\n\".join(raw)\n",
                "    return batch.map(formatter)\n",
                "\n",
                "@kd.udf(\"f<N: any>(x: N) -> string\")\n",
                "def extract_users(batch: pd.Series):\n",
                "    def get_users(raw):\n",
                "        users = [raw[\"user\"]]\n",
                "        # for user in json.loads(raw[\"reactions\"]).keys():\n",
                "        #     if user not in users:\n",
                "        #         users.append(user)\n",
                "        return json.dumps(users)\n",
                "    return batch.map(get_users)\n",
                "\n",
                "@kd.udf(\"f<N: any>(x: N) -> string\")\n",
                "def unique_users(batch: pd.Series):\n",
                "    def get_users(raw):\n",
                "        users = []\n",
                "        for user_set in raw:\n",
                "            users.extend(json.loads(user_set))\n",
                "        return json.dumps(list(set(users)))\n",
                "    return batch.map(get_users)\n",
                "\n",
                "conversations = kd.record({\n",
                "    # \"conversation\": messages.select(\"user\", \"text\", \"reactions\").pipe(format_message).collect(max=None).pipe(format_messages),\n",
                "    # \"users\": messages.select(\"user\", \"reactions\").pipe(extract_users).collect(max=None).pipe(unique_users),\n",
                "    \"conversation\": messages.select(\"user\", \"text\").pipe(format_message).collect(max=None).pipe(format_messages),\n",
                "    \"users\": messages.select(\"user\").pipe(extract_users).collect(max=None).pipe(unique_users),\n",
                "})\n",
                "conversations.preview(5)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "#### Create LlamaIndex documents for all the topics"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "from llama_index import Document\n",
                "\n",
                "topics = {\n",
                "    \"ID1\": \"Tell me about engineering discussions related to the Supply Chain Management project\",\n",
                "    \"ID2\": \"Alert me when people are making streaming technology decisions\",\n",
                "    \"ID3\": \"Poke me when there are people chatting about SRE topics like monitoring and alerting\",\n",
                "    \"ID4\": \"Let me know when people are talking about their weekends\",\n",
                "    \"ID5\": \"Inform me of any important discussions happening on the Fraud Detection project\",\n",
                "}\n",
                "\n",
                "topic_documents = []\n",
                "for id in topics.keys():\n",
                "    document = Document(\n",
                "        doc_id=id,\n",
                "        text=topics[id],\n",
                "        metadata={\n",
                "            \"id\": id,\n",
                "        })\n",
                "    document.excluded_llm_metadata_keys = [\"id\"]\n",
                "    document.excluded_embed_metadata_keys = [\"id\"]\n",
                "    topic_documents.append(document)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "#### Create LlamaIndex documents for all the conversations"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import json\n",
                "\n",
                "conversation_documents = []\n",
                "max_conversation_length = 0\n",
                "async for row in conversations.run_iter(results=kd.results.Snapshot(), kind=\"row\"):\n",
                "    id = json.dumps(row[\"_key\"])\n",
                "    conversation = row[\"conversation\"]\n",
                "    users = row[\"users\"]\n",
                "\n",
                "    max_conversation_length = max(max_conversation_length, len(conversation))\n",
                "\n",
                "    document = Document(\n",
                "        doc_id=id,\n",
                "        text=conversation,\n",
                "        metadata={\n",
                "            \"users\": users,\n",
                "        })\n",
                "    document.excluded_llm_metadata_keys = [\"users\"]\n",
                "    document.excluded_embed_metadata_keys = [\"users\"]\n",
                "    conversation_documents.append(document)\n",
                "\n",
                "print(f\"Max conversation length in chars: {max_conversation_length}\")"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "#### Using LlamaIndex and HuggingFace, create embeddings for all the topics and conversations using various different embedding models.\n",
                "\n",
                "* Store the embeddings for topics in a `topic_indexes` map of vector indexes\n",
                "* Store the embeddings for conversations in a `conversation_indexes` map of vector indexes"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "from llama_index import VectorStoreIndex, ServiceContext\n",
                "from llama_index.embeddings import HuggingFaceEmbedding, InstructorEmbedding\n",
                "from llama_index.embeddings.utils import EmbedType\n",
                "from llama_index.node_parser.simple import SimpleNodeParser\n",
                "\n",
                "topic_indexes = {}\n",
                "conversation_indexes = {}\n",
                "\n",
                "embedding_models: {str:EmbedType} = {\n",
                "    # This will use llama2-chat-13B from with LlamaCPP, and assumes you have llama-cpp-python installed\n",
                "    \"llama2\": \"local\",\n",
                "    # This will use open-ai embeddings, and assumes OpenAI has already been initialized\n",
                "    \"openai\": \"default\",\n",
                "\n",
                "    \"bge-small-en-v1.5\" : HuggingFaceEmbedding(model_name=\"BAAI/bge-small-en-v1.5\"),\n",
                "    \"bge-base-en-v1.5\" : HuggingFaceEmbedding(model_name=\"BAAI/bge-base-en-v1.5\"),\n",
                "    \"bge-large-en-v1.5\" : HuggingFaceEmbedding(model_name=\"BAAI/bge-large-en-v1.5\"),\n",
                "    \"ember-v1\" : HuggingFaceEmbedding(model_name=\"llmrails/ember-v1\"),\n",
                "    \"gte-large\" : HuggingFaceEmbedding(model_name=\"thenlper/gte-large\"),\n",
                "    \"gte-base\" : HuggingFaceEmbedding(model_name=\"thenlper/gte-base\"),\n",
                "    \"e5-large-v2\" : HuggingFaceEmbedding(model_name=\"intfloat/e5-large-v2\"),\n",
                "    \"e5-base-v2\" : HuggingFaceEmbedding(model_name=\"intfloat/e5-base-v2\"),\n",
                "}\n",
                "\n",
                "# limit documents to single-node\n",
                "node_parser = SimpleNodeParser.from_defaults(\n",
                "    chunk_size = max_conversation_length\n",
                ")\n",
                "\n",
                "for model_name in embedding_models.keys():\n",
                "    service_context = ServiceContext.from_defaults(\n",
                "        llm=None, # explicitly disable llm use\n",
                "        embed_model=embedding_models[model_name],\n",
                "        node_parser=node_parser,\n",
                "    )\n",
                "    topic_indexes[model_name] = VectorStoreIndex([], service_context=service_context)\n",
                "    conversation_indexes[model_name] = VectorStoreIndex([], service_context=service_context)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# insert docs and generate embeddings for everything\n",
                "\n",
                "for model_name in embedding_models.keys():\n",
                "    for topic_doc in topic_documents:\n",
                "        topic_indexes[model_name].insert(topic_doc)\n",
                "    for conversation_doc in conversation_documents:\n",
                "        conversation_indexes[model_name].insert(conversation_doc)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "#### For each embedding model, perform retrieval on each conversation embedding against the topic embeddings vector index\n",
                "\n",
                "Save score results to files."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "base_result = {}\n",
                "topic_count = len(topics.keys())\n",
                "\n",
                "for topic_id in topics.keys():\n",
                "    base_result[topic_id] = 0.0\n",
                "\n",
                "from llama_index.indices.vector_store.retrievers import VectorIndexRetriever\n",
                "from llama_index.indices.query.schema import QueryBundle\n",
                "\n",
                "from llama_index.indices.vector_store.retrievers import VectorIndexRetriever\n",
                "\n",
                "\n",
                "for model_name in embedding_models.keys():\n",
                "    topic_index = topic_indexes[model_name]\n",
                "    conversation_index = conversation_indexes[model_name]\n",
                "\n",
                "    retriever = VectorIndexRetriever(\n",
                "        index=topic_index,\n",
                "        similarity_top_k=topic_count,\n",
                "        vector_store_query_mode=\"default\",\n",
                "        alpha=None,\n",
                "        doc_ids=None,\n",
                "    )\n",
                "\n",
                "    with open(f\"Vectors_v1_{model_name}_scores.jsonl\", \"w\") as out_file:\n",
                "        for doc_id in conversation_index.ref_doc_info.keys():\n",
                "            parsed_doc_id = json.loads(doc_id)\n",
                "            for node_id in conversation_index.ref_doc_info[doc_id].node_ids:\n",
                "                embedding = conversation_index.vector_store.get(node_id)\n",
                "                text = conversation_index.docstore.docs[node_id].text\n",
                "\n",
                "                qb = QueryBundle(query_str=text, embedding=embedding)\n",
                "                result = base_result.copy()\n",
                "                for node_score in retriever.retrieve(qb):\n",
                "                    topic_id = node_score.metadata[\"id\"]\n",
                "                    result[topic_id] = node_score.score\n",
                "                out = {\"id\": parsed_doc_id, \"result\": result}\n",
                "                out_file.write(json.dumps(out) + \"\\n\")"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "#### Load the results from `ChatCompletions_v1` to use as a baseline for comparison to this method."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "chat_completion_results = {}\n",
                "\n",
                "with open(\"ChatCompletion_v1_results_avg.jsonl\") as input:\n",
                "    for line in input.readlines():\n",
                "        result = json.loads(line)\n",
                "        key = json.dumps(result[\"id\"])\n",
                "        chat_completion_results[key] = result[\"result\"]"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "#### Create methods to optimize the altering threshold in order to maximize the f1 score on each embedding model\n",
                "\n",
                "Make comparisons against the data from `chat_completion_results`"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "def get_results_array(file_name:str, threshold:float):\n",
                "    results = []\n",
                "    with open(file_name) as in_file:\n",
                "        for line in in_file.readlines():\n",
                "            row = json.loads(line)\n",
                "            true_ids = []\n",
                "            for id in row[\"result\"].keys():\n",
                "                value = row[\"result\"][id]\n",
                "                if isinstance(value, float):\n",
                "                    if value > threshold:\n",
                "                        true_ids.append(id)\n",
                "                elif isinstance(value, bool):\n",
                "                    if value:\n",
                "                        true_ids.append(id)\n",
                "                else:\n",
                "                    print(f\"unknown value type: {type(value)}\")\n",
                "            results.append(true_ids)\n",
                "    return results"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "from sklearn.preprocessing import MultiLabelBinarizer\n",
                "from sklearn.metrics import f1_score\n",
                "\n",
                "def get_f1_score(file_name, threshold):\n",
                "    mlb = MultiLabelBinarizer()\n",
                "    mlb.fit([topics.keys()])\n",
                "\n",
                "    test = get_results_array(\"ChatCompletion_v1_results_avg.jsonl\", 0.0)\n",
                "    pred = get_results_array(file_name, threshold)\n",
                "\n",
                "    y_test_transformed = mlb.transform(test)\n",
                "    y_pred_transformed = mlb.transform(pred)\n",
                "\n",
                "    return f1_score(y_test_transformed, y_pred_transformed, average='macro') # Or 'micro', 'weighted' based on need"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import numpy as np\n",
                "from scipy.optimize import minimize_scalar\n",
                "\n",
                "def optimize_f1_score(file_name) -> (float, float):\n",
                "    # Use the minimize function to maximize the objective function\n",
                "    result = minimize_scalar(fun=lambda x: -get_f1_score(file_name, x), bounds=(0.1, 1.0))\n",
                "\n",
                "    return (result.x, -result.fun)"
            ]
        },
        {
            "cell_type": "code",
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Model: llama2 \t\t\tf1: 0.36 \tthreshold: 0.82\n",
                        "Model: openai \t\t\tf1: 0.52 \tthreshold: 0.77\n",
                        "Model: bge-small-en-v1.5 \tf1: 0.42 \tthreshold: 0.65\n",
                        "Model: bge-base-en-v1.5 \tf1: 0.47 \tthreshold: 0.57\n",
                        "Model: bge-large-en-v1.5 \tf1: 0.52 \tthreshold: 0.63\n",
                        "Model: ember-v1 \t\tf1: 0.48 \tthreshold: 0.57\n",
                        "Model: gte-large \t\tf1: 0.29 \tthreshold: 0.44\n",
                        "Model: gte-base \t\tf1: 0.49 \tthreshold: 0.79\n",
                        "Model: e5-large-v2 \t\tf1: 0.29 \tthreshold: 0.44\n",
                        "Model: e5-base-v2 \t\tf1: 0.33 \tthreshold: 0.7\n"
                    ]
                }
            ],
            "source": [
                "for model_name in embedding_models.keys():\n",
                "    file_name = f\"Vectors_v1_{model_name}_scores.jsonl\"\n",
                "\n",
                "    (threshold, score) = optimize_f1_score(file_name)\n",
                "\n",
                "    print(f\"Model: {model_name} \\tf1: {round(score,2)} \\tthreshold: {round(threshold,2)}\")"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "#### Results\n",
                "\n",
                "OpenAI and bge-large-en-v1.5 perform the best, however their f1 score doesn't come close to the level we were able to obtain using fine-tuning."
            ]
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "Python 3",
            "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.11.4"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 2
}