Update pythainlp.wsd

PyThaiNLP · Jul 12, 2023 · d14d2b8 · d14d2b8
1 parent 0753174
commit d14d2b8
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 31 deletions.
diff --git a/notebooks/test_wsd.ipynb b/notebooks/test_wsd.ipynb
@@ -24,25 +24,13 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Setting ds_accelerator to cuda (auto detect)\n"
+      "Setting ds_accelerator to cuda (auto detect)\n",
+      "[('โปรแกรมคอมพิวเตอร์ใช้ในทางอินเทอร์เน็ตสำหรับเก็บข้อมูลของผู้ใช้งาน', 0.0974416732788086), ('ชื่อขนมชนิดหนึ่งจำพวกขนมเค้ก แต่ทำเป็นชิ้นเล็ก ๆ แบน ๆ แล้วอบให้กรอบ', 0.09319090843200684)]\n"
      ]
-    },
-    {
-     "data": {
-      "text/plain": [
-       "[('โปรแกรมคอมพิวเตอร์ใช้ในทางอินเทอร์เน็ตสำหรับเก็บข้อมูลของผู้ใช้งาน',\n",
-       "  0.0974416732788086),\n",
-       " ('ชื่อขนมชนิดหนึ่งจำพวกขนมเค้ก แต่ทำเป็นชิ้นเล็ก ๆ แบน ๆ แล้วอบให้กรอบ',\n",
-       "  0.09319090843200684)]"
-      ]
-     },
-     "execution_count": 2,
-     "metadata": {},
-     "output_type": "execute_result"
     }
    ],
    "source": [
-    "get_sense(\"เขากำลังอบขนมคุกกี้\",\"คุกกี้\")"
+    "print(get_sense(\"เขากำลังอบขนมคุกกี้\",\"คุกกี้\"))"
    ]
   },
   {
@@ -54,30 +42,34 @@
    },
    "outputs": [
     {
-     "data": {
-      "text/plain": [
-       "[('โปรแกรมคอมพิวเตอร์ใช้ในทางอินเทอร์เน็ตสำหรับเก็บข้อมูลของผู้ใช้งาน',\n",
-       "  0.1005704402923584),\n",
-       " ('ชื่อขนมชนิดหนึ่งจำพวกขนมเค้ก แต่ทำเป็นชิ้นเล็ก ๆ แบน ๆ แล้วอบให้กรอบ',\n",
-       "  0.12473666667938232)]"
-      ]
-     },
-     "execution_count": 3,
-     "metadata": {},
-     "output_type": "execute_result"
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[('โปรแกรมคอมพิวเตอร์ใช้ในทางอินเทอร์เน็ตสำหรับเก็บข้อมูลของผู้ใช้งาน', 0.1005704402923584), ('ชื่อขนมชนิดหนึ่งจำพวกขนมเค้ก แต่ทำเป็นชิ้นเล็ก ๆ แบน ๆ แล้วอบให้กรอบ', 0.12473666667938232)]\n"
+     ]
     }
    ],
    "source": [
-    "get_sense(\"เว็บนี้ต้องการคุกกี้ในการทำงาน\",\"คุกกี้\")"
+    "print(get_sense(\"เว็บนี้ต้องการคุกกี้ในการทำงาน\",\"คุกกี้\"))"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 4,
    "id": "b0ee35fc-f26e-4bce-b6fa-0e1efc863ae4",
    "metadata": {},
-   "outputs": [],
-   "source": []
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "None\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(get_sense(\"เว็บนี้ต้องการคุกกี้ในการทำงาน\",\"คน\"))"
+   ]
   }
  ],
  "metadata": {

diff --git a/pythainlp/wsd/core.py b/pythainlp/wsd/core.py
@@ -106,7 +106,7 @@ def get_sense(
     """
     global _MODEL
     _w = custom_tokenizer.word_tokenize(sentence)
-    if word not in _w:
+    if word not in set(custom_dict.keys()) or word not in sentence:
         return None
     if _MODEL == None:
         _MODEL = _SentenceTransformersModel(device=device)

diff --git a/tests/test_wsd.py b/tests/test_wsd.py
@@ -7,3 +7,4 @@ class TestWsdPackage(unittest.TestCase):
     def test_get_sense(self):
         self.assertIsNotNone(get_sense("เขากำลังอบขนมคุกกี้","คุกกี้"))
         self.assertIsNotNone(get_sense("เว็บนี้ต้องการคุกกี้ในการทำงาน","คุกกี้"))
+        self.assertIsNone(get_sense("เว็บนี้ต้องการคุกกี้ในการทำงาน","คน"))