自己紹介と LlamaIndex でSQL生成を試す

hmj

はじめまして。Hmj_kd です。
Twitter では白塗りのアイコンでやっております。

https://twitter.com/Hmj_kd

note では最近勉強と試しなどを始めた LLM (大規模言語モデル, Large Language Models) 関係のことを書きたくてはじめました。
これまでは、決済サービス内での不正対策チームに所属していまして、主に機械学習やグラフ理論を活用して施策を検討したり、実装や運用をしておりました。
4月からチーム変わり、急遽LLM周辺のことをキャッチアップしながら、日々小さく色々試しています。

LLM Meetup Tokyo #1

こちらに先日参加できたので、デモの内容をnoteにまとめたいと思います。

直近LLMを活用した社内の施策をいくつか検討していまして、その中でも社内にいっぱいあるドキュメント（例えば、仕様書、レポート、ミーティングのドキュメント etc）を横断的に検索できると良いし、対話的にやり取りや要約できるといいよね、というのがあります。

そのため、私にとってのスタートラインが"プロンプトを試行錯誤する"ではなく、外部データ（≒ 社内データ）を活用したいというのがモチベーションでありました。

そこで発見したのが、LlamaIndexやLangChainでした。
（他の方がわかりやすい解説記事を書いていらっしゃるので、割愛します）

試した結果はやりたいことがほぼできそうで、とても感動しています。
また、各種ドキュメントがベクトルになることで、問い合わせ内容のベクトル化をベクトル検索するだけでもいろいろ解決しそうで、このあたりの検索は今後熱くなりそうと感じました。

BigQueryのSchemaをdocumentとしてindexする

長くなりましたが、ここから本題で、社内文章をindexして色々試している時に、一番感動したのがSQL生成でした（まだこれからも感動がでそうですが）。

BigQueryのSchemaをドキュメントとみなして、indexさせます。たとえば「取引数に関するテーブルは？」で「このテーブルです、hogehoge….」というのは容易に想像がつきますし、できました。

その後、感動したのが、自然言語からSQL自体が作られたことでした。

まず、特にindexもせずに gpt-3.5-turbo に下記の内容でsubmitしてみます。

以下を実現するBigQueryのSQLを記述する。
```
2023年3月の購入金額の総和を取得する。
```
SQLを生テキストで出力します。

結果は、下記で、これだけでもすごいです。

以下が要求されたBigQueryのSQL文です。

``` 
SELECT SUM(purchase_amount) AS total_purchase_amount
FROM `table_name`
WHERE EXTRACT(YEAR FROM purchase_date) = 2023
AND EXTRACT(MONTH FROM purchase_date) = 3; 
```

上記クエリは、2023年3月に購入されたアイテムの購入金額の総和を取得します。`table_name`は、データが格納されているテーブルの名前に置き換えてください。また、`purchase_amount`は購入金額が格納されている列の名前に置き換えてください。`purchase_date`は購入日が格納されている列の名前に置き換えてください。

スキーマは社内情報のため、ここではダミーな一般的な内容に変えています。

Document と index の作成

# Python
schema1 = json.dumps(
    {
        'project': 'test_project',
        'dataset': 'test_dataset',
        'table': 'orders',
        'schema': [
            {'name': 'Id', 'type': 'INTEGER', 'mode': 'NULLABLE', 'description': 'id'},
            {'name': 'UserId', 'type': 'INTEGER', 'mode': 'NULLABLE', 'description': '顧客ID'},
            {'name': 'Amount', 'type': 'INTEGER', 'mode': 'NULLABLE', 'description': '金額'},
            {'name': 'CreatedAt', 'type': 'TIMESTAMP', 'mode': 'NULLABLE', 'description': '生成日時'},
        ]
    }
)

schema2 = json.dumps(
    {
        'project': 'test_project',
        'dataset': 'test_dataset',
        'table': 'customers',
        'schema': [
            {'name': 'Id', 'type': 'INTEGER', 'mode': 'NULLABLE', 'description': 'id'},
            {'name': 'Name', 'type': 'STRING', 'mode': 'NULLABLE', 'description': '氏名'},
            {'name': 'CreatedAt', 'type': 'TIMESTAMP', 'mode': 'NULLABLE', 'description': '生成日時'},
        ]
    }
)


schema3 = json.dumps(
    {
        'project': 'test_project',
        'dataset': 'test_dataset',
        'table': 'products',
        'schema': [
            {'name': 'Id', 'type': 'INTEGER', 'mode': 'NULLABLE', 'description': 'id'},
            {'name': 'Name', 'type': 'STRING', 'mode': 'NULLABLE', 'description': '商品名'},
            {'name': 'Category', 'type': 'STRING', 'mode': 'NULLABLE', 'description': '商品カテゴリー, カテゴリー'},
            {'name': 'CreatedAt', 'type': 'TIMESTAMP', 'mode': 'NULLABLE', 'description': '生成日時'},
        ]
    }
)

from langchain.chat_models import ChatOpenAI
from llama_index import (
    GPTSimpleVectorIndex, LLMPredictor, OpenAIEmbedding, ServiceContext, Document
)

service_context = ServiceContext.from_defaults(
    llm_predictor = LLMPredictor(llm=ChatOpenAI(temperature=0, model_name="gpt-3.5-turbo")),
    embed_model = OpenAIEmbedding(model="text-embedding-ada-002")
)

index = GPTSimpleVectorIndex.from_documents(
    documents=[
        Document(schema1),
        Document(schema2),
        Document(schema3),
    ],
    service_context=service_context
)

Embeddingには、現時点で推奨されている text-embedding-ada-002 を利用しています。

質問と応答

from llama_index.prompts.prompts import QuestionAnswerPrompt

fmt_qa_tmpl = '''
以下を実現するBigQueryのSQLを記述する。
```
{query_str}
```
対象テーブルのフォーマットは以下の通りです。
{context_str}
SQLを生テキストで出力します。
'''

qa_prompt = QuestionAnswerPrompt(fmt_qa_tmpl)

response = index.query(
    "2023年3月の購入金額の総和を取得する。",
    text_qa_template=qa_prompt,
    similarity_top_k=1,
)

qa_promptをデフォルトのフォーマットから少し変更しています。
これにより返ってくる結果は、以下のようになりました。

SELECT SUM(Amount)
FROM `test_project.test_dataset.orders`
WHERE EXTRACT(YEAR FROM CreatedAt) = 2023
AND EXTRACT(MONTH FROM CreatedAt) = 3

と、問い合わせ内容と保持しているindex内で類似性の高いテーブルのSchemaから上記のようにSQLが生成されました!! (すごい)

まとめ

  • LLMで外部データを活用する方法を模索していて、LlamaIndexなどを試すことができ、見通しが良くなりました

  • 副産物で、テーブルスキーマをドキュメントとして保持することで、簡単なSQLなら、生成もできるようになりました（ただし、ここは100%の精度ではない点に注意）

  • ここまで取り組みを通じて、目的に応じてプロンプトを変えることが必要と、結局社内情報はちゃんとデータ化してindex化して活用できるようにスタンバイするのが大事だなと感じています。

引き続き、LLM周りで試していきます。よろしくおねがいします！

この記事が気に入ったらサポートをしてみませんか？