CCCMKホールディングス TECH LABの Tech Blog

TECH LABのエンジニアが技術情報を発信しています

ブログタイトル

Snowflake Cortex Searchでセマンティック検索!

こんにちは、CCCMKホールディングス AIエンジニアの三浦です。気温が急に高くなって、まるで夏が来たみたいです。朝と夜はまだ涼しいので、出かける時に何を着ていくのか悩んでしまいます。

Cortex SearchはSnowflakeのセマンティック検索を実現するための機能です。セマンティック検索は自然言語処理を利用した、検索クエリと検索対象データの意味の近さを考慮した検索方法です。最近ではLLMの応用方法として有名なRAGの情報検索処理において欠かすことが出来ない要素となっています。

SnowflakeのCortex Searchは、Snowflakeにすでに格納されているTableから作ることが出来ます。検索はハイブリッド検索で行われ、埋め込みモデルによる埋め込み表現の近さによる検索とキーワード一致の検索の処理が走り、最後にそれぞれの結果とクエリとの関連性を評価し(ReRanking)、最も関連性の高い情報を取得することが出来るようになっています。埋め込みモデルはSnowflakeでホストされた複数のモデルから選択することが可能です。

今回の記事では、まずCortex Searchを作るところからはじめ、最終的にCortex SearchにPythonからクエリを実行するところまでまとめています。

コードの実行環境と検証に使用したデータ

コードの実行は、Azure DatabricksのNotebookから行いました。また、検証に使用したデータはHugging Faceに公開されている商品レビューデータセットです。

huggingface.co

データをダウンロードしてからSnowflakeのCortex Searchで利用出来るようにするまでにした作業は次の通りです。

まずデータセット全体から1,000件のサンプリングを行い、レビューのタイトルとテキストについて、英語から日本語に翻訳する前処理を実施しました。翻訳はDatabricksのFoundation Model APIで提供されているLLM"databricks-claude-3-7-sonnet"を利用して行いました。

その後、加工済みのデータを一度SnowflakeのTableとして作成し、そのTableをベースにCortex Searchを作成する、という流れを実施しました。

ライブラリのインストール

必要なライブラリのインストールを、以下のコマンドで行いました。

%pip install snowflake-core snowflake-connector-python databricks-langchain snowflake-snowpark-python
dbutils.library.restartPython()

翻訳処理の事前準備

英語で書かれたレビューのタイトルと本文を日本語に翻訳する処理を、LangChainのChainで定義します。翻訳したデータを取り出しやすくするように、"StructuredOutput"を使って処理結果を出力しています。StructuredOutputを使うと、自由なフォーマットのテキストではなく 事前に定めておいたスキーマ形式に従ってLLMから出力を得られるようになります。

from pydantic import BaseModel

from databricks_langchain.chat_models import ChatDatabricks
from langchain_core.prompts import PromptTemplate

class Ja_Title_Text(BaseModel):
    title: str
    text: str

prompt_template = """あなたはAmazonのレビューを日本語に翻訳するアシスタントです。
レビューのタイトルと本文テキストをそれぞれ日本語に翻訳してください。

title: {title}
text: {text}"""

prompt = PromptTemplate.from_template(prompt_template)
llm = ChatDatabricks(endpoint="databricks-claude-3-7-sonnet")\
        .with_structured_output(Ja_Title_Text)
chain = prompt|llm

データセットのダウンロードとサンプリング

Hugging Faceからデータセットをダウンロードします。ダウンロードしたのち、1,000件のサンプリングを行いました。

from datasets import load_dataset
import pandas as pd

dataset = load_dataset(
    "McAuley-Lab/Amazon-Reviews-2023", 
    "raw_review_All_Beauty", 
    trust_remote_code=True
)

sample = dataset["full"].shuffle(seed=42)[:1000]

# pandas.DataFrameに変換
sample_df = pd.DataFrame.from_dict(sample)

翻訳処理の実行

先ほど定義したLangChainの翻訳Chainを使ってデータセットに対して翻訳処理を実行します。

def translate_title_text(title, text):
    ja_title_text = chain.invoke({"title": title, "text": text})
    return ja_title_text

title_text_ja_lists = [
    translate_title_text(title, text) \
        for title, text in zip(sample_df["title"], sample_df["text"])
]

処理時間はだいたい1時間15分ほどかかりました。実行した結果を見てみると、以下のように日本語への翻訳が行われていました。

databricks-claude-3-7-sonnetによる英語→日本語翻訳結果

最終的に以下のようなデータセットを作成しました。

SnowflakeのTableに書き込む

完成したデータセットをSnowflakeのTableに書き込みます。Tableに書き込む前にカラム名を小文字表記から大文字表記に変換する処理を実行しています。結構重要なポイントで、Snowflakeに書き込むデータのカラム名に小文字が含まれていたりスペースが含まれていると、カラム名にダブルクオーテーションが付与されて登録されてしまいます

自動的に生成された"CREATE TABLE"ではカラム名が小文字だとすべてダブルクオートで囲まれている

この状態でCortex Searchを作ろうとすると、ProgrammingError: 399115 (42601): Invalid source query: quoted identifier or reserved word "ja_title" not allowed.のようなエラーが発生し、Cortex Searchの作成に失敗してしまいます。この現象を回避するため、Tableに書き込む前にカラム名をすべて大文字に変換する処理を実行しました。

import snowflake.connector
from snowflake.connector.pandas_tools import write_pandas

# 小文字のカラムをすべて大文字に変換する
sample_df.columns = map(lambda x: str(x).upper(), sample_df.columns)

with snowflake.connector.connect(**connection_info) as conn:
    write_pandas(
        conn, 
        sample_df, 
        'REVIEW_DATA',
        quote_identifiers=False,
        overwrite=True
    )
    cur = conn.cursor()
    cur.execute("ALTER TABLE REVIEW_DATA SET CHANGE_TRACKING = TRUE;")

最後のALTER TABLEは、Cortex Searchを作る際に必要になる設定です。以下のドキュメントを参照しています。

docs.snowflake.com

Cortex Searchを作成する

先ほど登録したTableをベースにCortex Searchを作成します。

with snowflake.connector.connect(**connection_info) as conn:
    cur = conn.cursor()
    cur.execute(\
        """
        CREATE OR REPLACE CORTEX SEARCH SERVICE REVIEW_SEARCH
        ON JA_TEXT
        ATTRIBUTES RATING
        TARGET_LAG = '1 day'
        WAREHOUSE = WH
        EMBEDDING_MODEL = 'snowflake-arctic-embed-l-v2.0'
        AS (
            SELECT
                JA_TITLE,
                JA_TEXT,
                RATING
            FROM REVIEW_DATA
        );
        """)
  • CREATE OR REPLACE CORTEX SEARCH SERVICE REVIEW_SEARCHで"REVIEW_SEARCH"という名前のCortex Searchを作成します。
  • ON JA_TEXTによって"JA_TEXT"というカラムを埋め込み、検索対象に指定しています。
  • ATTRIBUTES RATINGによって"RATING"カラムで検索時にフィルタがかけられるようになります。
  • TARGET_LAGはベースのTableの更新をどれくらいの頻度で確認するのかを指定します。
  • EMBEDDING_MODELは埋め込みに使うモデルです。

処理完了後、Cortex Searchにクエリを実行できるようになります。

Cortex Searchにクエリを実行する

Cortex SearchにPythonからクエリを実行するには、たとえば以下のようなコードで可能です。ここでは検索クエリとして"さっぱり感"を指定し、"RATING"が3.0以上のものだけを取得するようフィルタリングをかけています。検索結果は5件返ってくるようにしました。

from snowflake.core import Root
from snowflake.snowpark import Session

session = Session.builder.configs(connection_info).create()
root = Root(session)

# fetch service
my_service = (root
  .databases[DATABASE_NAME]
  .schemas[SCHEMA_NAME]
  .cortex_search_services["REVIEW_SEARCH"]
)

# query service
resp = my_service.search(
  query="さっぱり感",
  columns=["JA_TITLE", "JA_TEXT","RATING"],
  filter={"@gte": {"RATING": 3.0} },
  limit=5
)
pd.DataFrame.from_dict(resp.to_dict()["results"])

結果は以下のようになりました。なかなか良さそうな感じです!また、1,000件とは言えレスポンスがめちゃくちゃ速くてびっくりしました

Cortex Searchから得られた結果

まとめ

今回はSnowflakeのセマンティック検索機能を提供するCortex Searchの作成からクエリの実行までを試してみました。対象のTableのカラムに小文字が含まれていることによってカラム名にダブルクオーテーションが付与されてしまい、Cortex Searchが作れずに詰まってしまうところもあったのですが、それ以外はスムーズに進めることが出来ました。これでRAGの仕組みをSnowflakeで作ることが出来そうです。