{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b2000001",
   "metadata": {},
   "source": [
    "# Bulk Embedding Ingestion\n",
    "\n",
    "`ParallelLoader` sends batches of queries concurrently. In this example, a query generator supplies dish embeddings and metadata from the [Cookbook dataset](https://github.com/aperture-data/Cookbook). A sentence-transformer model converts dish descriptions to vectors before ingestion.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2000002",
   "metadata": {},
   "source": [
    "## Connect to ApertureDB\n",
    "\n",
    "**ApertureDB Cloud**\n",
    "The SDK connects with an API key from your [Cloud instance](https://cloud.aperturedata.io). Obtain the key from **Connect > Generate API Key** and add it to a `.env` file in this directory:\n",
    "```\n",
    "APERTUREDB_KEY=your_key_here\n",
    "```\n",
    "\n",
    "**Local evaluation:** follow the [deployment guide](https://docs.aperturedata.io/administration/deployment/DockerCompose), then configure your connection profile.\n",
    "\n",
    "\n",
    "See [client configuration options](https://docs.aperturedata.io/Setup/client/configuration) for all connection methods and [server setup options](https://docs.aperturedata.io/category/setup-server) for deployment choices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2000003",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install --quiet \"aperturedb==0.5.1\" \"python-dotenv==1.1.1\" sentence-transformers pandas\n",
    "\n",
    "# Encode literal data without interpreting a leading dollar sign as a property.\n",
    "def literal(value):\n",
    "    if isinstance(value, str) and value.startswith(\"$\"):\n",
    "        return \"$\" + value\n",
    "    if isinstance(value, (list, tuple)):\n",
    "        return {\"array\": [literal(item) for item in value]}\n",
    "    if isinstance(value, dict):\n",
    "        if len(value) == 1 and (\"date\" in value or \"_date\" in value):\n",
    "            return {\"to_date\": literal(next(iter(value.values())))}\n",
    "        return {\"json\": value}\n",
    "    return value\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2000004",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Option A: ApertureDB Cloud\n",
    "from dotenv import load_dotenv\n",
    "load_dotenv()  # loads APERTUREDB_KEY from .env into the environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2000005",
   "metadata": {},
   "outputs": [],
   "source": [
    "# For local evaluation, configure the profile from your versioned deployment package.\n",
    "# https://docs.aperturedata.io/administration/deployment/DockerCompose\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2000006",
   "metadata": {},
   "outputs": [],
   "source": [
    "from aperturedb.CommonLibrary import create_connector\n",
    "\n",
    "client = create_connector()\n",
    "response, _ = client.query([{\"GetStatus\": {}}])\n",
    "client.print_last_response()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2000007",
   "metadata": {},
   "source": [
    "## Load Dataset and Generate Embeddings\n",
    "\n",
    "Each input combines `dish_name` and `caption` into a description. `all-MiniLM-L6-v2` converts those descriptions to 384-dimensional embeddings and can run on a CPU.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2000008",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "from sentence_transformers import SentenceTransformer\n",
    "\n",
    "dishes = pd.read_csv(\n",
    "    \"https://raw.githubusercontent.com/aperture-data/Cookbook/51b3c3d9f495f8fd77108a270b88fe8ab8cc0004/images.adb.csv\"\n",
    ")\n",
    "dishes[\"description\"] = dishes[\"dish_name\"] + \" - \" + dishes[\"caption\"]\n",
    "print(f\"Loaded {len(dishes)} dishes\")\n",
    "\n",
    "model = SentenceTransformer(\"all-MiniLM-L6-v2\")\n",
    "embeddings = model.encode(dishes[\"description\"].tolist(), normalize_embeddings=True)\n",
    "print(f\"Embedding shape: {embeddings.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2000009",
   "metadata": {},
   "source": [
    "## Create the DescriptorSet\n",
    "\n",
    "The descriptor set defines the vector dimensions, index engine, and distance metric used for the dish embeddings.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2000010",
   "metadata": {},
   "outputs": [],
   "source": [
    "import uuid\n",
    "SET_NAME = \"cookbook_bulk_\" + uuid.uuid4().hex[:12]\n",
    "\n",
    "client.query([{\n",
    "    \"AddDescriptorSet\": {\n",
    "        \"name\":       SET_NAME,\n",
    "        \"dimensions\": 384,\n",
    "        \"engine\":     \"HNSW\",\n",
    "        \"metric\":     \"CS\",\n",
    "    }\n",
    "}])\n",
    "client.print_last_response()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2000011",
   "metadata": {},
   "source": [
    "## Bulk Ingest with ParallelLoader\n",
    "\n",
    "[ParallelLoader](https://docs.aperturedata.io/python_sdk/parallel_exec/ParallelLoader) consumes a `Subscriptable` that provides `(query, blobs)` pairs. The `DescriptorGenerator` below creates those pairs from dish metadata and precomputed embeddings. The loader processes batches of five records with four worker threads.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2000012",
   "metadata": {},
   "outputs": [],
   "source": [
    "from aperturedb.ParallelLoader import ParallelLoader\n",
    "\n",
    "class DescriptorGenerator:\n",
    "    \"\"\"Subscriptable generator of (query, blobs) pairs for ParallelLoader.\"\"\"\n",
    "\n",
    "    def __init__(self, dishes, embeddings, set_name):\n",
    "        self.dishes     = dishes\n",
    "        self.embeddings = embeddings\n",
    "        self.set_name   = set_name\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.dishes)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        # ParallelLoader calls __getitem__ with a slice for each batch;\n",
    "        # return a list of (query, blobs) pairs in that case.\n",
    "        if isinstance(idx, slice):\n",
    "            return [self[i] for i in range(*idx.indices(len(self)))]\n",
    "        row = self.dishes.iloc[idx]\n",
    "        emb = self.embeddings[idx].astype(\"float32\")\n",
    "        query = [{\n",
    "            \"AddDescriptor\": {\n",
    "                \"set\": self.set_name,\n",
    "                \"properties\": {\n",
    "                    \"dish_name\":   row[\"dish_name\"],\n",
    "                    \"cuisine\":     row[\"food_tags\"],\n",
    "                    \"caption\":     row[\"caption\"],\n",
    "                },\n",
    "                \"if_not_found\": ['$dish_name', '==', literal(row['dish_name'])],\n",
    "            }\n",
    "        }]\n",
    "        return query, [emb.tobytes()]\n",
    "\n",
    "generator = DescriptorGenerator(dishes, embeddings, SET_NAME)\n",
    "loader    = ParallelLoader(client)\n",
    "loader.ingest(generator, batchsize=5, numthreads=4, stats=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2000013",
   "metadata": {},
   "source": [
    "## Verify the Ingestion\n",
    "\n",
    "The count query reports the number of descriptors stored in the set.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2000014",
   "metadata": {},
   "outputs": [],
   "source": [
    "response, _ = client.query([{\n",
    "    \"FindDescriptor\": {\n",
    "        \"set\": SET_NAME,\n",
    "        \"results\":   {\"count\": True},\n",
    "    }\n",
    "}])\n",
    "client.print_last_response()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2000015",
   "metadata": {},
   "source": [
    "## Search the Bulk-Loaded Descriptors\n",
    "\n",
    "The query text is embedded with the same model as the dish descriptions. `FindDescriptor` returns up to three similar dishes with their metadata and distances.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2000016",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_text = \"creamy tomato curry\"\n",
    "query_emb  = model.encode([query_text], normalize_embeddings=True)[0].astype(\"float32\")\n",
    "\n",
    "response, _ = client.query([{\n",
    "    \"FindDescriptor\": {\n",
    "        \"set\":         SET_NAME,\n",
    "        \"k_neighbors\": 3,\n",
    "        \"distances\":   True,\n",
    "        \"results\":     {\"all_properties\": True},\n",
    "    }\n",
    "}], [query_emb.tobytes()])\n",
    "\n",
    "for entity in response[0][\"FindDescriptor\"].get(\"entities\", []):\n",
    "    score = 1 - entity[\"_distance\"]\n",
    "    print(f\"  {entity['dish_name']:<30} [{entity['cuisine']}]  score={score:.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2000017",
   "metadata": {},
   "source": [
    "## Cleanup\n",
    "\n",
    "Cleanup removes the descriptor set created for this run.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2000018",
   "metadata": {},
   "outputs": [],
   "source": [
    "client.query([{\"DeleteDescriptorSet\": {\"with_name\": SET_NAME}}])\n",
    "client.print_last_response()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2000019",
   "metadata": {},
   "source": [
    "## What's Next\n",
    "\n",
    "- [Hybrid Search](https://docs.aperturedata.io/HowToGuides/start/HybridSearch): combine KNN with metadata filters\n",
    "- [Recipe Text Search](https://docs.aperturedata.io/HowToGuides/start/TextSearch): single-item embedding flow with sentence-transformers\n",
    "- [Work with Descriptors](https://docs.aperturedata.io/HowToGuides/start/Embeddings): Add, Find, Update, Delete for individual descriptors\n",
    "- [Embeddings Extraction workflow](https://docs.aperturedata.io/workflows/embeddings_extraction): production-scale ingestion via the Workflows UI"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}