Update common.py

PyThaiNLP · Dec 12, 2023 · 92ad41a · 92ad41a
1 parent cf6647e
commit 92ad41a
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/pythainlp/corpus/common.py b/pythainlp/corpus/common.py
@@ -354,13 +354,20 @@ def find_synonyms(word: str) -> List[str]:
         from pythainlp.corpus import find_synonyms
 
         print(find_synonyms("หมู"))
-        # output: ['จรุก', 'วราห์', 'วราหะ', 'ศูกร', 'สุกร']
+        # output: ['จรุก', 'วราหะ', 'วราห์', 'ศูกร', 'สุกร']
     """
     synonyms = thai_synonyms()  # get a dictionary of {word, synonym}
     list_synonym = []
 
-    for idx, words in enumerate(synonyms["word"]):
+    if word in synonyms["word"]: # find by word
+        list_synonym.extend(synonyms["synonym"][synonyms["word"].index(word)])
+
+    for idx, words in enumerate(synonyms["synonym"]): # find by synonym
         if word in words:
             list_synonym.extend(synonyms["synonym"][idx])
+            list_synonym.append(synonyms["word"][idx])
+
+    if word in list_synonym: # remove same word
+        list_synonym.remove(word)
 
     return sorted(list(set(list_synonym)))