こんにちは、CCCMKホールディングス AIエンジニアの三浦です。気温が急に高くなって、まるで夏が来たみたいです。朝と夜はまだ涼しいので、出かける時に何を着ていくのか悩んでしまいます。
- Cortex Search
- コードの実行環境と検証に使用したデータ
- ライブラリのインストール
- 翻訳処理の事前準備
- データセットのダウンロードとサンプリング
- 翻訳処理の実行
- SnowflakeのTableに書き込む
- Cortex Searchを作成する
- Cortex Searchにクエリを実行する
- まとめ
Cortex Search
Cortex SearchはSnowflakeのセマンティック検索を実現するための機能です。セマンティック検索は自然言語処理を利用した、検索クエリと検索対象データの意味の近さを考慮した検索方法です。最近ではLLMの応用方法として有名なRAGの情報検索処理において欠かすことが出来ない要素となっています。
SnowflakeのCortex Searchは、Snowflakeにすでに格納されているTableから作ることが出来ます。検索はハイブリッド検索で行われ、埋め込みモデルによる埋め込み表現の近さによる検索とキーワード一致の検索の処理が走り、最後にそれぞれの結果とクエリとの関連性を評価し(ReRanking)、最も関連性の高い情報を取得することが出来るようになっています。埋め込みモデルはSnowflakeでホストされた複数のモデルから選択することが可能です。
今回の記事では、まずCortex Searchを作るところからはじめ、最終的にCortex SearchにPythonからクエリを実行するところまでまとめています。
コードの実行環境と検証に使用したデータ
コードの実行は、Azure DatabricksのNotebookから行いました。また、検証に使用したデータはHugging Faceに公開されている商品レビューデータセットです。
データをダウンロードしてから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分ほどかかりました。実行した結果を見てみると、以下のように日本語への翻訳が行われていました。
最終的に以下のようなデータセットを作成しました。
SnowflakeのTableに書き込む
完成したデータセットをSnowflakeのTableに書き込みます。Tableに書き込む前にカラム名を小文字表記から大文字表記に変換する処理を実行しています。結構重要なポイントで、Snowflakeに書き込むデータのカラム名に小文字が含まれていたりスペースが含まれていると、カラム名にダブルクオーテーションが付与されて登録されてしまいます。
この状態で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を作る際に必要になる設定です。以下のドキュメントを参照しています。
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件とは言えレスポンスがめちゃくちゃ速くてびっくりしました。
まとめ
今回はSnowflakeのセマンティック検索機能を提供するCortex Searchの作成からクエリの実行までを試してみました。対象のTableのカラムに小文字が含まれていることによってカラム名にダブルクオーテーションが付与されてしまい、Cortex Searchが作れずに詰まってしまうところもあったのですが、それ以外はスムーズに進めることが出来ました。これでRAGの仕組みをSnowflakeで作ることが出来そうです。