Customizing_embeddings.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "Vq31CdSRpgkI"
   },
   "source": [
    "# Customizing embeddings\n",
    "\n",
    "This notebook demonstrates one way to customize OpenAI embeddings to a particular task.\n",
    "\n",
    "The input is training data in the form of [text_1, text_2, label] where label is +1 if the pairs are similar and -1 if the pairs are dissimilar.\n",
    "\n",
    "The output is a matrix that you can use to multiply your embeddings. The product of this multiplication is a 'custom embedding' that will better emphasize aspects of the text relevant to your use case. In binary classification use cases, we've seen error rates drop by as much as 50%.\n",
    "\n",
    "In the following example, I use 1,000 sentence pairs picked from the SNLI corpus. Each pair of sentences are logically entailed (i.e., one implies the other). These pairs are our positives (label = 1). We generate synthetic negatives by combining sentences from different pairs, which are presumed to not be logically entailed (label = -1).\n",
    "\n",
    "For a clustering use case, you can generate positives by creating pairs from texts in the same clusters and generate negatives by creating pairs from sentences in different clusters.\n",
    "\n",
    "With other data sets, we have seen decent improvement with as little as ~100 training examples. Of course, performance will be better with  more examples."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "arB38jFwpgkK"
   },
   "source": [
    "# 0. Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ifvM7g4apgkK"
   },
   "outputs": [],
   "source": [
    "# imports\n",
    "from typing import List, Tuple # for type hints\n",
    "from typing_extensions import Annotated \n",
    "\n",
    "import numpy as np  # for manipulating arrays\n",
    "import pandas as pd  # for manipulating data in dataframes\n",
    "import pickle  # for saving the embeddings cache\n",
    "import plotly.express as px  # for plots\n",
    "import random  # for generating run IDs\n",
    "from sklearn.model_selection import train_test_split  # for splitting train & test data\n",
    "import torch  # for matrix optimization\n",
    "\n",
    "from zenml import step, pipeline\n",
    "\n",
    "from zenml.logger import get_logger\n",
    "\n",
    "logger = get_logger(__name__)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "DtBbryAapgkL"
   },
   "source": [
    "## 1. Inputs\n",
    "\n",
    "Most inputs are here. The key things to change are where to load your datset from, where to save a cache of embeddings to, and which embedding engine you want to use.\n",
    "\n",
    "Depending on how your data is formatted, you'll want to rewrite the process_input_data function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "import zipfile\n",
    "import os\n",
    "import pandas as pd \n",
    "from typing_extensions import Tuple \n",
    "\n",
    "@step\n",
    "def download_data(\n",
    "    snli_url: str = 'https://nlp.stanford.edu/projects/snli/snli_1.0.zip',\n",
    "    save_path: str = 'data/snli_1.0.zip',\n",
    "    extracted_dir: str = 'data/snli_dataset'\n",
    ") -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:\n",
    "    \"\"\"Returns the path to the downloaded dataset.\"\"\"\n",
    "    # Ensure the directory exists\n",
    "    os.makedirs(os.path.dirname(save_path), exist_ok=True)\n",
    "\n",
    "    # Check if the zip file has already been downloaded\n",
    "    if not os.path.exists(save_path):\n",
    "        # Download the file\n",
    "        response = requests.get(snli_url, stream=True)\n",
    "\n",
    "        # Check if the request was successful\n",
    "        if response.status_code == 200:\n",
    "            with open(save_path, 'wb') as f:\n",
    "                f.write(response.content)\n",
    "            print(f'Download complete. Dataset saved to {save_path}')\n",
    "        else:\n",
    "            print(f'Failed to download the dataset. Status code: {response.status_code}')\n",
    "            return None  # Return None or raise an exception if the download failed\n",
    "\n",
    "    # Unzip the downloaded file\n",
    "    with zipfile.ZipFile(save_path, 'r') as zip_ref:\n",
    "        zip_ref.extractall(extracted_dir)\n",
    "        print(f'Extraction complete. Dataset extracted to {extracted_dir}')\n",
    "\n",
    "    \n",
    "    # Read the JSON Lines file into a DataFrame\n",
    "    train_df = pd.read_json(os.path.join(extracted_dir, \"snli_1.0/snli_1.0_train.jsonl\"), lines=True)\n",
    "    test_df = pd.read_json(os.path.join(extracted_dir, \"snli_1.0/snli_1.0_test.jsonl\"), lines=True)\n",
    "    dev_df = pd.read_json(os.path.join(extracted_dir, \"snli_1.0/snli_1.0_dev.jsonl\"), lines=True)\n",
    "\n",
    "    # Return the dataframes\n",
    "    return train_df, test_df, dev_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "download_data.entrypoint()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "UzxcWRCkpgkM"
   },
   "outputs": [],
   "source": [
    "@step\n",
    "def process_input_data(df: pd.DataFrame, dataset_type: str, num_pairs_to_embed: int = 1000) -> pd.DataFrame:\n",
    "    # you can customize this to preprocess your own dataset\n",
    "    # output should be a dataframe with 3 columns: text_1, text_2, label (1 for similar, -1 for dissimilar)\n",
    "    df[\"label\"] = df[\"gold_label\"]\n",
    "    df = df[df[\"label\"].isin([\"entailment\"])]\n",
    "    df[\"label\"] = df[\"label\"].apply(lambda x: {\"entailment\": 1, \"contradiction\": -1}[x])\n",
    "    df = df.rename(columns={\"sentence1\": \"text_1\", \"sentence2\": \"text_2\"})\n",
    "    df = df[[\"text_1\", \"text_2\", \"label\"]]\n",
    "    df = df.head(num_pairs_to_embed)\n",
    "    \n",
    "    df.loc[:, \"dataset\"] = dataset_type\n",
    "    \n",
    "    return df\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "aBbH71hEpgkM"
   },
   "source": [
    "## 2. Load and process input data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@pipeline\n",
    "def load_and_process_data():\n",
    "    train_df, test_df, dev_df = download_data()\n",
    "    processed_train_df = process_input_data(train_df, dataset_type=\"train\")\n",
    "    processed_test_df = process_input_data(test_df, dataset_type=\"test\")\n",
    "    processed_dev_df = process_input_data(dev_df, dataset_type=\"dev\")\n",
    "    \n",
    "load_and_process_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kAKLjYG6pgkN",
    "outputId": "dc178688-e97d-4ad0-b26c-dff67b858966"
   },
   "outputs": [],
   "source": [
    "# load data\n",
    "from zenml.client import Client \n",
    "\n",
    "client = Client()\n",
    "\n",
    "latest_run = client.get_pipeline(\"load_and_process_data\").runs[-1]\n",
    "df = latest_run.steps[\"process_input_data\"].output.load()\n",
    "df"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "MzAFkA2opgkP"
   },
   "source": [
    "## 3. Generate synthetic negatives\n",
    "\n",
    "This is another piece of the code that you will need to modify to match your use case.\n",
    "\n",
    "If you have data with positives and negatives, you can skip this section.\n",
    "\n",
    "If you have data with only positives, you can mostly keep it as is, where it generates negatives only.\n",
    "\n",
    "If you have multiclass data, you will want to generate both positives and negatives. The positives can be pairs of text that share labels, and the negatives can be pairs of text that do not share labels.\n",
    "\n",
    "The final output should be a dataframe with text pairs, where each pair is labeled -1 or 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rUYd9V0zpgkP"
   },
   "outputs": [],
   "source": [
    "# generate negatives\n",
    "@step\n",
    "def dataframe_of_negatives(dataframe_of_positives: pd.DataFrame, dataset_type: str) -> pd.DataFrame:\n",
    "    \"\"\"Return dataframe of negative pairs made by combining elements of positive pairs.\"\"\"\n",
    "    texts = set(dataframe_of_positives[\"text_1\"].values) | set(\n",
    "        dataframe_of_positives[\"text_2\"].values\n",
    "    )\n",
    "    all_pairs = {(t1, t2) for t1 in texts for t2 in texts if t1 < t2}\n",
    "    positive_pairs = set(\n",
    "        tuple(text_pair)\n",
    "        for text_pair in dataframe_of_positives[[\"text_1\", \"text_2\"]].values\n",
    "    )\n",
    "    negative_pairs = all_pairs - positive_pairs\n",
    "    df_of_negatives = pd.DataFrame(list(negative_pairs), columns=[\"text_1\", \"text_2\"])\n",
    "    df_of_negatives[\"label\"] = -1\n",
    "    \n",
    "    df_of_negatives[\"dataset\"] = dataset_type\n",
    "    \n",
    "    return df_of_negatives"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Rkh8-J89pgkP"
   },
   "outputs": [],
   "source": [
    "@step \n",
    "def concatenate_dataframes(\n",
    "        train_df: pd.DataFrame,\n",
    "        train_df_negatives: pd.DataFrame, \n",
    "        test_df: pd.DataFrame, \n",
    "        test_df_negatives: pd.DataFrame, \n",
    "        random_seed: int = 93,\n",
    "        negatives_per_positive = (1), # it will work at higher values too, but more data will be slower\n",
    "    ) -> pd.DataFrame:\n",
    "    # sample negatives and combine with positives\n",
    "    train_df = pd.concat(\n",
    "        [\n",
    "            train_df,\n",
    "            train_df_negatives.sample(\n",
    "                n=len(train_df) * negatives_per_positive, random_state=random_seed\n",
    "            ),\n",
    "        ]\n",
    "    )\n",
    "    test_df = pd.concat(\n",
    "        [\n",
    "            test_df,\n",
    "            test_df_negatives.sample(\n",
    "                n=len(test_df) * negatives_per_positive, random_state=random_seed\n",
    "            ),\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    big_df = pd.concat([train_df, test_df])\n",
    "\n",
    "    # sample it to make this faster\n",
    "    big_df = big_df.sample(frac=0.3, random_state=random_seed).reset_index(drop=True)\n",
    "    \n",
    "    return big_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@pipeline\n",
    "def load_and_process_data():\n",
    "    train_df, test_df, _ = download_data()\n",
    "    processed_train_df = process_input_data(train_df, dataset_type=\"train\")\n",
    "    processed_test_df = process_input_data(test_df, dataset_type=\"test\")\n",
    "\n",
    "    train_df_negatives = dataframe_of_negatives(processed_train_df, dataset_type=\"train\")\n",
    "    test_df_negatives = dataframe_of_negatives(processed_test_df, dataset_type=\"test\")\n",
    "\n",
    "    big_df = concatenate_dataframes(\n",
    "        processed_train_df,\n",
    "        train_df_negatives, \n",
    "        processed_test_df, \n",
    "        test_df_negatives, \n",
    "    )\n",
    "    \n",
    "load_and_process_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load data\n",
    "from zenml.client import Client \n",
    "\n",
    "client = Client()\n",
    "\n",
    "latest_run = client.get_pipeline(\"load_and_process_data\").runs[-1]\n",
    "big_df = latest_run.steps[\"concatenate_dataframes\"].output.load()\n",
    "big_df.head()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "8MVSLMSrpgkQ"
   },
   "source": [
    "## 4. Calculate embeddings and cosine similarities\n",
    "\n",
    "Here, I create a cache to save the embeddings. This is handy so that you don't have to pay again if you want to run the code again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.embeddings_utils import get_embedding, cosine_similarity  # for embeddings\n",
    "\n",
    "# establish a cache of embeddings to avoid recomputing\n",
    "# cache is a dict of tuples (text, engine) -> embedding\n",
    "def get_embedding_cache(embedding_cache_path: str) -> dict:\n",
    "    try:\n",
    "        with open(embedding_cache_path, \"rb\") as f:\n",
    "            embedding_cache = pickle.load(f)\n",
    "            logger.info(\"Loaded embeddings from cache.\")\n",
    "    except FileNotFoundError:\n",
    "        precomputed_embedding_cache_path = \"https://cdn.openai.com/API/examples/data/snli_embedding_cache.pkl\"\n",
    "        embedding_cache = pd.read_pickle(precomputed_embedding_cache_path)\n",
    "        logger.info(\"Loaded embeddings from OpenAI.\")\n",
    "    return embedding_cache\n",
    "\n",
    "embedding_cache_path = \"data/get_embedding_cache.pkl\"\n",
    "embedding_cache = get_embedding_cache(embedding_cache_path)\n",
    "\n",
    "# this function will get embeddings from the cache and save them there afterward\n",
    "@step\n",
    "def calculate_embeddings_and_cos_similarity(\n",
    "    big_df: pd.DataFrame,\n",
    "    default_embedding_engine: str = \"babbage-similarity\"\n",
    ") -> pd.DataFrame:\n",
    "    \"\"\"Calculate embeddings and cosine similarity for each pair of texts in the dataframe.\"\"\"\n",
    "    \n",
    "    # Set OpenAI key\n",
    "    secret = Client().get_secret(\"openai\")\n",
    "    os.environ[\"OPENAI_API_KEY\"] = secret.secret_values[\"api_key\"]\n",
    "\n",
    "    def get_embedding_with_cache(\n",
    "        text: str,\n",
    "        engine: str,\n",
    "        embedding_cache_path: str,\n",
    "    ) -> list:\n",
    "        if (text, engine) not in embedding_cache.keys():\n",
    "            # if not in cache, call API to get embedding\n",
    "            logger.info(text)\n",
    "            embedding_cache[(text, engine)] = get_embedding(text, engine)\n",
    "            # save embeddings cache to disk after each update\n",
    "            with open(embedding_cache_path, \"wb\") as embedding_cache_file:\n",
    "                pickle.dump(embedding_cache, embedding_cache_file)\n",
    "        return embedding_cache[(text, engine)]\n",
    "\n",
    "    \n",
    "    # create column of embeddings\n",
    "    big_df = big_df.dropna(subset=['text_1', 'text_2'])\n",
    "    for column in [\"text_1\", \"text_2\"]:\n",
    "        big_df[f\"{column}_embedding\"] = big_df[column].apply(\n",
    "            lambda text: get_embedding_with_cache(text, embedding_cache_path=embedding_cache_path, engine=default_embedding_engine)\n",
    "        )\n",
    "\n",
    "    # create column of cosine similarity between embeddings\n",
    "    big_df[\"cosine_similarity\"] = big_df.apply(\n",
    "        lambda row: cosine_similarity(row[\"text_1_embedding\"], row[\"text_2_embedding\"]),\n",
    "        axis=1,\n",
    "    )\n",
    "    \n",
    "    \n",
    "    # Return the big_df with embeddings and cosine similarity and the final embedding cache\n",
    "    return big_df\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@pipeline\n",
    "def load_and_process_data():\n",
    "    train_df, test_df, _ = download_data()\n",
    "    processed_train_df = process_input_data(train_df, dataset_type=\"train\")\n",
    "    processed_test_df = process_input_data(test_df, dataset_type=\"test\")\n",
    "\n",
    "    train_df_negatives = dataframe_of_negatives(processed_train_df, dataset_type=\"train\")\n",
    "    test_df_negatives = dataframe_of_negatives(processed_test_df, dataset_type=\"test\")\n",
    "\n",
    "    big_df = concatenate_dataframes(\n",
    "        train_df,\n",
    "        train_df_negatives, \n",
    "        test_df, \n",
    "        test_df_negatives, \n",
    "    )\n",
    "    \n",
    "    big_df_with_cos = calculate_embeddings_and_cos_similarity(big_df)\n",
    "    \n",
    "load_and_process_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load data\n",
    "from zenml.client import Client \n",
    "\n",
    "client = Client()\n",
    "\n",
    "latest_run = client.get_pipeline(\"load_and_process_data\").runs[-1]\n",
    "df = latest_run.steps[\"calculate_embeddings_and_cos_similarity\"].output.load()\n",
    "df"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "4pwn608LpgkQ"
   },
   "source": [
    "## 6. Plot distribution of cosine similarity\n",
    "\n",
    "Here we measure similarity of text using cosine similarity. In our experience, most distance functions (L1, L2, cosine similarity) all work about the same. Note that our embeddings are already normalized to length 1, so cosine similarity is equivalent to dot product.\n",
    "\n",
    "The graphs show how much the overlap there is between the distribution of cosine similarities for similar and dissimilar pairs. If there is a high amount of overlap, that means there are some dissimilar pairs with greater cosine similarity than some similar pairs.\n",
    "\n",
    "The accuracy I compute is the accuracy of a simple rule that predicts 'similar (1)' if the cosine similarity is above some threshold X and otherwise predicts 'dissimilar (0)'."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SoeDF8vqpgkQ",
    "outputId": "17db817e-1702-4089-c4e8-8ca32d294930"
   },
   "outputs": [],
   "source": [
    "# calculate accuracy (and its standard error) of predicting label=1 if similarity>x\n",
    "# x is optimized by sweeping from -1 to 1 in steps of 0.01\n",
    "def accuracy_and_se(cosine_similarity: float, labeled_similarity: int) -> Tuple[float]:\n",
    "    accuracies = []\n",
    "    for threshold_thousandths in range(-1000, 1000, 1):\n",
    "        threshold = threshold_thousandths / 1000\n",
    "        total = 0\n",
    "        correct = 0\n",
    "        for cs, ls in zip(cosine_similarity, labeled_similarity):\n",
    "            total += 1\n",
    "            if cs > threshold:\n",
    "                prediction = 1\n",
    "            else:\n",
    "                prediction = -1\n",
    "            if prediction == ls:\n",
    "                correct += 1\n",
    "        accuracy = correct / total\n",
    "        accuracies.append(accuracy)\n",
    "    a = max(accuracies)\n",
    "    n = len(cosine_similarity)\n",
    "    standard_error = (a * (1 - a) / n) ** 0.5  # standard error of binomial\n",
    "    return a, standard_error\n",
    "\n",
    "\n",
    "# check that training and test sets are balanced\n",
    "px.histogram(\n",
    "    df,\n",
    "    x=\"cosine_similarity\",\n",
    "    color=\"label\",\n",
    "    barmode=\"overlay\",\n",
    "    width=500,\n",
    "    facet_row=\"dataset\",\n",
    ").show()\n",
    "\n",
    "for dataset in [\"train\", \"test\"]:\n",
    "    data = df[df[\"dataset\"] == dataset]\n",
    "    a, se = accuracy_and_se(data[\"cosine_similarity\"], data[\"label\"])\n",
    "    print(f\"{dataset} accuracy: {a:0.1%} ± {1.96 * se:0.1%}\")\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "zHLxlnsApgkR"
   },
   "source": [
    "## 7. Optimize the matrix using the training data provided"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "z52V0x8IpgkR"
   },
   "outputs": [],
   "source": [
    "def embedding_multiplied_by_matrix(\n",
    "    embedding: List[float], matrix: torch.tensor\n",
    ") -> np.array:\n",
    "    embedding_tensor = torch.tensor(embedding).float()\n",
    "    modified_embedding = embedding_tensor @ matrix\n",
    "    modified_embedding = modified_embedding.detach().numpy()\n",
    "    return modified_embedding\n",
    "\n",
    "\n",
    "# compute custom embeddings and new cosine similarities\n",
    "@step\n",
    "def apply_matrix_to_embeddings_dataframe(matrix: torch.tensor, df: pd.DataFrame) -> pd.DataFrame:\n",
    "    for column in [\"text_1_embedding\", \"text_2_embedding\"]:\n",
    "        df[f\"{column}_custom\"] = df[column].apply(\n",
    "            lambda x: embedding_multiplied_by_matrix(x, matrix)\n",
    "        )\n",
    "    df[\"cosine_similarity_custom\"] = df.apply(\n",
    "        lambda row: cosine_similarity(\n",
    "            row[\"text_1_embedding_custom\"], row[\"text_2_embedding_custom\"]\n",
    "        ),\n",
    "        axis=1,\n",
    "    )\n",
    "    return df\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nlcUW-zEpgkS",
    "outputId": "4bd4bdff-628a-406f-fffe-aedbfad66446"
   },
   "outputs": [],
   "source": [
    "def optimize_matrix(\n",
    "    df: pd.DataFrame,\n",
    "    modified_embedding_length: int = 2048,  # in my brief experimentation, bigger was better (2048 is length of babbage encoding)\n",
    "    batch_size: int = 100,\n",
    "    max_epochs: int = 100,\n",
    "    learning_rate: float = 100.0,  # seemed to work best when similar to batch size - feel free to try a range of values\n",
    "    dropout_fraction: float = 0.0,  # in my testing, dropout helped by a couple percentage points (definitely not necessary)\n",
    "    print_progress: bool = True,\n",
    "    save_results: bool = True,\n",
    ") -> torch.tensor:\n",
    "    \"\"\"Return matrix optimized to minimize loss on training data.\"\"\"\n",
    "    run_id = random.randint(0, 2 ** 31 - 1)  # (range is arbitrary)\n",
    "    # convert from dataframe to torch tensors\n",
    "    # e is for embedding, s for similarity label\n",
    "    def tensors_from_dataframe(\n",
    "        df: pd.DataFrame,\n",
    "        embedding_column_1: str,\n",
    "        embedding_column_2: str,\n",
    "        similarity_label_column: str,\n",
    "    ) -> Tuple[torch.tensor]:\n",
    "        e1 = np.stack(np.array(df[embedding_column_1].values))\n",
    "        e2 = np.stack(np.array(df[embedding_column_2].values))\n",
    "        s = np.stack(np.array(df[similarity_label_column].astype(\"float\").values))\n",
    "\n",
    "        e1 = torch.from_numpy(e1).float()\n",
    "        e2 = torch.from_numpy(e2).float()\n",
    "        s = torch.from_numpy(s).float()\n",
    "\n",
    "        return e1, e2, s\n",
    "\n",
    "    e1_train, e2_train, s_train = tensors_from_dataframe(\n",
    "        df[df[\"dataset\"] == \"train\"], \"text_1_embedding\", \"text_2_embedding\", \"label\"\n",
    "    )\n",
    "    e1_test, e2_test, s_test = tensors_from_dataframe(\n",
    "        df[df[\"dataset\"] == \"test\"], \"text_1_embedding\", \"text_2_embedding\", \"label\"\n",
    "    )\n",
    "\n",
    "    # create dataset and loader\n",
    "    dataset = torch.utils.data.TensorDataset(e1_train, e2_train, s_train)\n",
    "    train_loader = torch.utils.data.DataLoader(\n",
    "        dataset, batch_size=batch_size, shuffle=True\n",
    "    )\n",
    "\n",
    "    # define model (similarity of projected embeddings)\n",
    "    def model(embedding_1, embedding_2, matrix, dropout_fraction=dropout_fraction):\n",
    "        e1 = torch.nn.functional.dropout(embedding_1, p=dropout_fraction)\n",
    "        e2 = torch.nn.functional.dropout(embedding_2, p=dropout_fraction)\n",
    "        modified_embedding_1 = e1 @ matrix  # @ is matrix multiplication\n",
    "        modified_embedding_2 = e2 @ matrix\n",
    "        similarity = torch.nn.functional.cosine_similarity(\n",
    "            modified_embedding_1, modified_embedding_2\n",
    "        )\n",
    "        return similarity\n",
    "\n",
    "    # define loss function to minimize\n",
    "    def mse_loss(predictions, targets):\n",
    "        difference = predictions - targets\n",
    "        return torch.sum(difference * difference) / difference.numel()\n",
    "\n",
    "    # initialize projection matrix\n",
    "    embedding_length = len(df[\"text_1_embedding\"].values[0])\n",
    "    matrix = torch.randn(\n",
    "        embedding_length, modified_embedding_length, requires_grad=True\n",
    "    )\n",
    "\n",
    "    epochs, types, losses, accuracies, matrices = [], [], [], [], []\n",
    "    for epoch in range(1, 1 + max_epochs):\n",
    "        # iterate through training dataloader\n",
    "        for a, b, actual_similarity in train_loader:\n",
    "            # generate prediction\n",
    "            predicted_similarity = model(a, b, matrix)\n",
    "            # get loss and perform backpropagation\n",
    "            loss = mse_loss(predicted_similarity, actual_similarity)\n",
    "            loss.backward()\n",
    "            # update the weights\n",
    "            with torch.no_grad():\n",
    "                matrix -= matrix.grad * learning_rate\n",
    "                # set gradients to zero\n",
    "                matrix.grad.zero_()\n",
    "        # calculate test loss\n",
    "        test_predictions = model(e1_test, e2_test, matrix)\n",
    "        test_loss = mse_loss(test_predictions, s_test)\n",
    "\n",
    "        # compute custom embeddings and new cosine similarities\n",
    "        apply_matrix_to_embeddings_dataframe(matrix, df)\n",
    "\n",
    "        # calculate test accuracy\n",
    "        for dataset in [\"train\", \"test\"]:\n",
    "            data = df[df[\"dataset\"] == dataset]\n",
    "            a, se = accuracy_and_se(data[\"cosine_similarity_custom\"], data[\"label\"])\n",
    "\n",
    "            # record results of each epoch\n",
    "            epochs.append(epoch)\n",
    "            types.append(dataset)\n",
    "            losses.append(loss.item() if dataset == \"train\" else test_loss.item())\n",
    "            accuracies.append(a)\n",
    "            matrices.append(matrix.detach().numpy())\n",
    "\n",
    "            # optionally print accuracies\n",
    "            if print_progress is True:\n",
    "                print(\n",
    "                    f\"Epoch {epoch}/{max_epochs}: {dataset} accuracy: {a:0.1%} ± {1.96 * se:0.1%}\"\n",
    "                )\n",
    "\n",
    "    data = pd.DataFrame(\n",
    "        {\"epoch\": epochs, \"type\": types, \"loss\": losses, \"accuracy\": accuracies}\n",
    "    )\n",
    "    data[\"run_id\"] = run_id\n",
    "    data[\"modified_embedding_length\"] = modified_embedding_length\n",
    "    data[\"batch_size\"] = batch_size\n",
    "    data[\"max_epochs\"] = max_epochs\n",
    "    data[\"learning_rate\"] = learning_rate\n",
    "    data[\"dropout_fraction\"] = dropout_fraction\n",
    "    data[\n",
    "        \"matrix\"\n",
    "    ] = matrices  # saving every single matrix can get big; feel free to delete/change\n",
    "    if save_results is True:\n",
    "        data.to_csv(f\"{run_id}_optimization_results.csv\", index=False)\n",
    "\n",
    "    return data\n",
    "\n",
    "\n",
    "@step\n",
    "def tune_matrix(big_df: pd.DataFrame) -> pd.DataFrame:\n",
    "    # example hyperparameter search\n",
    "    # I recommend starting with max_epochs=10 while initially exploring\n",
    "    results = []\n",
    "    max_epochs = 30\n",
    "    dropout_fraction = 0.2\n",
    "    for batch_size, learning_rate in [(10, 10), (100, 100), (1000, 1000)]:\n",
    "        result = optimize_matrix(\n",
    "            df=big_df,\n",
    "            batch_size=batch_size,\n",
    "            learning_rate=learning_rate,\n",
    "            max_epochs=max_epochs,\n",
    "            dropout_fraction=dropout_fraction,\n",
    "            save_results=False,\n",
    "        )\n",
    "        results.append(result)\n",
    "        \n",
    "    runs_df = pd.concat(results)\n",
    "    return runs_df\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@pipeline\n",
    "def load_and_process_data():\n",
    "    train_df, test_df, _ = download_data()\n",
    "    processed_train_df = process_input_data(train_df, dataset_type=\"train\")\n",
    "    processed_test_df = process_input_data(test_df, dataset_type=\"test\")\n",
    "\n",
    "    train_df_negatives = dataframe_of_negatives(processed_train_df, dataset_type=\"train\")\n",
    "    test_df_negatives = dataframe_of_negatives(processed_test_df, dataset_type=\"test\")\n",
    "\n",
    "    big_df = concatenate_dataframes(\n",
    "        train_df,\n",
    "        train_df_negatives, \n",
    "        test_df, \n",
    "        test_df_negatives, \n",
    "    )\n",
    "    \n",
    "    big_df_with_cos, embedding_cache = calculate_embeddings_and_cos_similarity(big_df)\n",
    "    \n",
    "    big_df_with_embeddings = apply_matrix_to_embeddings_dataframe(big_df_with_cos)\n",
    "    \n",
    "    tune_matrix(big_df_with_embeddings)\n",
    "    \n",
    "load_and_process_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load data\n",
    "from zenml.client import Client \n",
    "\n",
    "client = Client()\n",
    "\n",
    "latest_run = client.get_pipeline(\"load_and_process_data\").runs[-1]\n",
    "runs_df = latest_run.steps[\"tune_matrix\"].output.load()\n",
    "runs_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PoTZWC1SpgkS",
    "outputId": "207360e5-fd07-4180-a143-0ec5dd27ffe1"
   },
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "# plot training loss and test loss over time\n",
    "px.line(\n",
    "    runs_df,\n",
    "    line_group=\"run_id\",\n",
    "    x=\"epoch\",\n",
    "    y=\"loss\",\n",
    "    color=\"type\",\n",
    "    hover_data=[\"batch_size\", \"learning_rate\", \"dropout_fraction\"],\n",
    "    facet_row=\"learning_rate\",\n",
    "    facet_col=\"batch_size\",\n",
    "    width=500,\n",
    ").show()\n",
    "\n",
    "# plot accuracy over time\n",
    "px.line(\n",
    "    runs_df,\n",
    "    line_group=\"run_id\",\n",
    "    x=\"epoch\",\n",
    "    y=\"accuracy\",\n",
    "    color=\"type\",\n",
    "    hover_data=[\"batch_size\", \"learning_rate\", \"dropout_fraction\"],\n",
    "    facet_row=\"learning_rate\",\n",
    "    facet_col=\"batch_size\",\n",
    "    width=500,\n",
    ").show()\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "MiBQDcMPpgkS"
   },
   "source": [
    "## 8. Plot the before & after, showing the results of the best matrix found during training\n",
    "\n",
    "The better the matrix is, the more cleanly it will separate the similar and dissimilar pairs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "hzjoyLDOpgkS"
   },
   "outputs": [],
   "source": [
    "# apply result of best run to original data\n",
    "best_run = runs_df.sort_values(by=\"accuracy\", ascending=False).iloc[0]\n",
    "best_matrix = best_run[\"matrix\"]\n",
    "apply_matrix_to_embeddings_dataframe.entrypoint(best_matrix, df)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nLnvABnXpgkS",
    "outputId": "0c070faa-6e3e-4765-b082-565c72a609be"
   },
   "outputs": [],
   "source": [
    "# plot similarity distribution BEFORE customization\n",
    "px.histogram(\n",
    "    df,\n",
    "    x=\"cosine_similarity\",\n",
    "    color=\"label\",\n",
    "    barmode=\"overlay\",\n",
    "    width=500,\n",
    "    facet_row=\"dataset\",\n",
    ").show()\n",
    "\n",
    "test_df = df[df[\"dataset\"] == \"test\"]\n",
    "a, se = accuracy_and_se(test_df[\"cosine_similarity\"], test_df[\"label\"])\n",
    "print(f\"Test accuracy: {a:0.1%} ± {1.96 * se:0.1%}\")\n",
    "\n",
    "# plot similarity distribution AFTER customization\n",
    "px.histogram(\n",
    "    df,\n",
    "    x=\"cosine_similarity_custom\",\n",
    "    color=\"label\",\n",
    "    barmode=\"overlay\",\n",
    "    width=500,\n",
    "    facet_row=\"dataset\",\n",
    ").show()\n",
    "\n",
    "a, se = accuracy_and_se(test_df[\"cosine_similarity_custom\"], test_df[\"label\"])\n",
    "print(f\"Test accuracy after customization: {a:0.1%} ± {1.96 * se:0.1%}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XO7iqiVjpgkT",
    "outputId": "a100a9e0-d5aa-46ab-b8a7-4ec6f7bd1cec"
   },
   "outputs": [],
   "source": [
    "best_matrix  # this is what you can multiply your embeddings by\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "name": "customized_embeddings_example_with_synthetic_negatives.ipynb",
   "provenance": []
  },
  "interpreter": {
   "hash": "365536dcbde60510dc9073d6b991cd35db2d9bac356a11f5b64279a5e6708b97"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}