use halfvec

codelibs · May 3, 2024 · 64b9d8d · 64b9d8d
1 parent 780be11
commit 64b9d8d
Showing 1 changed file with 17 additions and 5 deletions.
diff --git a/run-pgvector.ipynb b/run-pgvector.ipynb
@@ -59,6 +59,7 @@
     "    hnsw_ef_construction: int\n",
     "    hnsw_ef: int\n",
     "    update_docs_per_sec: int\n",
+    "    quantization: str\n",
     "\n",
     "    pgvector_name: str = \"benchmark_pgvector\"\n",
     "    pgvector_host: str = \"localhost\"\n",
@@ -87,6 +88,7 @@
     "            \"hnsw_ef_construction\": 200,\n",
     "            \"hnsw_ef\": 100,\n",
     "            \"update_docs_per_sec\": 0,\n",
+    "            \"quantization\": \"halfvec\", # \"vector\",\n",
     "        },\n",
     "        \"1m-768-m49-ef100-ip\": {\n",
     "            \"content_path\": Path(\"dataset/passages-c400-jawiki-20230403\"),\n",
@@ -101,6 +103,7 @@
     "            \"hnsw_ef_construction\": 200,\n",
     "            \"hnsw_ef\": 100,\n",
     "            \"update_docs_per_sec\": 0,\n",
+    "            \"quantization\": \"halfvec\", # \"vector\",\n",
     "        },\n",
     "        \"5m-768-m49-ef100-ip\": {\n",
     "            \"content_path\": Path(\"dataset/passages-c400-jawiki-20230403\"),\n",
@@ -115,6 +118,7 @@
     "            \"hnsw_ef_construction\": 200,\n",
     "            \"hnsw_ef\": 100,\n",
     "            \"update_docs_per_sec\": 0,\n",
+    "            \"quantization\": \"halfvec\", # \"vector\",\n",
     "        },\n",
     "    }\n",
     "    return DataSetConfig(**setting.get(target_name))\n"
@@ -278,10 +282,14 @@
    "outputs": [],
    "source": [
     "def create_index(config):\n",
-    "    print(F\"Creating {config.index_name}... \", end=\"\")\n",
-    "    vector_ops = \"vector_ip_ops\" if config.distance == \"<#>\" else \"vector_cosine_ops\"\n",
+    "    print(F\"Creating {config.index_name} with {config.quantization}... \", end=\"\")\n",
+    "    vector_ops = f\"{config.quantization}_ip_ops\" if config.distance == \"<#>\" else f\"{config.quantization}_cosine_ops\"\n",
     "    with psycopg.connect(f\"dbname={config.pgvector_dbname} {config.pgvector_conninfo}\") as conn:\n",
     "        conn.execute(\"CREATE EXTENSION IF NOT EXISTS vector\")\n",
+    "        if config.quantization == \"halfvec\":\n",
+    "            create_index = f\"CREATE INDEX ON {config.index_name} USING hnsw ((embedding::{config.quantization}({config.dimension})) {vector_ops}) WITH (m = {config.hnsw_m}, ef_construction = {config.hnsw_ef_construction});\"\n",
+    "        else:\n",
+    "            create_index = f\"CREATE INDEX ON {config.index_name} USING hnsw (embedding {vector_ops}) WITH (m = {config.hnsw_m}, ef_construction = {config.hnsw_ef_construction});\"\n",
     "        conn.execute(f\"\"\"\n",
     "            CREATE TABLE {config.index_name} (\n",
     "                doc_id integer PRIMARY KEY,\n",
@@ -290,9 +298,9 @@
     "                /* title character(200), */\n",
     "                section character(128),\n",
     "                /* text text, */\n",
-    "                embedding vector({config.dimension})\n",
+    "                embedding {config.quantization}({config.dimension})\n",
     "            );\n",
-    "            CREATE INDEX ON {config.index_name} USING hnsw (embedding {vector_ops}) WITH (m = {config.hnsw_m}, ef_construction = {config.hnsw_ef_construction});\n",
+    "            {create_index}\n",
     "        \"\"\")\n",
     "        print(\" [OK]\")\n",
     "        return\n"
@@ -521,10 +529,14 @@
     "def search(config, embedding, page_size, pre_filter):\n",
     "    # print(query)\n",
     "    where = f\"WHERE {next(pre_filter)}\" if pre_filter is not None else \"\"\n",
+    "    if config.quantization == \"halfvec\":\n",
+    "        column_name = f\"embedding::{config.quantization}({config.dimension})\"\n",
+    "    else:\n",
+    "        column_name = \"embedding\"\n",
     "    query = f\"\"\"\n",
     "        SELECT\n",
     "            doc_id,\n",
-    "            embedding {config.distance} %s AS distance\n",
+    "            {column_name} {config.distance} %s AS distance\n",
     "            /*, section*/\n",
     "        FROM {config.index_name}\n",
     "        {where}\n",