{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Get embeddings from dataset\n", "\n", "This notebook gives an example on how to get embeddings from a large dataset.\n", "\n", "\n", "## 1. Load the dataset\n", "\n", "The dataset used in this example is [fine-food reviews](https://www.kaggle.com/snap/amazon-fine-food-reviews) from Amazon. The dataset contains a total of 568,454 food reviews Amazon users left up to October 2012. We will use a subset of this dataset, consisting of 1,000 most recent reviews for illustration purposes. The reviews are in English and tend to be positive or negative. Each review has a ProductId, UserId, Score, review title (Summary) and review body (Text).\n", "\n", "We will combine the review summary and review text into a single combined text. The model will encode this combined text and it will output a single vector embedding." ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "To run this notebook, you will need to install: pandas, openai, transformers, plotly, matplotlib, scikit-learn, torch (transformer dep), torchvision, and scipy." ] }, { "cell_type": "code", "execution_count": 6, "metadata": {}, "outputs": [], "source": [ "# imports\n", "import pandas as pd\n", "import tiktoken\n", "\n", "from utils.embeddings_utils import get_embedding\n" ] }, { "cell_type": "code", "execution_count": 7, "metadata": {}, "outputs": [], "source": [ "# embedding model parameters\n", "embedding_model = \"text-embedding-ada-002\"\n", "embedding_encoding = \"cl100k_base\" # this the encoding for text-embedding-ada-002\n", "max_tokens = 8000 # the maximum for text-embedding-ada-002 is 8191\n" ] }, { "cell_type": "code", "execution_count": 8, "metadata": {}, "outputs": [ { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
TimeProductIdUserIdScoreSummaryTextcombined
01351123200B003XPF9BOA3R7JR3FMEBXQB5where does one start...and stop... with a tre...Wanted to save some to bring to my Chicago fam...Title: where does one start...and stop... wit...
11351123200B003JK537SA3JBPC3WFUT5ZP1Arrived in piecesNot pleased at all. When I opened the box, mos...Title: Arrived in pieces; Content: Not pleased...
\n", "
" ], "text/plain": [ " Time ProductId UserId Score \\\n", "0 1351123200 B003XPF9BO A3R7JR3FMEBXQB 5 \n", "1 1351123200 B003JK537S A3JBPC3WFUT5ZP 1 \n", "\n", " Summary \\\n", "0 where does one start...and stop... with a tre... \n", "1 Arrived in pieces \n", "\n", " Text \\\n", "0 Wanted to save some to bring to my Chicago fam... \n", "1 Not pleased at all. When I opened the box, mos... \n", "\n", " combined \n", "0 Title: where does one start...and stop... wit... \n", "1 Title: Arrived in pieces; Content: Not pleased... " ] }, "execution_count": 8, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# load & inspect dataset\n", "input_datapath = \"data/fine_food_reviews_1k.csv\" # to save space, we provide a pre-filtered dataset\n", "df = pd.read_csv(input_datapath, index_col=0)\n", "df = df[[\"Time\", \"ProductId\", \"UserId\", \"Score\", \"Summary\", \"Text\"]]\n", "df = df.dropna()\n", "df[\"combined\"] = (\n", " \"Title: \" + df.Summary.str.strip() + \"; Content: \" + df.Text.str.strip()\n", ")\n", "df.head(2)\n" ] }, { "cell_type": "code", "execution_count": 9, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "1000" ] }, "execution_count": 9, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# subsample to 1k most recent reviews and remove samples that are too long\n", "top_n = 1000\n", "df = df.sort_values(\"Time\").tail(top_n * 2) # first cut to first 2k entries, assuming less than half will be filtered out\n", "df.drop(\"Time\", axis=1, inplace=True)\n", "\n", "encoding = tiktoken.get_encoding(embedding_encoding)\n", "\n", "# omit reviews that are too long to embed\n", "df[\"n_tokens\"] = df.combined.apply(lambda x: len(encoding.encode(x)))\n", "df = df[df.n_tokens <= max_tokens].tail(top_n)\n", "len(df)\n" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "## 2. Get embeddings and save them for future reuse" ] }, { "cell_type": "code", "execution_count": 10, "metadata": {}, "outputs": [], "source": [ "# Ensure you have your API key set in your environment per the README: https://github.com/openai/openai-python#usage\n", "\n", "# This may take a few minutes\n", "df[\"embedding\"] = df.combined.apply(lambda x: get_embedding(x, engine=embedding_model))\n", "df.to_csv(\"data/fine_food_reviews_with_embeddings_1k.csv\")\n" ] } ], "metadata": { "kernelspec": { "display_name": "openai", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.9.16" }, "orig_nbformat": 4, "vscode": { "interpreter": { "hash": "365536dcbde60510dc9073d6b991cd35db2d9bac356a11f5b64279a5e6708b97" } } }, "nbformat": 4, "nbformat_minor": 2 }