From 60d406acaa25725225ca6f988e4a0c18ecff238b Mon Sep 17 00:00:00 2001 From: H <43509927+guoyuhao2330@users.noreply.github.com> Date: Mon, 15 Jul 2024 17:45:40 +0800 Subject: [PATCH] Set wikipedia lang (#1515) ### What problem does this PR solve? ### Type of change - [x] New Feature (non-breaking change which adds functionality) --- graph/component/wikipedia.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/graph/component/wikipedia.py b/graph/component/wikipedia.py index 9e6787596..48abefaf6 100644 --- a/graph/component/wikipedia.py +++ b/graph/component/wikipedia.py @@ -30,9 +30,16 @@ class WikipediaParam(ComponentParamBase): def __init__(self): super().__init__() self.top_n = 10 + self.lang = 'en' def check(self): self.check_positive_integer(self.top_n, "Top N") + self.check_valid_value(self.lang, "Wikipedia languages", + ['af', 'pl', 'ar', 'ast', 'az', 'bg', 'nan', 'bn', 'be', 'ca', 'cs', 'cy', 'da', 'de', + 'et', 'el', 'en', 'es', 'eo', 'eu', 'fa', 'fr', 'gl', 'ko', 'hy', 'hi', 'hr', 'id', + 'it', 'he', 'ka', 'lld', 'la', 'lv', 'lt', 'hu', 'mk', 'arz', 'ms', 'min', 'my', 'nl', + 'ja', 'nb', 'nn', 'ce', 'uz', 'pt', 'kk', 'ro', 'ru', 'ceb', 'sk', 'sl', 'sr', 'sh', + 'fi', 'sv', 'ta', 'tt', 'th', 'tg', 'azb', 'tr', 'uk', 'ur', 'vi', 'war', 'zh', 'yue']) class Wikipedia(ComponentBase, ABC): @@ -45,9 +52,11 @@ class Wikipedia(ComponentBase, ABC): return Wikipedia.be_output(self._param.no) wiki_res = [] - for wiki_key in wikipedia.search(ans, results=self._param.top_n): + wikipedia.set_lang(self._param.lang) + wiki_engine = wikipedia + for wiki_key in wiki_engine.search(ans, results=self._param.top_n): try: - page = wikipedia.page(title=wiki_key, auto_suggest=False) + page = wiki_engine.page(title=wiki_key, auto_suggest=False) wiki_res.append({"content": '' + page.title + ' ' + page.summary}) except Exception as e: print(e)