↗️

簡単なRAGチャットボットを作って、RAG、Embedding、VetcorDBについてさっくり理解しよう

に公開

はじめに

こんにちは。Kouchと申します。
以前こんな記事を書きました。
https://zenn.dev/kouch/articles/9338930c27f02c
今回はこの中でもRAGについて、実際にPythonのコードをお見せしながら解説していこうと思います。
せっかくなので、小さなチャットボットを作り、RAGが動いた感動ポイントを皆さんにも味わってもらえればと思います。

RAG実装の全体像

今回のゴール

事前にある程度のテキストを与えて、そのテキストに紐づいた回答を生成するチャットボット、を今回のゴールとします。

今回インプットする事前情報はこちらです。

texts = [
  "私はオランダで生まれて、東京で育ちました",
  "趣味は自転車とポーカーです。",
  "AIについて勉強しています。",
  "休日はカフェで作業することが多いよ。",
  "猫を飼っていて、まろんという名前です",
  "2年前に結婚しました",
  "旅行が好きで、先月はシリコンバレーに行きました",
  "英語の勉強をしてて、もっと話せるようになりたいな",
  "赤い色が好きです!",
  "この前自転車を盗まれて、新しい自転車を買いました。悲しい…",
  "最近夜更かししてるから早く寝ないとなあ",
]

色々な質問をしてみましょう。

質問: ペットは飼ってる?
はい、まろんという名前の猫を飼っています。

質問: 先月は何してた?
先月はシリコンバレーに行きました。

質問: 勉強してることはある?
手元のメモには、英語とAIを勉強していると書かれています。

質問: 身長は?
手元のメモには情報がありません。

このようにハルシネーションすることなく、与えた情報から回答を作成してくれます。
これをさらに応用することで、企業内のNotionなどのWikiを参照して自社のチャットボットを作ったり、問い合わせ対応のチャットボットなどを作ることができますね。

実際のコード

全体像のコードをここでお見せします。
ここから細かく解説していきます。
(環境変数は.env などに保存してください。)

from dotenv import load_dotenv
import os
from openai import OpenAI
import chromadb
from chromadb.utils import embedding_functions

load_dotenv()
embedding_model = "text-embedding-3-large"
openai_client = OpenAI(
  api_key = os.getenv("OPENAI_API_KEY")
)

chromadb_client = chromadb.PersistentClient(path="./chroma_db")
openai_ef = embedding_functions.OpenAIEmbeddingFunction(
  api_key = os.getenv("OPENAI_API_KEY"),
  model_name = embedding_model
)

collection = chromadb_client.get_or_create_collection(
  name = "test_collection",
  embedding_function = openai_ef
)

texts = [
  "私はオランダで生まれて、東京で育ちました",
  "趣味は自転車とポーカーです。",
  "AIについて勉強しています。",
  "休日はカフェで作業することが多いよ。",
  "猫を飼っていて、まろんという名前です",
  "2年前に結婚しました",
  "旅行が好きで、先月はシリコンバレーに行きました",
  "英語の勉強をしてて、もっと話せるようになりたいな",
  "赤い色が好きです!",
  "この前自転車を盗まれて、新しい自転車を買いました。悲しい…",
  "最近夜更かししてるから早く寝ないとなあ",
]

collection.upsert(
  documents = texts, 
  ids = [f"doc_{i}" for i in range(len(texts))]
)

def rag_answer(question: str):
  response = collection.query(
    query_texts = [question],
    n_results = 3,
  )   

  contexts = []
  for i in range(3):
    contexts.append(response['documents'][0][i])

  prompt = f"""
    以下の情報をもとに質問に答えてください。
    あなたは厳密なリサーチアシスタントです。
    以下の「参照メモ」に書かれている内容【だけ】を根拠に、日本語で簡潔に答えてください。
    参照に無い内容は推測せず、「手元のメモには情報がありません」と述べてください。

    参考メモ: {contexts}
    質問: {question}
  """

  chat_response = openai_client.chat.completions.create(
    model="gpt-4o",
    messages=[
        {"role": "system", "content": "あなたは厳密なリサーチアシスタントです。"},
        {"role": "user", "content": prompt}
    ]
  )

  return chat_response.choices[0].message.content

print("質問を入力してください")
user_question = input("質問: ")
answer = rag_answer(user_question)
print(answer)

ざっくり全体的な流れ

  1. 事前の外部情報を数値のベクトルに変換(Embedding)する
  2. 変換されたベクトル情報を、VectorDBに保存する
  3. 入力されたクエリから、参考になりそうな情報(コンテキスト)をDBから検索する
  4. 元々用意していたプロンプトと、DBから引っ張ってきたコンテキストを合わせてLLMに投げる
  5. 生成された回答をユーザーに表示する
    という流れになります。

それぞれ説明していきましょう!

主要なプロセスの解説

Embeddingとは?

まずはテキストを数字化するプロセス、Embeddingのプロセスです。
普段使っているのはOpenAIのChatをするためのAPI(gpt-4oなど)が多いと思いますが、それとは別にEmbeddingするためのAPIがあります。
今回はそれを利用してベクトル化をやってみましょう。

本題のコードに入る前に、例としてある文章をOpenAIのEmbeddingAPIを使ってベクトル化してみましょう

client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
embedding_model = "text-embedding-3-large"

text = "猫を飼っていて、まろんという名前です"

emb = client.embeddings.create(
  model = embedding_model,
  input = text
)
vec = emb.data[0].embedding
print(vec)

このコードを実行すると、テキストが変換されたベクトルが見れます。

[0.028040073812007904, -0.008439182303845882, 0.008930659852921963, ....]

今回は、 text-embedding-3-large というモデルを利用しているので3072次元のベクトルに変換されました。
別のtext-embedding-3-small という軽量のモデルを使うと、1536次元になります。
このベクトルはその単語や文章の意味を表しています。
より高次元数のベクトルの方が、意味を細かく表現できるようになります。

こちらの動画は意味ベクトルを三次元空間で可視化しながら説明しているわかりやすい動画なので、もっと詳しく知りたい方はこちらもご覧ください。
https://www.youtube.com/watch?v=KlZ-QmPteqM

どうしてベクトル化する必要がある?

DBに保存するときに、わざわざベクトルの形で保存する理由は、「意味で検索できるようにするため」です。
例えば、テキストで「猫を飼っていて、まろんという名前です」と保存している場合、「猫は飼っていますか?」という問いに対しては「猫」という単語が入っているので一致検索で見つけることができると思います。
ですが「ペットは飼っていますか?」という問いに対しては関連性を見出すことができません。

Embeddingをすることで、3072次元(1536次元)のベクトルになりますが、これによって「猫を飼っている」「ペットを飼っている」の2つの文章が意味的に近い、ということが検索できるようになります。
これによって「ペットの話してるけど、猫の話題が意味として近そうだな」と参照できるようになるんですね。

つまり、柔軟に検索して、より関連性のあるコンテキストを抽出するために、単語の表層だけではなく、「意味」まで理解させるためのプロセスになります。

VertorDB

さて、テキストをベクトルの形にしたところで、それをDBに保存します。
このようなベクトルを保存し、類似検索を可能にするDBをVectorDBといい、この記事ではVectorDBの一種であり、Pythonで動かしやすいChromaDBというライブラリを使っています。

MySQLなどの一般的なDBは、一致検索や条件検索が主な使用用途でした。
例えば「user_idが17のArticleを取ってくる」「nameが"Kouch"のUser一覧を取得する」などですね。
VectorDBの場合は類似検索に特化されており、ベクトル同士の距離が近いものを返す、という挙動を目的として作られています。

前半まとめ

はい、ここまでで前半が終了しました。
事前にテキストを与えて、それをEmbeddingし、VectorDBに保存する、というプロセスです。
コードだと、こちらの部分になります。

embedding_model = "text-embedding-3-large"
openai_client = OpenAI(
  api_key = os.getenv("OPENAI_API_KEY")
)

# ChromaDBの保存先をローカルファイルに指定して、永続化する
chromadb_client = chromadb.PersistentClient(path="./chroma_db")
openai_ef = embedding_functions.OpenAIEmbeddingFunction(
  api_key = os.getenv("OPENAI_API_KEY"),
  model_name = embedding_model
)

# collectionという、SQLでいうテーブルを作成
collection = chromadb_client.get_or_create_collection(
  name = "test_collection",
  embedding_function = openai_ef
)

texts = [
  "私はオランダで生まれて、東京で育ちました",
  "趣味は自転車とポーカーです。",
  "AIについて勉強しています。",
  "休日はカフェで作業することが多いよ。",
  "猫を飼っていて、まろんという名前です",
  "2年前に結婚しました",
  "旅行が好きで、先月はシリコンバレーに行きました",
  "英語の勉強をしてて、もっと話せるようになりたいな",
  "赤い色が好きです!",
  "この前自転車を盗まれて、新しい自転車を買いました。悲しい…",
  "最近夜更かししてるから早く寝ないとなあ",
]

# ベクトル化してcollectionに挿入
collection.upsert(
  documents = texts, 
  ids = [f"doc_{i}" for i in range(len(texts))]
)

コンテキストの抽出

ここから後半、ユーザーから質問が入力された後の処理について解説していきます。
コードでいうと rag_answer の関数の中身です。

今回は、ChromaDBに対してユーザーの質問を渡して検索し、意味が近いテキストを3つ出力させています。

実際のコードとしてはシンプルで、collectionに対して query関数を使うことで検索ができます。

response = collection.query(
    query_texts = [question],
    n_results = 3,
)   
    
contexts = []
for i in range(3):
    contexts.append(response['documents'][0][i])

例えば、「好きなことを教えて」という質問をすると、 contexts には

趣味は自転車とポーカーです
英語の勉強をしてて、もっと話せるようになりたいな
赤い色が好きです!

の3つが入っています。

コンテキストをLLMに渡そう!

あとは、このコンテキストをプロンプトに埋め込んで、LLMに投げるだけです。
該当するコードはこちら

  prompt = f"""
    以下の情報をもとに質問に答えてください。
    あなたは厳密なリサーチアシスタントです。
    以下の「参照メモ」に書かれている内容【だけ】を根拠に、日本語で簡潔に答えてください。
    参照に無い内容は推測せず、「手元のメモには情報がありません」と述べてください。

    参考メモ: {contexts}
    質問: {question}
  """

  chat_response = openai_client.chat.completions.create(
    model="gpt-4o",
    messages=[
        {"role": "system", "content": "あなたは厳密なリサーチアシスタントです。"},
        {"role": "user", "content": prompt}
    ]
  )

  return chat_response.choices[0].message.content

実際にLLMに投げられるプロンプトとしてはこのような形になっています。

以下の情報をもとに質問に答えてください。
あなたは厳密なリサーチアシスタントです。
以下の「参照メモ」に書かれている内容【だけ】を根拠に、日本語で簡潔に答えてください。
参照に無い内容は推測せず、「手元のメモには情報がありません」と述べてください。

参考メモ:['趣味は自転車とポーカーです。', '英語の勉強をしてて、もっと話せるようになりたいな', '赤い色が好きです!']
質問: 好きなことを教えて

これにより、LLMが本来知らない情報をコンテキストに加えて回答を生成させることができました。

まとめ

今回だと、texts が11行くらいしかないので、原文のままプロンプトに混ぜてLLMに渡しても同じ結果になると思います。
ただ、社内ドキュメントやNotionの情報など、情報が膨大になると、毎回プロンプトに全て埋め込むわけにはいきません。
なので、今回は練習として、VectorDBに意味検索をさせるプロセスを踏んでみました。

今回はフレームワークなどを利用せずにやっていましたが、LangChainなどを使えば、もっと綺麗に書くことができます。
ちょうどLangChainを使ってRAGを実装するチュートリアルもあるので、どこかのタイミングでこちらの解説もやってみたいですね!
https://docs.langchain.com/oss/python/langchain/rag

Discussion