StreamlitとOpenAI APIを使用したRAGの実装

はじめに

人工知能（AI）の進化に伴い、情報検索と自然言語処理（NLP）の領域は新たな時代を迎えています。今回は、StreamlitとOpenAI APIを使用して、Retrieval-Augmented Generation (RAG) モデルを構築し、運用するプロセスを紹介します。

データセットについて

今回は例としてkaggleから取得したスターバックスデータセットをjson形式に変換したもの使用します。内容としては以下の通りで、ビバレッジとそれに対応する商品情報が格納されたものを使用していきます。
データセットリンク：https://www.kaggle.com/datasets/henryshan/starbucks

 {
    "Beverage_category": "Coffee",
    "Beverage": "Brewed Coffee",
    "Beverage_prep": "Tall",
    "Calories": 4,
    " Total Fat (g)": "0.1",
    "Trans Fat (g) ": 0.0,
    "Saturated Fat (g)": 0.0,
    " Sodium (mg)": 0,
    " Total Carbohydrates (g) ": 10,
    "Cholesterol (mg)": 0,
    " Dietary Fibre (g)": 0,
    " Sugars (g)": 0,
    " Protein (g) ": 0.5,
    "Vitamin A (% DV) ": "0%",
    "Vitamin C (% DV)": "0%",
    " Calcium (% DV) ": "0%",
    "Iron (% DV) ": "0%",
    "Caffeine (mg)": "260"
  },

技術スタック

以下のツールとライブラリを使用します。今回は、費用は最低限に抑えるためにオープンソースを中心に使用しており、かかるコストはOpenAI APIのみとなっています。

Streamlit: データサイエンスと機械学習アプリケーションの迅速な構築を支援するPythonライブラリです。
FAISS: 高速な類似性検索を提供するライブラリで、Facebook AI Researchにより開発されました。
Sentence Transformers: 文書を意味的に類似したベクトルに変換するライブラリです。
OpenAI API: 自然言語処理と生成のためのAPIで、GPTモデルへのアクセスに利用されます。

アプリケーションの構築プロセス

1. Streamlitでのユーザーインターフェースの構築

最初に、Streamlitを使用してユーザーインターフェースを設計します。APIキーの入力と質問の入力ができるようにします。

import streamlit as st

st.title('RAG with STARBUCKS GPT')

api_key = st.text_input("Enter your OpenAI API Key:", type="password")
question = st.text_input("Enter your question:")

2. データの準備と文書のベクトル化

次に、スターバックスデータセットを読み込み、Sentence Transformersを用いて文書をベクトル化します。

import json
from sentence_transformers import SentenceTransformer

# JSONファイルの読み込み
file_path = 'starbucks_data.json'
with open(file_path, 'r') as file:
    data = json.load(file)

# 文書生成
docs = [f"{d['Beverage_category']} {d['Beverage']} {d['Beverage_prep']}" for d in data]

# 文書をベクトル化
model = SentenceTransformer("sentence-transformers/all-MiniLM-l6-v2")
embeddings = model.encode(docs)

3. FAISSを使用した類似性検索

ベクトル化した文書をFAISSインデックスに追加し、ユーザーの質問に対する関連文書を検索します。

import faiss
import numpy as np

# FAISSインデックスの作成と保存
d = embeddings.shape[1]
index = faiss.IndexFlatL2(d)
index.add(np.array(embeddings).astype('float32'))

# 質問のベクトル化と関連文書の検索
question_embedding = model.encode([question], convert_to_tensor=True)
question_embedding = question_embedding.cpu().numpy()
_, I = index.search(question_embedding, 5)

related_docs = [docs[i] for i in I[0]]

4. OpenAI APIを利用したテキスト生成

関連文書をもとに、OpenAI GPTを使用してユーザーの質問に対する回答を生成します。

import openai

# 入力されたOpenAIのAPIキーの変数設定
openai.api_key = api_key

# 関連する文書を基にしてプロンプトを作成
prompt = question + "\n\n" + "\n\n".join(related_docs)

# テキスト生成
response = openai.ChatCompletion.create(
    model="gpt-3.5-turbo",
    messages=[
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": prompt}
    ]
)

st.write(response.choices[0].message['content'])

実際のアプリケーション画面では、以下のように質問すると
ベクトル化されたデータから質問に適した内容を検索し解答を生成します。

まとめ

この記事では、StreamlitとOpenAI APIを組み合わせてRAGモデルを実装しました。この手法は、任意のドメインのデータに対して適用可能で、形態素解析などの前処理を適切に行うことで、テキストデータに対応させられます。LLMとRAGを利用し、情報検索と知識の抽出を活用していきましょう。