SemanticSearchRedis.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install redis numpy pandas\n",
    "%pip install -U sentence-transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import time\n",
    "from redis import Redis\n",
    "from redis.commands.search.field import VectorField\n",
    "from redis.commands.search.field import TextField\n",
    "from redis.commands.search.field import TagField\n",
    "from redis.commands.search.query import Query\n",
    "from redis.commands.search.result import Result\n",
    "\n",
    "class color:\n",
    "   PURPLE = '\\033[95m'\n",
    "   CYAN = '\\033[96m'\n",
    "   DARKCYAN = '\\033[36m'\n",
    "   BLUE = '\\033[94m'\n",
    "   GREEN = '\\033[92m'\n",
    "   YELLOW = '\\033[93m'\n",
    "   RED = '\\033[91m'\n",
    "   BOLD = '\\033[1m'\n",
    "   UNDERLINE = '\\033[4m'\n",
    "   END = '\\033[0m'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MAX_TEXT_LENGTH=512\n",
    "NUMBER_PRODUCTS=1000\n",
    "\n",
    "def auto_truncate(val):\n",
    "    return val[:MAX_TEXT_LENGTH]\n",
    "\n",
    "#Load Product data and truncate long text fields\n",
    "all_prods_df = pd.read_csv(\"data/product_data.csv\", converters={'bullet_point': auto_truncate,'item_keywords':auto_truncate,'item_name':auto_truncate})\n",
    "all_prods_df['primary_key'] = all_prods_df['item_id'] + '-' + all_prods_df['domain_name']\n",
    "all_prods_df['item_keywords'].replace('', np.nan, inplace=True)\n",
    "all_prods_df.dropna(subset=['item_keywords'], inplace=True)\n",
    "all_prods_df.reset_index(drop=True,inplace=True)\n",
    "\n",
    "#get the first 1000 products with non-empty item keywords\n",
    "product_metadata = all_prods_df.head(NUMBER_PRODUCTS).to_dict(orient='index')\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_prods_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "redis_conn = Redis(\n",
    "  host='XXXX',\n",
    "  port=12345,\n",
    "  password=\"XXXXXXXX\")\n",
    "print ('Connected to redis')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sentence_transformers import SentenceTransformer\n",
    "model = SentenceTransformer('sentence-transformers/all-distilroberta-v1')\n",
    "\n",
    "%%time\n",
    "\n",
    "item_keywords =  [product_metadata[i]['item_keywords']  for i in product_metadata.keys()]\n",
    "item_keywords_vectors = [ model.encode(sentence) for sentence in item_keywords]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(item_keywords_vectors)\n",
    "len(product_metadata)\n",
    "# Check one of the products\n",
    "product_metadata[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_vectors(client:Redis, product_metadata, vector_dict, vector_field_name):\n",
    "    p = client.pipeline(transaction=False)\n",
    "    for index in product_metadata.keys():    \n",
    "        #hash key\n",
    "        key='product:'+ str(index)+ ':' + product_metadata[index]['primary_key']\n",
    "        \n",
    "        #hash values\n",
    "        item_metadata = product_metadata[index]\n",
    "        item_keywords_vector = vector_dict[index].astype(np.float32).tobytes()\n",
    "        item_metadata[vector_field_name]=item_keywords_vector\n",
    "        \n",
    "        # HSET\n",
    "        p.hset(key,mapping=item_metadata)\n",
    "            \n",
    "    p.execute()\n",
    "\n",
    "def create_flat_index (redis_conn,vector_field_name,number_of_vectors, vector_dimensions=512, distance_metric='L2'):\n",
    "    redis_conn.ft().create_index([\n",
    "        VectorField(vector_field_name, \"FLAT\", {\"TYPE\": \"FLOAT32\", \"DIM\": vector_dimensions, \"DISTANCE_METRIC\": distance_metric, \"INITIAL_CAP\": number_of_vectors, \"BLOCK_SIZE\":number_of_vectors }),\n",
    "        TagField(\"product_type\"),\n",
    "        TextField(\"item_name\"),\n",
    "        TextField(\"item_keywords\"),\n",
    "        TagField(\"country\")        \n",
    "    ])\n",
    "\n",
    "def create_hnsw_index (redis_conn,vector_field_name,number_of_vectors, vector_dimensions=512, distance_metric='L2',M=40,EF=200):\n",
    "    redis_conn.ft().create_index([\n",
    "        VectorField(vector_field_name, \"HNSW\", {\"TYPE\": \"FLOAT32\", \"DIM\": vector_dimensions, \"DISTANCE_METRIC\": distance_metric, \"INITIAL_CAP\": number_of_vectors, \"M\": M, \"EF_CONSTRUCTION\": EF}),\n",
    "        TagField(\"product_type\"),\n",
    "        TextField(\"item_keywords\"),        \n",
    "        TextField(\"item_name\"),\n",
    "        TagField(\"country\")     \n",
    "    ])    "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FLAT index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "ITEM_KEYWORD_EMBEDDING_FIELD='item_keyword_vector'\n",
    "TEXT_EMBEDDING_DIMENSION=768\n",
    "NUMBER_PRODUCTS=1000\n",
    "\n",
    "print ('Loading and Indexing + ' +  str(NUMBER_PRODUCTS) + ' products')\n",
    "\n",
    "#flush all data\n",
    "redis_conn.flushall()\n",
    "\n",
    "#create flat index & load vectors\n",
    "create_flat_index(redis_conn, ITEM_KEYWORD_EMBEDDING_FIELD,NUMBER_PRODUCTS,TEXT_EMBEDDING_DIMENSION,'COSINE')\n",
    "load_vectors(redis_conn,product_metadata,item_keywords_vectors,ITEM_KEYWORD_EMBEDDING_FIELD)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "topK=5\n",
    "product_query='beautifully crafted present for her. a special occasion'\n",
    "#product_query='cool way to pimp up my cell'\n",
    "\n",
    "#vectorize the query\n",
    "query_vector = model.encode(product_query).astype(np.float32).tobytes()\n",
    "\n",
    "#prepare the query\n",
    "q = Query(f'*=>[KNN {topK} @{ITEM_KEYWORD_EMBEDDING_FIELD} $vec_param AS vector_score]').sort_by('vector_score').paging(0,topK).return_fields('vector_score','item_name','item_id','item_keywords').dialect(2)\n",
    "params_dict = {\"vec_param\": query_vector}\n",
    "\n",
    "\n",
    "#Execute the query\n",
    "results = redis_conn.ft().search(q, query_params = params_dict)\n",
    "\n",
    "#Print similar products found\n",
    "for product in results.docs:\n",
    "    print ('***************Product  found ************')\n",
    "    print (color.BOLD + 'hash key = ' +  color.END + product.id)\n",
    "    print (color.YELLOW + 'Item Name = ' +  color.END  + product.item_name)\n",
    "    print (color.YELLOW + 'Item Id = ' +  color.END  + product.item_id)\n",
    "    print (color.YELLOW + 'Item keywords = ' +  color.END  + product.item_keywords)\n",
    "    print (color.YELLOW + 'Score = ' +  color.END  + product.vector_score)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# HNSW"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "print ('Loading and Indexing + ' +  str(NUMBER_PRODUCTS) + ' products')\n",
    "\n",
    "ITEM_KEYWORD_EMBEDDING_FIELD='item_keyword_vector'\n",
    "NUMBER_PRODUCTS=1000\n",
    "TEXT_EMBEDDING_DIMENSION=768\n",
    "\n",
    "#flush all data\n",
    "redis_conn.flushall()\n",
    "\n",
    "#create flat index & load vectors\n",
    "create_hnsw_index(redis_conn, ITEM_KEYWORD_EMBEDDING_FIELD,NUMBER_PRODUCTS,TEXT_EMBEDDING_DIMENSION,'COSINE',M=40,EF=200)\n",
    "load_vectors(redis_conn,product_metadata,item_keywords_vectors,ITEM_KEYWORD_EMBEDDING_FIELD)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "topK=5\n",
    "product_query='beautifully crafted present for her. a special occasion'\n",
    "#product_query='cool way to pimp up my cell'\n",
    "\n",
    "#vectorize the query\n",
    "query_vector = model.encode(product_query).astype(np.float32).tobytes()\n",
    "\n",
    "#prepare the query\n",
    "q = Query(f'*=>[KNN {topK} @{ITEM_KEYWORD_EMBEDDING_FIELD} $vec_param AS vector_score]').sort_by('vector_score').paging(0,topK).return_fields('vector_score','item_name','item_id','item_keywords','country').dialect(2)\n",
    "params_dict = {\"vec_param\": query_vector}\n",
    "\n",
    "#Execute the query\n",
    "results = redis_conn.ft().search(q, query_params = params_dict)\n",
    "\n",
    "#Print similar products found\n",
    "for product in results.docs:\n",
    "    print ('***************Product  found ************')\n",
    "    print (color.BOLD + 'hash key = ' +  color.END + product.id)\n",
    "    print (color.YELLOW + 'Item Name = ' +  color.END  + product.item_name)\n",
    "    print (color.YELLOW + 'Item Id = ' +  color.END  + product.item_id)\n",
    "    print (color.YELLOW + 'Item keywords = ' +  color.END  + product.item_keywords)\n",
    "    print (color.YELLOW + 'Country = ' +  color.END  + product.country)\n",
    "    print (color.YELLOW + 'Score = ' +  color.END  + product.vector_score)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We highly recommend checking out [this repo](https://github.com/RedisAI/vecsim-demo) for more information and a guide on how to do a similarity search like we did here with images."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11.1"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "5ffd7eb2cebf9ac436b5021ba01877e9cee6b03524e01bf8c8637d3e64111215"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}