es-ingest-data.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from PIL import Image\n",
    "from transformers import AutoImageProcessor, AutoModel, AutoTokenizer\n",
    "import numpy as np\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "model = AutoModel.from_pretrained(\"openai/clip-vit-base-patch16\").to(device)\n",
    "processor = AutoImageProcessor.from_pretrained(\"openai/clip-vit-base-patch16\")\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"openai/clip-vit-base-patch16\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "dataset = pd.read_json(\"description.json\")\n",
    "dataset[\"new_name\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "\n",
    "dataset['text_embeddings'] = dataset.apply(\n",
    "    lambda x: model.get_text_features(\n",
    "        **tokenizer([x[\"description\"]], truncation=True, return_tensors=\"pt\").to(device)\n",
    "    )[0]\n",
    "    .detach()\n",
    "    .cpu()\n",
    "    .numpy(), \n",
    "    axis=1\n",
    ")\n",
    "dataset['image_embedding'] = dataset.apply(\n",
    "    lambda x: model.get_image_features(**processor([cv2.imread(x[\"new_name\"])], return_tensors=\"pt\").to(device))[0]\n",
    "        .detach()\n",
    "        .cpu()\n",
    "        .numpy(),\n",
    "        axis=1\n",
    ")\n",
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = dataset.apply(\n",
    "    lambda x: cv2.imread(x[\"new_name\"]), axis=1\n",
    ")\n",
    "\n",
    "a "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from elasticsearch import Elasticsearch\n",
    "\n",
    "# Connect to Elasticsearch\n",
    "es = Elasticsearch([{'host': 'es-test.aws.primehub.io', 'port': 9200, 'scheme': 'http'}],\n",
    "                   http_auth=('elastic', ''))\n",
    "\n",
    "# Define the index settings and mappings\n",
    "index_settings = {\n",
    "    \"settings\": {\n",
    "        \"number_of_shards\": 1,\n",
    "        \"number_of_replicas\": 1\n",
    "    },\n",
    "    \"mappings\": {\n",
    "        \"properties\": {\n",
    "            \"content\": {\n",
    "                \"type\": \"text\",\n",
    "                \"analyzer\": \"smartcn\"\n",
    "            },\n",
    "            \"content_vector\": {\n",
    "                \"type\": \"dense_vector\",\n",
    "                \"dims\": 512  # Replace with the actual number of dimensions\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "}\n",
    "index_name = 'image-text-search-v1'\n",
    "# Delete the index if it exists\n",
    "if es.indices.exists(index=index_name):\n",
    "    es.indices.delete(index=index_name)\n",
    "    print(f\"Deleted existing index: {index_name}\")\n",
    "else:\n",
    "    print(f\"Index {index_name} does not exist.\")\n",
    "\n",
    "\n",
    "es.indices.create(index=index_name, body=index_settings)\n",
    "print(f\"Created index: {index_name}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os \n",
    "\n",
    "folder_path = './'  # Replace with your folder path\n",
    "\n",
    "for index, row in dataset.iterrows():\n",
    "    # Construct the document\n",
    "    doc = {\n",
    "        \"content\": row[\"new_name\"],\n",
    "        \"content_vector\": row[\"image_embedding\"]\n",
    "    }\n",
    "\n",
    "    # # Index the document\n",
    "    es.index(index=index_name, body=doc)\n",
    "\n",
    "    print(f\"Indexed row {index+1}/{len(dataset)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query_vector = model.get_text_features(\n",
    "        **tokenizer([\"fish\"], truncation=True, return_tensors=\"pt\")).detach().cpu().numpy()\n",
    "search_query = {\n",
    "    \"query\": {\n",
    "        \"knn\": {\n",
    "            \"field\": \"content_vector\",\n",
    "            \"query_vector\": query_vector,\n",
    "            \"num_candidates\": 1\n",
    "        }\n",
    "    }\n",
    "}\n",
    "\n",
    "response = es.search(index=index_name, body=search_query)\n",
    "for hit in response['hits']['hits']:\n",
    "    print(hit['_source'])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dimon",
   "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.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}