DatabricksとOllamaによるLLMアプリの実験と検証 - CCCMKホールディングス TECH LABの Tech Blog

こんにちは、CCCMKホールディングスTECH LABの三浦です。

だいぶ気温が下がって、空気が乾燥してきたように感じます。すぐに手がしもやけになってしまうので、外に出る時は手袋をするようにしないと・・・と思います。

はじめに

先日Databricksの年次のイベント"Data + AI World Tour 2024 Tokyo"に参加しました。他社様のLLMを使った取り組みを聞くことが出来て、とても刺激になりました。特にLLMの動作の検証は単に「いい感じの結果が生成された」という感覚的な評価でなく、定量的な評価もしている、という他社様の取り組みを聞き、自身の今のLLMへの取り組み方を見直さないと、と考えさせられました。

そのようなきっかけから、最近自分が取り組んでいるオープンソースLLMを使ったアプリケーションの実験検証環境をDatabricksで作ってみたいな、と考えました。

オープンソースLLMは、Ollamaというアプリケーションを使って動かしています。Ollamaは量子化されたLLMモデルファイルを読み込んで、少ないリソース上でも大きなLLMを高速に稼働させることが出来るアプリケーションです。

ollama.com

私が最近取り組んでいたLLMアプリケーションはRAG(Retrieval-Augmented Generation)を使った情報検索アプリケーションなのですが、特にRetriever(情報検索)の部分の検索精度に課題を感じていました。精度向上に向けたアプローチとして、次の2つを考えていました。

ベクトルデータベースに格納する時のデータ形式
使用するデータは表形式のデータで、そのうち1つのカラムにテキスト情報が格納されています。このテキスト情報だけ埋め込みベクトルを求め、検索対象にすべきか、あるいは1行まるごとjson文字列にして埋め込みベクトルにすべきか？
アプリケーションへの入力に対する最適な検索クエリ
アプリケーションへのユーザーの入力でそのまま検索をかける方法がネイティブなRAGですが、入力が検索クエリ向きではないこともあります。そのため、ユーザーの入力をLLMを使って検索クエリ用に書き換える方法を取りたいのですが、具体的にどのようにしたらいいのか？(プロンプトの探索もありますが、今回はどのLLMを使うべきか、というアプローチで対応。)

これを実際に試し、かつ定量的に評価しようとすると準備に結構手間がかかります。

Ollamaを動かすサーバを用意する。
評価用のデータを準備する。
実験をし、結果をどこに記録するのか。

といったところが特に手間がかかりそうです。今回はこれらの準備を全てDatabricksで行い、実験と検証まで試してみました。個人的にはかなりいい感じに出来たな、と思っているのでぜひこの記事でご紹介させていただきます！

実験内容

検証したいこと

今回はRAGのRetriever部分の精度向上のための実験と検証です。検索対象のデータはユーザーレビューを含む商品情報で、この中からユーザーのインプットにマッチした商品を検索出来るようにしたいです。そのためにどのようにデータを持つべきか、どう検索を行うかが検証したい内容になります。

データの持ち方
- レビュー本文だけ埋め込んでベクトル検索の対象にする
- レビュータイトル＋レビュー本文を埋め込んで検索対象にする
- 商品情報全て埋め込んで検索対象にする
検索の仕方
- ユーザーのインプットをそのまま検索クエリに使う
- "llama3.2:1b"を使ってインプットを検索クエリに書き換え検索する
- "llama3.2:3b"を使ってインプットを検索クエリに書き換え検索する
- "gemma2:9b"を使ってインプットを検索クエリに書き換え検索する

これらの組み合わせで決まる計12パターンのRetrieverの処理で、最も精度がいい組み合わせを探したい、ということが今回の最終ゴールです。

使用するデータ

検索対象のデータはテキスト型のユーザーレビュー項目を含む商品リストで、表形式です。データはgpt-4oを使って今回の実験の用途のために生成しました。

表は以下の項目で構成されています。

カラム名	説明
review_date	レビューが投稿された日時(yyyymmdd HH:mm:ss形式)
category	レビュー対象の商品が属する商品カテゴリ
product_name	レビュー対象の商品の名前
review_title	レビューのタイトル
review_text	日本語50文字数以上の長さのレビュー本文
cateratingory	レビューの評価(1-5)

検証方法

テスト用のユーザーインプットデータに対し、商品を合計3件検索します。ユーザーのインプットと検索された商品をgpt-4oに与え、関連度を示すスコアを0.0～5.0の範囲で出力させます。

ユーザーインプットは100件用意し、それぞれで求めたスコアを平均し、対象のRetrieverの処理の評価値とします。

Notebookの構成

今回の実験のため、4つのNotebookを作成しました。

Notebook	説明
run_ollama_server	OllamaをインストールしてOllamaサーバを起動します。その後必要なLLMをPullするNotebook。
create_dataset	Azure OpenAI Serviceのgpt-4oを使って実験に使うデータセットを作成し、評価時に使用するテスト用インプットデータを作成するNotebook。
create_vectordb	データセットをベクトルデータベースに格納するNotebook。検証項目にデータのどこを埋め込み化したらよいのかが含まれているため、埋め込み方法ごとにインデックス(RDBにおけるテーブルに当たる)を作成しました。
execute_experiment	検証項目に応じた処理を実行し、評価を行い、Databricksに統合されているMLflowに実験結果を記録するNotebook。

Notebook詳細

run_ollama_server

OllamaサーバをComputeの中に立ち上げるNotebookです。

Computeが起動している間は、同じComputeで起動した他のNotebookからもこのサーバにアクセス出来るようになります。

Ollamaサーバをインストールし、起動するコマンドをシェルコードで実行しました。

%sh
curl -fsSL https://ollama.com/install.sh | sh

検証に使用するLLMをダウンロードします。埋め込み計算に使用するLLMはHugging Faceのモデルを使用しました。

huggingface.co

このモデルを量子化したこちらのファイルを使用しました。

huggingface.co

最近Hugging Faceに公開されているGGUFファイル(LLMを量子化したファイル)を直接Ollamaで取り込めるようになりました。この方法に従って埋め込みモデルをOllamaに取り込みました。

huggingface.co

%sh
ollama pull hf.co/nnch/multilingual-e5-large-Q4_K_M-GGUF:Q4_K_M
ollama pull llama3.2:1b
ollama pull llama3.2:3b
ollama pull gemma2:9b

create_dataset

実験に使うデータをAzure OpenAI Serviceのgpt-4oを使って生成しました。何か決まった形式の出力をLLMにさせる場合はLangChainの"StructuredOutput"を使用するととても便利です。ただしLLMはTool Callingに対応している必要があります。

具体的には、まず以下の様に出力したい形式をPydanticのBaseModelを使って定義します。レビュー1件の情報を定義するReviewを構成した後、それを複数持つデータセットとしてReviewDatasetというクラスを定義しました。

from pydantic import BaseModel, Field

class Review(BaseModel):
  review_date: str = Field(description="レビューが投稿された日時(yyyymmdd HH:mm:ss形式)")
  category: str = Field(description="レビュー対象の商品が属する商品カテゴリ")
  product_name: str = Field(description="レビュー対象の商品の名前")
  review_title: str = Field(description="レビューのタイトル")
  review_text: str = Field(description="日本語50文字数以上の長さのレビュー本文")
  rating: float = Field(description="レビューの評価(1-5)")

class ReviewDataset(BaseModel):
  reviews: list[Review] = Field(description="レビューのリスト")

その後、使用するLLMを定義して、with_structured_outputを呼び出してLLMとクラスの紐づけをします。

from langchain_openai import AzureChatOpenAI

llm = AzureChatOpenAI(    
    model="gpt-4o",
    api_version="2024-02-15-preview"
)

dataset_creater = llm.with_structured_output(ReviewDataset)

あとは通常のLLMと同様にプロンプトを与えて回答を生成させることが出来ます。

from langchain_core.prompts  import PromptTemplate

prompt = PromptTemplate.from_template(
  "商品カテゴリ'{category}'の商品レビューデータを25件作成して"
)

dataset_creater_chain = prompt|dataset_creater

categories = ["小説", "コミック", "ライトノベル", "エッセイ"]

datasets = []

for category in categories:
  dataset = dataset_creater_chain.invoke({"category":category})
  datasets.extend(dataset.reviews)

表形式にするため、Pandasを使って整形しました。

import pandas as pd

datasets_df = pd.DataFrame()

for data in datasets:
  data_df = pd.DataFrame(
    {
      "review_date": [data.review_date],
      "category": [data.category],
      "product_name": [data.product_name],
      "review_title": [data.review_title],
      "review_text": [data.review_text],
      "rating": [data.rating],
    }
  )
  datasets_df = pd.concat([datasets_df, data_df])

以下のようなデータを生成することが出来ました。

create_vectordb

生成したデータセットをベクトルデータベースに格納するためのNotebookです。

埋め込みのやり方を3パターン検証するため、3つのインデックス(RDBのテーブルに該当する)を生成しました。ベクトルデータベースは普段Chromaを使って作ることが多いのですが、Chromaのデータの永続化の仕様が理解しきれず、今回は別のライブラリを使ってベクトルデータベースを構築することにしました。Chromaと同じくポピュラーな"Faiss(Facebook AI Similarity Search)"です。

github.com

必要なライブラリのインストールをNotebookで実行しました。

%pip install chromadb "langchain<0.3"
%pip install -qU langchain-community faiss-cpu langchain-ollama pydantic
dbutils.library.restartPython()

使用する埋め込みモデルを呼び出します。Ollamaで稼働しているものを使用します。

from langchain_ollama import OllamaEmbeddings

model_name = "hf.co/nnch/multilingual-e5-large-Q4_K_M-GGUF:Q4_K_M"
embedding = OllamaEmbeddings(model=model_name)

レビュータイトルとレビュー本文を埋め込んだインデックスの作成処理は次のようにしました。

import json
from uuid import uuid4

import faiss
from langchain_community.docstore.in_memory import InMemoryDocstore
from langchain_community.vectorstores import FAISS
from langchain_core.documents import Document


# インデックスの生成
index = faiss.IndexFlatL2(len(embedding.embed_query("hello world")))
vector_store = FAISS(
    embedding_function=embedding,
    index=index,
    docstore=InMemoryDocstore(),
    index_to_docstore_id={},
)

# Pandas DataFrameのRow(Series)をDocumentに変換する処理を関数にまとめる。
def create_document(row_num: int, row: pd.Series):
  content_json = {
    "review_title":row["review_title"], 
    "review_text":row["review_text"]
  }
  content = json.dumps(content_json, ensure_ascii=False)
  return Document(
    page_content=content,
    metadata={
      "review_date": row["review_date"],
      "category": row["category"],
      "product_name": row["product_name"],
      "review_title":row["review_title"], 
      "review_text":row["review_text"],
      "rating": row["rating"]
    },
    id=row_num,
  )

# Pandas DataFrameをDocumentのリストに変換
documents_for_db = []
for i, row in dataset_df.iterrows():
  documents_for_db.append(create_document(i, row))

# ID列の生成
ids = [str(uuid4()) for _ in range(len(documents_for_db))]

# Documentの登録
vector_store.add_documents(
  documents=documents_for_db, ids=ids
)

# インデックスの保存
vector_store.save_local("./vector_db/review_title")

同様にレビュー本文だけ、商品情報全てを埋め込んだインデックスを生成し、保存しておきました。

execute_experiment

これで準備が出来たので、いよいよ実験を行います。

まず実験を記録するMLflowのExperimentを作成しました。DatabricksのMLflowのExperiment名は、"/Users/{自身のメールアドレス}/{Experimentの名前}"で構成されます。

import mlflow
mlflow.create_experiment("/Users/{mailaddress}/retriever_exp")

次にユーザーのインプットと検索結果の商品との関連度スコアを計算する処理を作りました。 gpt-4oを内部で呼び出します。ここでもLangChainの"StructuredOutput"を使って実装しました。

from pydantic import BaseModel, Field

from langchain_core.prompts import PromptTemplate
from langchain_openai import AzureChatOpenAI

class JudgeScore(BaseModel):
  score: float = Field(
    description="クエリと検索結果の関連度(0.0～5.0までの数値)",
    ge=0.0, 
    le=5.0
  )

llm = AzureChatOpenAI(    
    model="gpt-4o",
    api_version="2024-02-15-preview",
    temperature=0.0
)


prompt = PromptTemplate.from_template(
"""
与えられたクエリと、検索された関連商品の類似度を計算してください。
類似度は0.0～5.0までの数値で出力してください。

# クエリ
{query}

# 関連商品
{products}
"""
)
judgement = prompt|llm.with_structured_output(JudgeScore)

正しく機能するか試してみました。まずユーザーのインプットとして次のように与えました。

"心が温かくなる物語"

それに対する関連商品として、あえてインプットの内容と異なる商品情報を与えてみました。

{"review_date": "20231001 14:23:45", "category": "小説", "product_name": "幻想の城", "review_title": "驚くべき結末", "review_text": "物語の展開が素晴らしく、最後まで飽きずに読み進めることができました。特に、結末には驚かされました。", "rating": 5.0}

関連スコアは2.0で、関連度が低い、と見なしていることが分かります。

商品情報はそのままに、インプットを"展開に驚かされる物語"に変えて実行すると、5.0という高いスコアが出力されました。判定処理は上手く機能していることが確認出来ました。

続いて実験に関連する設定を進めました。まず動かすパラーメータを定義しました。このパラメータの組み合わせを総当たりで実行していき、最も高いスコアを出した組み合わせを見つけ出します。

rewrite_query = """次のユーザーの入力に対し、情報を検索するために最適なクエリを生成し、クエリだけを出力してください。
入力: {input}
クエリ："""

parameters = {
  "indexs":[
    "./vector_db/all",
    "./vector_db/review",
    "./vector_db/review_title"
  ],
  "rewrite_model":[
    None,
    "llama3.2:1b",
    "llama3.2:3b",
    "gemma2:9b"
  ],
  "rewrite_query":[rewrite_query]
}

使用するインデックス、クエリ書き換えモデルによってRetrieverのパイプラインを都度構築する必要があるため、処理を関数にまとめておきました。

def create_chain(index, rewrite_model, rewrite_query):
  vector_store = FAISS.load_local(
    index, 
    embeddings=embedding, 
    allow_dangerous_deserialization=True
  )
  retriever = vector_store.as_retriever(search_kwargs={"k": 3})
  
  if rewrite_model is not None:
    # クエリの書き換えをする場合
    llm = ChatOllama(model=rewrite_model)
    prompt = PromptTemplate.from_template(rewrite_query)
    chain = prompt|llm|StrOutputParser()|retriever
  else:
    # 書き換えをしない場合はそのまま検索クエリとして渡す。
    chain = itemgetter("input")|retriever
  return chain

そして次が実験のメイン処理です。

import json
import time

import mlflow
import numpy as np

mlflow.set_experiment("/Users/{email}/retriever_exp")

# パラメータを総当たりで取得できるようにする。
param_combination = itertools.product(
  parameters["indexs"], 
  parameters["rewrite_model"],
  parameters["rewrite_query"]
)

for index, rewrite_model, rewrite_query in param_combination:

  with mlflow.start_run():
    # MLflow Experimentへの記録(Run)の開始
    chain = create_chain(index,rewrite_model,rewrite_query)

    print("Start Retrieval Chain.")
    start_time = time.time()
    # バッチで実行する。
    results = chain.batch([{"input":query} for query in test_queries_ls])
    end_time = time.time()
    exec_time = end_time - start_time
    print(f"End Retrieval Chain.time: {exec_time}s")
    
    result_items = []
    for result in results:
      # 取得された関連商品情報を文字列に変換する。
      result_items_per_query = [
        json.dumps(r.metadata,ensure_ascii=False) for r in result
      ]
      # 3つ関連商品が取得されるので、"\n"で改行する。
      result_items_per_query_str = "\n".join(result_items_per_query)
      result_items.append(result_items_per_query_str)
    
    # 関連スコアを計算する
    print("Start calculating score.")
    scores = judgement.batch(
      [
        {
          "query": test_query,
          "products": test_sample
        } 
        for test_query, test_sample in zip(test_queries_ls, result_items)
      ])
    scores = [s.score for s in scores]
    print("End calculating score.")

    # MLflowに結果を記録する
    print("Start logging.")
    mlflow.log_table(
      data={
        "query": test_queries_ls,
        "result": result_items,
        "relevant_score": scores
      },
      artifact_file="query_result_score_table.json"
    )

    # 実行パラメータの記録
    mlflow.log_params(
      {
        "index": index,
        "rewrite_model": rewrite_model,
        "rewrite_query": rewrite_query
      }
    )

    # 関連スコアの平均の標準偏差の記録
    mlflow.log_metric("relevant_score_mean", np.mean(scores))
    mlflow.log_metric("relevant_score_std", np.std(scores))
    # 検索にかかった時間の記録
    mlflow.log_metric("exec_time", exec_time)
    print("End logging.")