SwiftでRAG実装 Part 1：テキストをコンテキストベクトル化 の続き。
前回はドキュメントデータから埋め込み生成。今回はそこから検索クエリに類似するドキュメントを抽出してみる。先に結論を言うと、結果はいまひとつなので調整が必要そう。
今回実装する計算ロジックは以下（現状の筆者の理解であることに注意）。
- ドキュメントからそれぞれ埋め込み生成し、ベクトル化する
- ドキュメントごとのベクトル表現(D次元)を、ドキュメントの数(M個)分並べて MxD 行列を生成する（ドキュメント行列）
- クエリ文字列をベクトル化する。1. と同じ手法で生成しベクトル次元は一致(D次元)。（クエリベクトル）
ドキュメント行列・クエリベクトル（内積）を計算する (MxD・Dx1 = Mx1)- つまりドキュメントごとのベクトルとクエリベクトルとの内積が結果として求まる。内積ベクトルのノルム(コサイン類似度)が1に近いほど類似度高く、0、-1に近づくほど類似度が低いと判定できる
- 4. の結果で得られた類似度の配列(要素数M)を降順にソート、上から任意の件数を上位ヒットとする
var docTensors: [MLTensor] = ...
// docTensors に含むドキュメントごとのクエリとの類似度を検索し、上位任意件数(maxCount)をヒット結果として返す
// (返すのは該当文書テンソルの docTensors 配列内におけるインデックス番号)
func search(query: String, maxCount: Int) async -> [Int] {
// ドキュメントを集積した MxD 行列
let flatteneds = tensors.map { $0.flattened() }
let docsMat = MLTensor(stacking: flatteneds)
// クエリベクトル Mx1 (encode関数は前回記事参照)
guard let queryVec = try? embedder.encode(text: query, asColumn: true /*列ベクトルで出力*/) else { return [] }
// ドキュメント行列・クエリベクトル (MxD・Dx1 = Mx1)
let mulResult = docsMat.matmul(queryTensor)
// 類似度スコア配列に変換
let calcScores = await mulResult.shapedArray(of: Float.self).scalars
// ソートして、トップの結果をインデックス番号として抽出
let map = zip(Array(0..<docTensors.count), calcScores)
let sorted = map.sorted { $0.1 > $1.1 }.prefix(maxCount)
return sorted.map { $0.0 }
}これでサンプルの日記データを検索したところ、肌感としてあまり精度高くなかった。
ただし、クエリの何かしらの特徴は反映していそうだった。たとえば「I watched a movie」のクエリで検索すると、映画を見た日記エントリはヒットしないが、上位ヒットは軒並み「I (動詞過去形)」のパターンで始まっていたり、更に分かりやすい例を挙げると、次のように「at a cafe」クエリに対し「At a/the」で始まるエントリが抽出されたりした。まったくの出鱈目ではなさそうだが、肝心の文章的なコンテキストは明らかに落ちているように見える。
追記（2025/10/01）：後日性能改善できたので後々投稿するが、取り急ぎこちらの資料にソースコードを添付した。（enumerateTokenVectors を使うようにしただけ）
