概要
Neptune AnalyticsとLlama-Indexライブラリを使用し、ベクトル検索を行おうとした際に、思うような挙動をしなかったため修正した概要を記載
各種リンク
※実行環境はSageMaker Notebookですが、環境構築などは別の記事をご参照ください。
目次
- 現状
- 課題
- 改善点
- 修正後
現状
目的
クエリキーワードに近い意味の言葉のノードを検索する
要点
- Llama-Indexのプロパティグラフを用いてグラフを構築
- 検索時は Retrieverである
VectorContextRetrieverを使用 -
VectorContextRetrieverで検索を行うと、最終的に以下ファイルのvector_query関数でNeptune Analyticsに対してクエリを実行する
下記コードはvector_query関数のクエリ実行箇所
data = self.structured_query(
f"""
MATCH (e:`{BASE_ENTITY_LABEL}`)
WHERE ({filters})
CALL neptune.algo.vectors.get(e)
YIELD embedding
WHERE embedding IS NOT NULL
CALL neptune.algo.vectors.topKByNode(e)
YIELD node, score
WITH e, score
ORDER BY score DESC LIMIT $limit
RETURN e.id AS name,
[l in labels(e) WHERE l <> '{BASE_ENTITY_LABEL}' | l][0] AS type,
e{{.* , embedding: Null, name: Null, id: Null}} AS properties,
score""",
param_map={
"embedding": query.query_embedding,
"dimension": len(query.query_embedding),
"limit": query.similarity_top_k,
},
)
課題
上記クエリの内容はざっと以下の挙動となる
引数
- embedding：クエリキーワードをベクトル化したもの
- dimension：embeddingの次元数
- limit：ベクトル検索結果の取得件数（SQLのlimitのようなもの）
挙動
- ノードを取得し、各ノードのベクトルを取得
- 取得した各ノードのベクトル値と近いベクトルを持つノードを検索し、ノードとスコアを取得
- 全体のスコアを降順に整列
- 結果の整形
課題点
- 引数の
embeddingが使用されていない -
neptune.algo.vectors.topKByNode(e)の使用が不適切 - スコアのソートが不適切
改善点
各課題の改善点は以下の通り
引数のembeddingが使用されていない
本来クエリパラメータをクエリ内で使用するには、$embeddingとして使用してあげる必要があるが、当該クエリでは一切使用されている様子が見られない。
neptune.algo.vectors.topKByNode(e)の使用が不適切
公式ドキュメントでは、下記引用のように記されている。（deeplによる翻訳）
.vectors.topKByNodeアルゴリズムは、ノードからのベクトル埋め込み距離に基づいて、ノードの上位K個の最近傍を見つける。
つまり、当該クエリにおいては、ベクトルを持つ各ノードと近いベクトルを持つノードとそのスコアを取得していることになる。
この挙動は今回の目的からは逸脱してしまっているため、.vectors.topKByEmbeddingを使用する。
公式ドキュメントでは、下記のように述べられている。（deeplによる翻訳）
.vectors.topKByEmbeddingアルゴリズムは、埋め込みベクトルの距離に基づいて、埋め込みベクトルの上位K個の最近傍を見つける。
つまり、このアルゴリズムの引数にクエリパラメータの$embeddingを渡せば、上位K件のベクトルに近いノードを検索することが可能になる。
WITH $embedding AS embedding
.vectors.topKByEmbedding(embedding)
スコアのソートが不適切
公式ドキュメントでは、Neptune Analyticsのベクトル検索のスコアについて以下の様に述べられている(deeplによる翻訳)
.vectors.distanceアルゴリズムは，埋め込みに基づいて2つのノード間の距離を計算します．この距離は，埋め込みベクトルのL2ノルムの2乗です．
つまり、「スコアが小さいほど類似している」といえる。
そのため、当該クエリでは「昇順」が正しい。
修正後
上記改善点を踏まえると以下のクエリとなる。
注意
筆者はこのクエリに条件などを追加して動作確認しているため、このクエリ単体で動作するかは未検証です。
ご使用になる際は、別途動作確認をお願いします。
data = self.structured_query(
f"""
- MATCH (e:`{BASE_ENTITY_LABEL}`)
- WHERE ({filters})
- CALL neptune.algo.vectors.get(e)
- YIELD embedding
- WHERE embedding IS NOT NULL
- CALL neptune.algo.vectors.topKByNode(e)
- YIELD node, score
- WITH e, score
+ WITH $embedding AS embedded_query
+ CALL neptune.algo.vectors.topKByEmbedding(embedded_query)
+ YIELD embedding, node, score
+ WITH node, score
- ORDER BY score DESC LIMIT $limit
- RETURN e.id AS name,
- [l in labels(e) WHERE l <> '{BASE_ENTITY_LABEL}' | l][0] AS type,
+ ORDER BY score ASC LIMIT $limit
+ RETURN node.id AS name,
+ [l in labels(node) WHERE l <> '{BASE_ENTITY_LABEL}' | l][0] AS type,
e{{.* , embedding: Null, name: Null, id: Null}} AS properties,
score""",
param_map={
"embedding": query.query_embedding,
- "dimension": len(query.query_embedding),
"limit": query.similarity_top_k,
},
)