diff --git a/lib/model/yake_keywords.py b/lib/model/yake_keywords.py index 7dd7001..ded4a63 100644 --- a/lib/model/yake_keywords.py +++ b/lib/model/yake_keywords.py @@ -8,6 +8,7 @@ import yake import cld3 +import jieba class Model(Model): @@ -36,6 +37,9 @@ 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, @@ -58,6 +62,10 @@ def run_yake(self, text: str, language = cld3.get_language(text).language ### normalize special characters text = self.normalize_special_characters(text) + # Segmentation for mandarin + if language[:2]=="zh": + 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, diff --git a/requirements.txt b/requirements.txt index 8beeca6..e857c31 100644 --- a/requirements.txt +++ b/requirements.txt @@ -24,4 +24,5 @@ numpy==1.26.4 protobuf==3.20.2 openai==1.35.6 anthropic==0.31.1 -pycld3==0.22 \ No newline at end of file +pycld3==0.22 +jieba==0.42.1 \ No newline at end of file diff --git a/test/lib/model/test_yake_keywords.py b/test/lib/model/test_yake_keywords.py index 9a82e72..f31c241 100644 --- a/test/lib/model/test_yake_keywords.py +++ b/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": {