add wikipedia reader (#31)

examples/test_wiki/TestWikiReader.ipynb

+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c5d167a5-81f8-4d2c-b42f-0a190577132f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# My OpenAI Key\n",
+    "import os\n",
+    "os.environ['OPENAI_API_KEY'] = \"INSERT OPENAI KEY\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 36,
+   "id": "5f60348e-731d-4a95-bae2-426e184a914e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from gpt_index import GPTKeywordTableIndex, WikipediaReader"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 37,
+   "id": "952c4659-7fbb-447e-8caf-06916412cc37",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "page: Covid-19\n"
+     ]
+    }
+   ],
+   "source": [
+    "wiki_docs = WikipediaReader().load_data(pages=['Covid-19'])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3be202db-a4c7-41d2-ba7d-446d1f934830",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "index = GPTKeywordTableIndex(wiki_docs)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 42,
+   "id": "7f5667a9-6758-447b-9af2-5e5a4d008a29",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# save index to docs\n",
+    "index.save_to_disk('index_covid.json')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 45,
+   "id": "77340460-8319-474f-91eb-545ea5790127",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "new_index = GPTKeywordTableIndex.load_from_disk('index_covid.json')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 46,
+   "id": "28d7163e-f26f-4ad8-89d5-9cb7662c4d9c",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "> Starting query: Which country included tocilizumab in treatment for covid-19?\n",
+      "Extracted keywords: ['tocilizumab', 'treatment', 'covid-19', 'covid', '19']\n",
+      "> Querying with idx: 1105763466456338724: of age or older weighing at least 40 kilograms ...\n",
+      "> Querying with idx: 2820318727532393752: Coronavirus disease 2019 (COVID-19) is a contag...\n",
+      "> Querying with idx: 897499143815831368: if the mask includes an exhalation valve, a wea...\n",
+      "> Querying with idx: 8628144746434065339: pulmonary fibrosis, cystic fibrosis. Evidence s...\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "'\\n\\nChina'"
+      ]
+     },
+     "execution_count": 46,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# GPT doesn't find the corresponding evidence in the leaf node, but still gives the correct answer\n",
+    "# try verbose=True for more detailed outputs\n",
+    "new_index.query(\"Which country included tocilizumab in treatment for covid-19?\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "60ae24ad-0aa1-4844-8e9a-caf30c74643e",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "myvenv",
+   "language": "python",
+   "name": "myvenv"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

gpt_index/__init__.py

@@ -17,11 +17,13 @@ from gpt_index.prompts.base import Prompt
 
 # readers
 from gpt_index.readers.simple_reader import SimpleDirectoryReader
+from gpt_index.readers.wikipedia import WikipediaReader
 
 __all__ = [
     "GPTKeywordTableIndex",
     "GPTListIndex",
     "GPTTreeIndex",
     "Prompt",
+    "WikipediaReader",
     "SimpleDirectoryReader",
 ]

gpt_index/readers/base.py

 """Base reader class."""
 from abc import abstractmethod
-from typing import List
+from typing import Any, List
 
 from gpt_index.schema import Document
 
@@ -9,5 +9,5 @@ class BaseReader:
     """Utilities for loading data from a directory."""
 
     @abstractmethod
-    def load_data(self) -> List[Document]:
+    def load_data(self, **load_kwargs: Any) -> List[Document]:
         """Load data from the input directory."""

gpt_index/readers/simple_reader.py

 """Simple reader that ."""
 from pathlib import Path
-from typing import List
+from typing import Any, List
 
 from gpt_index.readers.base import BaseReader
 from gpt_index.schema import Document
@@ -22,7 +22,7 @@ class SimpleDirectoryReader(BaseReader):
                 raise ValueError(f"Expected {input_file} to be a file.")
         self.input_files = input_files
 
-    def load_data(self) -> List[Document]:
+    def load_data(self, **load_kwargs: Any) -> List[Document]:
         """Load data from the input directory."""
         data = ""
         for input_file in self.input_files:

gpt_index/readers/wikipedia.py

+"""Simple reader that ."""
+from typing import Any, List
+
+from gpt_index.readers.base import BaseReader
+from gpt_index.schema import Document
+
+
+class WikipediaReader(BaseReader):
+    """Wikipedia reader.
+
+    Reads a page.
+
+    """
+
+    def __init__(self) -> None:
+        """Initialize with parameters."""
+        try:
+            import wikipedia  # noqa: F401
+        except ImportError:
+            raise ValueError(
+                "`wikipedia` package not found, please run `pip install wikipedia`"
+            )
+
+    def load_data(self, **load_kwargs: Any) -> List[Document]:
+        """Load data from the input directory."""
+        import wikipedia
+
+        pages: List[str] = load_kwargs.pop("pages", None)
+        if pages is None:
+            raise ValueError('Must specify a "page" in `load_kwargs`.')
+
+        results = []
+        for page in pages:
+            page_content = wikipedia.page(page).content
+            results.append(Document(page_content))
+        return results

requirements.txt

 -e .
+
+# third-party
+wikipedia
+
 # linting
 black
 isort