opensearch_ingest.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "220ba1c4-a628-4c81-afd8-50ae27f8ab56",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Creating a vectorstore with Amazon Bedrock multimodal-embeddings\n",
    "\n",
    "This notebook gives a step-by-step tutorial to populate a vector database in [Opensearch Serverless](https://aws.amazon.com/opensearch-service/features/serverless/). These vector embeddings will be used by the Bedrock Agent to search for similar images in the provided vectorstore.\n",
    "\n",
    "This notebook is required if you would like to the agent to be able to take the `/image_look_up` action, otherwise you can directly run the `Create_Fashion_Agent.ipynb` notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "193be29c-1a8c-422e-b3e6-ae1f9abccd53",
   "metadata": {},
   "source": [
    "#### Download the dataset locally\n",
    "\n",
    "For demo purposes, we will be using an external [dataset](https://github.com/orbitalsonic/Fashion-Dataset-Images-Eastern-Dress) from GitHub."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e02d7c71-383c-431e-a2eb-bc83f0480297",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!git clone https://github.com/orbitalsonic/Fashion-Dataset-Images-Western-Dress.git"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b60b8a5-0c89-48d0-a452-9653f315b250",
   "metadata": {},
   "source": [
    "### Add all the dependencies/imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c22c7e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7944bd29-ddae-4b12-965d-9e7499b7ae5e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import boto3\n",
    "from opensearchpy import AWSV4SignerAuth, OpenSearch, RequestsHttpConnection\n",
    "from pathlib import Path\n",
    "import yaml\n",
    "import json\n",
    "from tqdm.auto import tqdm\n",
    "from pathlib import Path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d1a65168",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join(Path(os.path.abspath(\"\")), \"config.yml\"), \"r\") as ymlfile:\n",
    "    config = yaml.load(ymlfile, Loader=yaml.SafeLoader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "105aad0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"variables.json\", \"r\") as f:\n",
    "    variables = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "403b640c-c641-46d1-b9ce-a64f4337566f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "boto3_session = boto3.Session(profile_name=\"alexhrn-Admin\", region_name=\"us-east-1\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "486b8d87",
   "metadata": {},
   "outputs": [],
   "source": [
    "import io\n",
    "import json\n",
    "import boto3\n",
    "import base64\n",
    "from PIL import Image\n",
    "\n",
    "\n",
    "# Define output vector size – 1,024 (default), 384, 256\n",
    "\n",
    "EMBEDDING_CONFIG = {\n",
    "    \"embeddingConfig\": {\"outputEmbeddingLength\": int(config[\"embeddingSize\"])}\n",
    "}\n",
    "\n",
    "\n",
    "class OpensearchIngestion:\n",
    "    def __init__(self, client, session=None):\n",
    "        self.client = client\n",
    "        self.session = session if session else boto3.Session()\n",
    "        self.region = self.session.region_name\n",
    "\n",
    "    def put_bulk_in_opensearch(self, docs):\n",
    "        print(f\"Putting {len(docs)} documents in OpenSearch\")\n",
    "        success, failed = self.client.bulk(docs)\n",
    "        return success, failed\n",
    "\n",
    "    def check_index_exists(self, index_name):\n",
    "        return self.client.indices.exists(index=index_name)\n",
    "\n",
    "    def create_index(self, index_name):\n",
    "        if not self.check_index_exists(index_name):\n",
    "            settings = {\n",
    "                \"settings\": {\n",
    "                    \"index.knn\": True,\n",
    "                }\n",
    "            }\n",
    "            response = self.client.indices.create(index=index_name, body=settings)\n",
    "            return bool(response[\"acknowledged\"])\n",
    "        return False\n",
    "\n",
    "    def create_index_mapping(self, index_name):\n",
    "        response = self.client.indices.put_mapping(\n",
    "            index=index_name,\n",
    "            body={\n",
    "                \"properties\": {\n",
    "                    \"vector_field\": {\n",
    "                        \"type\": \"knn_vector\",\n",
    "                        \"dimension\": config[\"embeddingSize\"],\n",
    "                        \"method\": {\n",
    "                            \"name\": \"hnsw\",\n",
    "                            \"engine\": \"nmslib\",\n",
    "                        },\n",
    "                    },\n",
    "                    \"image_b64\": {\"type\": \"text\"},\n",
    "                }\n",
    "            },\n",
    "        )\n",
    "        return bool(response[\"acknowledged\"])\n",
    "\n",
    "    def get_bedrock_client(self):\n",
    "        return self.session.client(\"bedrock-runtime\", region_name=self.region)\n",
    "\n",
    "    def create_titan_multimodal_embeddings(\n",
    "        self,\n",
    "        image_path: str = \"None\",\n",
    "        text: str = \"None\",\n",
    "    ):\n",
    "        \"\"\"Creates the titan embeddings from the provided image and/or text.\"\"\"\n",
    "        payload_body = {}\n",
    "\n",
    "        if image_path and image_path != \"None\":\n",
    "            payload_body[\"inputImage\"] = self.get_encoded_image(image_path)\n",
    "        if text and (text != \"None\"):\n",
    "            payload_body[\"inputText\"] = text\n",
    "        if (image_path == \"None\") and (text == \"None\"):\n",
    "            raise \"please provide either an image and/or a text description\"\n",
    "\n",
    "        bedrock_client = self.get_bedrock_client()\n",
    "\n",
    "        response = bedrock_client.invoke_model(\n",
    "            body=json.dumps({**payload_body, **EMBEDDING_CONFIG}),\n",
    "            modelId=\"amazon.titan-embed-image-v1\",\n",
    "            accept=\"application/json\",\n",
    "            contentType=\"application/json\",\n",
    "        )\n",
    "        vector = json.loads(response[\"body\"].read())\n",
    "        return (payload_body, vector)\n",
    "\n",
    "    def get_encoded_image(self, image_path: str):\n",
    "        max_height, max_width = 1024, 1024  # Conservative Limit. Can increase to 2048\n",
    "        # Open the image and compress it if greater than the defined max size.\n",
    "        with Image.open(image_path) as image:\n",
    "            if (image.size[0] * image.size[1]) > (max_height * max_width):\n",
    "                image.thumbnail((max_height, max_width))\n",
    "                resized_img = image.copy()\n",
    "            else:\n",
    "                resized_img = image\n",
    "            img_byte_array = io.BytesIO()\n",
    "            resized_img.save(img_byte_array, format=image.format)\n",
    "            img_bytes = img_byte_array.getvalue()\n",
    "\n",
    "        # Encode the image to base64\n",
    "        image_encoded = base64.b64encode(img_bytes).decode(\"utf8\")\n",
    "        return image_encoded"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b29ad34d-43ae-4230-9062-1a699134e11c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# create a client for OSS\n",
    "client = boto3_session.client(\"opensearchserverless\")\n",
    "service = \"aoss\"\n",
    "region = boto3_session.region_name\n",
    "credentials = boto3_session.get_credentials()\n",
    "AWSAUTH = AWSV4SignerAuth(credentials, region, \"aoss\")\n",
    "for key in variables[\"FashionAgentStack\"].keys():\n",
    "    if key.startswith(\"OpenSearchServerlessConstructsFashionAgentStackOSSEndpoint\"):\n",
    "        host = variables[\"FashionAgentStack\"][key].removeprefix(\"https://\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ce0c3d4-885b-44ee-94af-e35c4761445e",
   "metadata": {},
   "source": [
    "#### Initialize an Opensearch client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "bcec490b-5724-4335-85e0-75e87e7fc84a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Create the client with SSL/TLS enabled.\n",
    "OSSclient = OpenSearch(\n",
    "    hosts=[{\"host\": host, \"port\": 443}],\n",
    "    http_auth=AWSAUTH,\n",
    "    use_ssl=True,\n",
    "    verify_certs=True,\n",
    "    connection_class=RequestsHttpConnection,\n",
    "    pool_maxsize=20,\n",
    "    timeout=3000,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73b97446-cec3-4166-abd6-db40e75936a0",
   "metadata": {},
   "source": [
    "### Create an index for the Opensearch ingestion\n",
    "Opensearch Ingestion class (created in opensearch_utils.py) contains helper functions for the document processing and ingestion into the index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "6fbf1e17-81cf-4734-b0b7-aedcc03169b4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "oss_instance = OpensearchIngestion(client=OSSclient, session=boto3_session)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77689c83-0c8e-4825-9103-1b121b273c30",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "oss_instance.create_index(config[\"opensearch\"][\"opensearch_index_name\"])\n",
    "oss_instance.create_index_mapping(config[\"opensearch\"][\"opensearch_index_name\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd340ec2-b945-4a64-b949-7c034b57c33b",
   "metadata": {},
   "source": [
    "### Ingest the images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "eaee504f-934f-40b8-9295-c72203dfadb5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "dataset_path = Path(\"Fashion-Dataset-Images-Western-Dress/WesternDress_Images\")\n",
    "image_count = sum(\n",
    "    1\n",
    "    for item in dataset_path.iterdir()\n",
    "    if item.is_file() and not item.name.startswith(\".\")\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e97f1da8-c511-4ad6-b112-f8bb8c231524",
   "metadata": {
    "jupyter": {
     "outputs_hidden": true
    },
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "failed = []\n",
    "for image_path in tqdm(dataset_path.iterdir(), total=image_count):\n",
    "    try:\n",
    "        (data, embedding) = oss_instance.create_titan_multimodal_embeddings(\n",
    "            image_path=image_path\n",
    "        )\n",
    "        img_id = str(image_path).rsplit(\"/\", 1)[1].split(\".\")[0]\n",
    "        body = {\n",
    "            \"vector_field\": embedding[\"embedding\"],\n",
    "            \"image_b64\": data[\"inputImage\"],\n",
    "        }\n",
    "    except Exception as e:\n",
    "        print(f\"Exception thrown in image {image_path}: {e}\")\n",
    "        continue\n",
    "    # Ingest the images one by one.\n",
    "    status = oss_instance.client.index(\n",
    "        index=config[\"opensearch\"][\"opensearch_index_name\"],\n",
    "        body=body,\n",
    "    )\n",
    "    if status[\"result\"] != \"created\":\n",
    "        failed.append(image_path)\n",
    "\n",
    "print(f\"Ingestion Complete. Failed ingestion for the following: {failed}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26a48958",
   "metadata": {},
   "source": [
    "##### Clean up will be done together by destroying the CDK\n",
    "\n",
    "run the `cdk destroy` command in the Terminal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abb4e5aa",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}