LlamaIndexのWikipedia Loaderを試してみた

Google Colab でLlamaIndexの[Wikipedia Loader](https://llamahub.ai/l/wikipedia)を試したので紹介します。
試しに[「ゆるキャン△」のWikipedia](https://ja.wikipedia.org/wiki/%E3%82%86%E3%82%8B%E3%82%AD%E3%83%A3%E3%83%B3%E2%96%B3)をインデックスにして、その内容に関する質疑応答ができるようにしました。

# Wikipedia Loader
このローダーは、wikipedia Pythonパッケージを使用して、Wikipediaの記事からテキストをフェッチします。引数は、pages(読み込むWikipediaのページタイトル)と **lang(言語設定。既定値は'en')** です。現在のところ、このローダーはテキストのみを抽出し、画像やテーブルなどを無視します。

## 使用方法
このローダーを使うには、Wikipediaのページタイトルを配列で渡す必要があります。
```python
from llama_index import download_loader

WikipediaReader = download_loader("WikipediaReader")

loader = WikipediaReader()
documents = loader.load_data(pages=['ゆるキャン△'], lang='en')
```

# Google Colab
```bash:pip install
!pip install llama-index
!pip install langchain
!pip install openai
```

```python:シークレットキーを登録する
import os
os.environ["OPENAI_API_KEY"] = "<API KEY を貼ってください>"
```

```python:Templateをカスタマイズ
from llama_index import QuestionAnswerPrompt, RefinePrompt


text_qa_template = """あなたはWikipediaのAIヘルプデスクです。以下のコンテキストをもとに質問へ回答してください。
コンテキストに書かれていないことを質問された場合、「ご質問に関する情報が提供されていないため、回答することができません」と答えてください。

コンテキスト：
---------------------
{context_str}
---------------------

質問: {query_str}
"""
text_qa_prompt = QuestionAnswerPrompt(text_qa_template)

refine_template = """元の質問は次のとおりです: {query_str}
    既存の回答を提供しました: {existing_answer}
    既存の答えを洗練する機会があります
    (必要な場合のみ)以下にコンテキストを追加します。
    ------------
    {context_msg}
    ------------
    新しいコンテキストを考慮して、元の答えをより良く洗練して質問に答えてください。
    コンテキストが役に立たない場合は、元の回答と同じものを返します。
    """
refine_prompt = RefinePrompt(refine_template) ```
```

```python:言語モデルを指定する
from llama_index import LLMPredictor, ServiceContext, PromptHelper
from langchain.chat_models import ChatOpenAI


max_input_size = 4096
num_output = 2048
max_chunk_overlap = 20
prompt_helper = PromptHelper(max_input_size, num_output, max_chunk_overlap)
llm_predictor = LLMPredictor(llm= ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0, max_tokens=2048))
service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor, prompt_helper=prompt_helper)
```

```python:インデックスの準備
from llama_index import SimpleDirectoryReader, GPTSimpleVectorIndex, QuestionAnswerPrompt, download_loader
from google.colab import drive


drive.mount('/content/drive')
base_path = '/content/drive/MyDrive/Colab Notebooks/index/'
wiki_title = 'ゆるキャン△'
wiki_lang = 'en'
documents = SimpleDirectoryReader(base_path + f'index_wiki_{wiki_title}_{wiki_lang}').load_data()
isIndexed = True if documents else False
if not isIndexed:
  WikipediaReader = download_loader("WikipediaReader")
  loader = WikipediaReader()
  documents = loader.load_data(pages=[wiki_title], lang=wiki_lang)

index = GPTSimpleVectorIndex.from_documents(documents,  service_context=service_context)

if not isIndexed:
  # save to disk
  index_json = base_path +  f'index_wiki_{wiki_title}_{wiki_lang}.json'
  index.save_to_disk(index_json)
```

```python:質疑応答をする
query_str = "各務原なでしこの性格を教えてください。？"

answer = index.query(
      query_str,
      text_qa_template=text_qa_prompt,
      refine_template=refine_prompt,
      similarity_top_k=1
    )
print(answer)
```

```bash:answer
各務原なでしこは、元気で明るい性格の持ち主であり、アウトドア活動が大好きです。彼女はキャンプや登山などのアウトドア活動を楽しんでおり、物語の中心となるリンとの友情も魅力的です。また、彼女の姓は岐阜県の各務原市の名前に由来しています。彼女は自分のペースで物事を進めることができ、周りの人々にも優しく接することができます。Laid-Back Campは、自然と友情を称賛する素晴らしいアニメであり、2010年代の最高のアニメの1つとしてIGNやCrunchyrollなどのメディアから高い評価を受けています。さらに、アニメは登場する場所の観光客数を3倍に増加させるなど、地域の観光振興にも貢献しています。
```

# 感想
Wikipediaの記事をChatGPTに認識させることができました。Refine時に、徐々に回答にノイズが含まれるようになります。これはRefinePromptを工夫することである程度軽減できると思いますので、試してみる予定です。例えば、質問と関連しないコンテキストは無視するような指示をすることが挙げられます。