feat: Added an option to choose the language in Wikipedia texts (#142)

This commit is contained in:
Emmanuel 2023-03-27 22:56:54 -03:00 committed by GitHub
parent d0185fc543
commit e43264cbbc
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -12,17 +12,18 @@ class WikipediaReader(BaseReader):
"""
def load_data(self, pages: List[str], **load_kwargs: Any) -> List[Document]:
def load_data(self, pages: List[str], lang: str = "en", **load_kwargs: Any) -> List[Document]:
"""Load data from the input directory.
Args:
pages (List[str]): List of pages to read.
lang (str): language of wikipedia texts (default English)
"""
import wikipedia
results = []
for page in pages:
wikipedia.set_lang(lang)
page_content = wikipedia.page(page, **load_kwargs).content
results.append(Document(page_content))
return results