add search_ef

codelibs · May 5, 2024 · 5a64781 · 5a64781
1 parent 1acad9e
commit 5a64781
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/run-chroma.ipynb b/run-chroma.ipynb
@@ -81,7 +81,7 @@
     "            \"exact\": False,\n",
     "            \"hnsw_m\": 32,\n",
     "            \"hnsw_ef_construction\": 200,\n",
-    "            \"hnsw_ef\": 100, # TODO\n",
+    "            \"hnsw_ef\": 100,\n",
     "            \"update_docs_per_sec\": 0,\n",
     "            \"quantization\": \"\",\n",
     "        },\n",
@@ -97,7 +97,7 @@
     "            \"exact\": False,\n",
     "            \"hnsw_m\": 48,\n",
     "            \"hnsw_ef_construction\": 200,\n",
-    "            \"hnsw_ef\": 100, # TODO\n",
+    "            \"hnsw_ef\": 100,\n",
     "            \"update_docs_per_sec\": 0,\n",
     "            \"quantization\": \"\",\n",
     "        },\n",
@@ -113,7 +113,7 @@
     "            \"exact\": False,\n",
     "            \"hnsw_m\": 48,\n",
     "            \"hnsw_ef_construction\": 200,\n",
-    "            \"hnsw_ef\": 100, # TODO\n",
+    "            \"hnsw_ef\": 100,\n",
     "            \"update_docs_per_sec\": 0,\n",
     "            \"quantization\": \"\",\n",
     "        },\n",
@@ -272,6 +272,7 @@
     "            metadata={\n",
     "                \"hnsw:space\": config.distance,\n",
     "                \"hnsw:construction_ef\": config.hnsw_ef_construction,\n",
+    "                \"hnsw:search_ef\": config.hnsw_ef,\n",
     "                \"hnsw:M\": config.hnsw_m,\n",
     "            }\n",
     "        )\n",
@@ -497,8 +498,8 @@
     "        took = time.time() - now\n",
     "        # print(query_result)\n",
     "\n",
-    "        product_ids = query_result.get(\"ids\")[0]\n",
-    "        scores = query_result.get(\"distances\")[0]\n",
+    "        product_ids = [int(x) for x in query_result.get(\"ids\")[0]]\n",
+    "        scores = [1 - x for x in query_result.get(\"distances\")[0]]\n",
     "        return took * 1000, len(product_ids), product_ids, scores\n",
     "    except Exception as e:\n",
     "        print(f\"[FAIL] {e}\")\n",