3-run-vector-search.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the open CLIP model and create function to generate embedding for a single text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import clip\n",
    "import torch\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "model, preprocess = clip.load(\"ViT-B/32\", device=device)\n",
    "\n",
    "def get_single_embedding(text):\n",
    "    with torch.no_grad():\n",
    "        # Encode the text to compute the feature vector and normalize it\n",
    "        text_input = clip.tokenize([text]).to(device)\n",
    "        text_features = model.encode_text(text_input)\n",
    "        text_features /= text_features.norm(dim=-1, keepdim=True)\n",
    "\n",
    "    # Return the feature vector\n",
    "    return text_features.cpu().numpy()[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a function to perform KNN search in OpenSearch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from dotenv import load_dotenv\n",
    "load_dotenv()\n",
    "SERVICE_URI = os.getenv(\"SERVICE_URI\")\n",
    "index_name = \"photos\"  # Update with your index name\n",
    "\n",
    "from opensearchpy import OpenSearch\n",
    "opensearch = OpenSearch(SERVICE_URI, use_ssl=True)\n",
    "\n",
    "def knn_search(text):\n",
    "    vector = get_single_embedding(text)\n",
    "\n",
    "    body = {\n",
    "        \"query\": {\n",
    "            \"knn\": {\n",
    "                \"embedding\": {\n",
    "                    \"vector\": vector.tolist(),  # Convert to list\n",
    "                    \"k\": 4  # Number of nearest neighbors to retrieve\n",
    "                }\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "\n",
    "    # Perform search\n",
    "    result = opensearch.search(index=index_name, body=body)\n",
    "    return result\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create function to display image as a result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import display, Image\n",
    "import pandas as pd\n",
    "\n",
    "# Read the TSV file into a DataFrame\n",
    "\n",
    "\n",
    "def display_images(result):\n",
    "    # Check if hits are present in the result\n",
    "    if 'hits' in result and 'hits' in result['hits']:\n",
    "        hits = result['hits']['hits']\n",
    "        \n",
    "        # Loop through each hit, up to a maximum of 4\n",
    "        for i, hit in enumerate(hits[:4]):\n",
    "            if '_source' in hit and 'image_url' in hit['_source']:\n",
    "                image_url = hit['_source']['image_url']\n",
    "                \n",
    "                # Display the image\n",
    "                print(f\"Displaying image {i+1}:\")\n",
    "                display(Image(url=image_url))\n",
    "            else:\n",
    "                print(f\"Hit {i+1} does not contain an 'image_url' key.\")\n",
    "                \n",
    "    else:\n",
    "        print(\"Invalid result format or no hits found.\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "text_input = \"room\"  # Provide your text input here\n",
    "result = knn_search(text_input)\n",
    "display_images(result)"
   ]
  }
 ],
 "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.12.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}