data_upload.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Processing and Vector Generation using FiftyOne and Pinecone"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook generates text embeddings for the COCO2017 dataset, using CLIP. These are both loaded from FiftyOne. Execute each cell, and create a config.py file in your directory, in which you should place your PINECONE_KEY. Make sure you have a Pinecone Index available to create a new one. \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/site-packages/pinecone/index.py:4: TqdmExperimentalWarning: Using `tqdm.autonotebook.tqdm` in notebook mode. Use `tqdm.tqdm` instead to force console mode (e.g. in jupyter console)\n",
      "  from tqdm.autonotebook import tqdm\n"
     ]
    }
   ],
   "source": [
    "import fiftyone.zoo as foz\n",
    "import pinecone\n",
    "import numpy as np\n",
    "from pkg_resources import packaging\n",
    "import torch\n",
    "\n",
    "from config import PINECONE_KEY"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading the data and the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading split 'train' to '/Users/laurendurivault/fiftyone/coco-2017/train' if necessary\n",
      "Found annotations at '/Users/laurendurivault/fiftyone/coco-2017/raw/instances_train2017.json'\n",
      "Downloading images to '/Users/laurendurivault/fiftyone/coco-2017/tmp-download/train2017.zip'\n",
      " 100% |████|  144.1Gb/144.1Gb [11.1m elapsed, 0s remaining, 228.3Mb/s]      \n",
      "Extracting images to '/Users/laurendurivault/fiftyone/coco-2017/train/data'\n",
      "Writing annotations to '/Users/laurendurivault/fiftyone/coco-2017/train/labels.json'\n",
      "Dataset info written to '/Users/laurendurivault/fiftyone/coco-2017/info.json'\n",
      "Loading 'coco-2017' split 'train'\n",
      " 100% |███████████| 118287/118287 [6.9m elapsed, 0s remaining, 340.3 samples/s]      \n",
      "Dataset 'coco-2017-train' created\n"
     ]
    }
   ],
   "source": [
    "dataset = foz.load_zoo_dataset(\"coco-2017\", split=\"train\")\n",
    "model = foz.load_zoo_model(\"clip-vit-base32-torch\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "if packaging.version.parse(\n",
    "  torch.__version__\n",
    ") < packaging.version.parse(\"1.8.0\"):\n",
    "  dtype = torch.long\n",
    "else:\n",
    "  dtype = torch.int"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 100% |███████████| 118287/118287 [1.5h elapsed, 0s remaining, 20.5 samples/s]      \n"
     ]
    }
   ],
   "source": [
    "# generating the embeddings\n",
    "dataset.compute_embeddings(\n",
    "    model, \n",
    "    embeddings_field=\"embedding\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# keep the data in my computer\n",
    "dataset.persistent = True"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initializing the pinecone index and upserting the vectors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize pinecone client\n",
    "pinecone.init(api_key=PINECONE_KEY, environment=\"us-east4-gcp\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create index with appropriate metric and name\n",
    "index_name = \"clip-image-search\"\n",
    "pinecone.create_index(\n",
    "    index_name, \n",
    "    dimension=512, \n",
    "    metric=\"cosine\", \n",
    "    pod_type=\"p1\"\n",
    ")\n",
    "# initialize index\n",
    "index = pinecone.Index(index_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "# convert numpy arrays to lists for pinecone\n",
    "embeddings = [arr.tolist() for arr in dataset.values(\"embedding\")]\n",
    "ids = [\"http://images.cocodataset.org/train2017/\" + file.split('/')[-1] for file in dataset.values(\"filepath\")]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create tuples of (id, embedding) for each sample\n",
    "index_vectors = list(zip(ids, embeddings))\n",
    "\n",
    "# upsert vectors in batches of 100\n",
    "def upsert_vectors(index, vectors):\n",
    "    num_vectors = len(vectors)\n",
    "    num_vectors_per_step = 100\n",
    "    num_steps = int(np.ceil(num_vectors/num_vectors_per_step))\n",
    "    for i in range(num_steps):\n",
    "        min_ind = num_vectors_per_step * i\n",
    "        max_ind = min(num_vectors_per_step * (i+1), num_vectors)\n",
    "        index.upsert(index_vectors[min_ind:max_ind])\n",
    "\n",
    "upsert_vectors(index, index_vectors)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Testing the image search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_text_embedding(prompt, clip_model):\n",
    "    tokenizer = clip_model._tokenizer\n",
    "\n",
    "    # standard start-of-text token\n",
    "    sot_token = tokenizer.encoder[\"<|startoftext|>\"]\n",
    "\n",
    "    # standard end-of-text token\n",
    "    eot_token = tokenizer.encoder[\"<|endoftext|>\"]\n",
    "\n",
    "    prompt_tokens = tokenizer.encode(prompt)\n",
    "    all_tokens = [[sot_token] + prompt_tokens + [eot_token]]\n",
    "\n",
    "    text_features = torch.zeros(\n",
    "        len(all_tokens),\n",
    "        clip_model.config.context_length,\n",
    "        dtype=dtype,\n",
    "        device=device,\n",
    "    )\n",
    "\n",
    "    # insert tokens into feature vector\n",
    "    text_features[0, : len(all_tokens[0])] = torch.tensor(all_tokens)\n",
    "\n",
    "    # encode text\n",
    "    embedding = clip_model._model.encode_text(text_features).to(device)\n",
    "\n",
    "    # convert to list for Pinecone\n",
    "    return embedding.tolist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'matches': [{'id': 'http://images.cocodataset.org/train2017/000000506187.jpg',\n",
       "              'score': 0.268776685,\n",
       "              'values': []},\n",
       "             {'id': 'http://images.cocodataset.org/train2017/000000546963.jpg',\n",
       "              'score': 0.267506778,\n",
       "              'values': []},\n",
       "             {'id': 'http://images.cocodataset.org/train2017/000000110798.jpg',\n",
       "              'score': 0.267011523,\n",
       "              'values': []},\n",
       "             {'id': 'http://images.cocodataset.org/train2017/000000048665.jpg',\n",
       "              'score': 0.258640438,\n",
       "              'values': []},\n",
       "             {'id': 'http://images.cocodataset.org/train2017/000000423327.jpg',\n",
       "              'score': 0.257749081,\n",
       "              'values': []},\n",
       "             {'id': 'http://images.cocodataset.org/train2017/000000445351.jpg',\n",
       "              'score': 0.249218613,\n",
       "              'values': []},\n",
       "             {'id': 'http://images.cocodataset.org/train2017/000000362138.jpg',\n",
       "              'score': 0.248322085,\n",
       "              'values': []},\n",
       "             {'id': 'http://images.cocodataset.org/train2017/000000478724.jpg',\n",
       "              'score': 0.248261213,\n",
       "              'values': []},\n",
       "             {'id': 'http://images.cocodataset.org/train2017/000000012228.jpg',\n",
       "              'score': 0.248198986,\n",
       "              'values': []},\n",
       "             {'id': 'http://images.cocodataset.org/train2017/000000260547.jpg',\n",
       "              'score': 0.247932509,\n",
       "              'values': []}],\n",
       " 'namespace': ''}"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prompt = \"a smile\"\n",
    "query_vector = get_text_embedding(prompt, model)\n",
    "top_k_samples = index.query(\n",
    "    vector=query_vector,\n",
    "    top_k=10,\n",
    "    include_values=False\n",
    ")\n",
    "\n",
    "top_k_samples"
   ]
  }
 ],
 "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.10.5"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}