Parametrize slice and threshold

mc-cat-tty · Jan 29, 2024 · 91f74d2 · 91f74d2
1 parent 1b2d198
commit 91f74d2
Showing 1 changed file with 7 additions and 11 deletions.
diff --git a/query_expansion.ipynb b/query_expansion.ipynb
@@ -90,17 +90,9 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 195,
+   "execution_count": 199,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "['house', 'property', 'is', 'building', 'residence']\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "from transformers import BertModel\n",
     "import torch\n",
@@ -135,13 +127,17 @@
     "        .value()\n",
     ")\n",
     "\n",
+    "THRESHOLD = 0.8\n",
+    "SLICE = 5\n",
+    "\n",
     "synonyms = (\n",
     "    pydash.chain(candidates)\n",
     "        .map(itemgetter('token_str'))\n",
     "        .zip(similarities)\n",
+    "        .filter(lambda t: t[1] > THRESHOLD)\n",
     "        .sort(key = itemgetter(1), reverse = True)\n",
     "        .map(itemgetter(0))\n",
-    "        .take(5)\n",
+    "        .take(SLICE)\n",
     "        .value()\n",
     ")\n",
     "\n",