アルバイトの Sasaokaです。今回はChatGPTにCSVやYouTube動画などの独自データを情報を付与して回答させることができるLlamaIndexというライブラリを試してみました。

準備

次の3つのライブラリをインストールします。

$ pip install llama-index langchain youtube_transcript_api

OpenAIのAPIキーを環境変数OPENAI_API_KEYに設定しておきます。

$ export OPENAI_API_KEY='sk-xxxxxx'

CSVデータを与える

ここでは弊社のアドクロールの広告タイトル文を読み込ませてみます。次のような、id、飲料の広告の会社、カテゴリ、ブランド、広告タイトルをカラムとする258レコードのdrink.csvファイルを用いました。

"id","company","category","brand","title"
13556981,"株式会社アルプロン","プロテイン飲料","Alpron","【アルプロン】🎉年に一度の大感謝SALE🎊 目玉商品を多数ご用意！  #プロテイン   #トレーニング #ソイプロテイン　#ホエイプロテイン #アルプロン合併記念　＃サプリメント #セール #アルプロン #ALPRON"
13674926,"The Hut.com Limited","プロテイン飲料","BCAA 4:1:1 パウダー","【最大80%OFF💥】今年最大級セールついに開催！大人気ホエイプロテインやEAA等、割引コード使用で半額以下🤩なんと新発売のImpact ホエイプロテインのメロンミルクフレーバーやブルーベリーレイヤードバー、【数量限定】人気商品お試しBOXなども割引対象に！ 数量限定なのでお買い求めはお早めに！ 売り切れる前に今すぐカートへ🏃💨 ------ #マイプロテイン #Myprotien #プロテイン #アミノ酸 #筋トレ"

CSVファイルをインデックスとして保存する際、追加情報をextra_infoとして辞書型で与えることができます。今回は、広告のidをextra_infoに保存することにします。それを実現するために、以下のCSVReaderを作成しました。

import csv
from pathlib import Path
from typing import Any, List

from llama_index.readers.base import BaseReader
from llama_index.readers.schema.base import Document

class CSVReader(BaseReader):

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)

    def load_data(self, file: Path) -> List[Document]:
        text_list = []
        extra_info_list = []
        with open(file, "r") as f:
            for row in csv.DictReader(f):
                extra_info_list.append({"id": row["id"], "content": "image"})
                text_list.append(", ".join([v for k, v in row.items() if k != "id"]))

        return [Document(text, extra_info=extra_info) for text, extra_info in zip(text_list, extra_info_list)]

from llama_index import download_loader, ServiceContext, GPTVectorStoreIndex, LLMPredictor
from langchain.chat_models import ChatOpenAI

loader = CSVReader()
documents =  loader.load_data('./drink.csv')

llm_predictor = LLMPredictor(llm=ChatOpenAI(model_name='gpt-3.5-turbo'))
service_context = ServiceContext.from_defaults(llm_predictor)

index = GPTVectorStoreIndex.from_documents(documents, service_context=service_context)
index.storage_context.persist('./drink')

このスクリプトを実行するとdrinkというディレクトリができ、その中に3つJSONファイルが保存されます。docstore.jsonには読み込んだテキストとextra_infoが、vector_store.jsonにはEmbedding APIでベクトル化された各テキストのベクトル表現が保存されます。

次に保存したインデックスを読み込み、質問に回答してもらうスクリプトは以下になります。ここではインデックスを保存したディレクトリと質問をコマンドライン引数として受け取るようにしています。

import sys
from llama_index import StorageContext, load_index_from_storage

storage_context = StorageContext.from_defaults(persist_dir=sys.argv[1])
index = load_index_from_storage(storage_context)

query_engine = index.as_query_engine()
response = query_engine.query(sys.argv[2])
print(response)

このスクリプトをask.pyとして保存し、"サントリーの商品の特徴を教えて"というプロンプトを送ると、読み込んだCSVの内容を踏まえた回答が返ってきました。

$ python ask.py drink "サントリーの商品の特徴を教えて"

サントリーの商品の特徴は、日本の四季を表現した六種の和素材をブレンドした酒類や、Netflixの人気コンテンツをモチーフにした限定デザインTシャツなど、様々な商品を提供しています。また、サントリーのウイスキーは、自由なスタイルで楽しめるように、オンラインライブを開催しています。

YouTube動画を与える

LlamaHub に公開されているYoutube Transcript Loaderを用いてYouTube動画の字幕データを与えることができます。ローダーのload_dataメソッドの引数に対象のYouTube動画のURLと字幕言語を指定して使用します。字幕言語のデフォルトは英語(en)です。字幕の存在しない動画URLや指定言語の字幕が存在しない場合はエラーになります。

今回は日本語の自動生成字幕が付けられている漫画広告動画を使用してみました。

from llama_index import download_loader, ServiceContext, GPTVectorStoreIndex, LLMPredictor
from langchain.chat_models import ChatOpenAI

YoutubeTranscriptReader = download_loader('YoutubeTranscriptReader')
loader = YoutubeTranscriptReader()
documents = loader.load_data(
    ytlinks=['https://www.youtube.com/watch?v=O-7qqkmUhXM'],
    languages=['ja']
)

llm_predictor = LLMPredictor(llm=ChatOpenAI(model_name='gpt-3.5-turbo'))
service_context = ServiceContext.from_defaults(llm_predictor)

index = GPTVectorStoreIndex.from_documents(documents, service_context=service_context)
index.storage_context.persist('./youtube_sample')

「この動画を要約して」というプロンプトを送ると漫画部分を省いた、宣伝部分を要約してくれました。

$ python ask.py youtube_sample "この動画を要約して"

この動画では、業界初の医薬部外品レーザーコスメを紹介しています。通常価格11,080円のところ、初回限定82％オフの1980円でお試しできます。定期縛りなしで、試すことができ、シミを消すことができます。また、楽天で1位を獲得している高評価のクリームであり、プルプルテクスチャーで肌をごくごく飲み込むように吸収し、保湿もしみ対策もできます。500人限定であるため、在庫切れ

おわりに

LlamaindexとOpenAIのAPIを用いることで、ChatGPTに独自データに基づいた回答をさせることができます。回答の多様性が多少失われますが、いい加減なことを言わなくなるというメリットがあります。興味を持った方は使ってみてください。