adding tests to jieba and chinese text with yake

lib/model/yake_keywords.py

@@ -37,6 +37,8 @@ def normalize_special_characters(self, text):
             text = text.replace(k, v)
         return text
 
+    def run_chinese_segmentation_with_jieba(self, text):
+        return " ".join(list(jieba.cut_for_search(text)))
     def run_yake(self, text: str,
                  language: str,
                  max_ngram_size: int,
@@ -61,7 +63,8 @@ def run_yake(self, text: str,
         text = self.normalize_special_characters(text)
         # Segmentation for mandarin
         if language == 'zh-CN' or language == 'zh' or language == 'zh-TW':
-            text = " ".join(list(jieba.cut_for_search(text)))
+            text = self.run_chinese_segmentation_with_jieba(text)
+            # text = " ".join(list(jieba.cut_for_search(text)))
         ### extract keywords
         custom_kw_extractor = yake.KeywordExtractor(lan=language, n=max_ngram_size, dedupLim=deduplication_threshold,
                                                     dedupFunc=deduplication_algo, windowsSize=window_size,

test/lib/model/test_yake_keywords.py

@@ -53,6 +53,22 @@ def test_normalize_special_characters(self):
         expected = "'''\"\""
         self.assertEqual(self.yake_model.normalize_special_characters(text), expected)
 
+    def test_run_chinese_segmentation_with_jieba(self):
+        test_text = '''哈里斯同意与特朗普再进行一次美大选辩论'''
+        expected = "哈里 里斯 哈里斯 同意 与 特朗普 再 进行 一次 美 大选 辩论"
+        self.assertEqual(self.yake_model.run_chinese_segmentation_with_jieba(test_text), expected)
+
+    def test_run_yake_real_with_chinese(self):
+        message = schemas.parse_input_message({
+            "body": {
+                "id": "1234",
+                "text": "哈里斯同意与特朗普再进行一次美大选辩论",
+            },
+            "model_name": "yake_keywords__Model"
+        })
+        results = self.yake_model.run_yake(**self.yake_model.get_params(message))
+        self.assertEqual(results, {"keywords":  [('哈里斯', 0.04491197687864554), ('特朗普', 0.04491197687864554)]})
+
     def test_get_params_with_defaults(self):
         message = schemas.parse_input_message({
             "body": {