SwiftでRAG実装 Part 2：クエリに類似するドキュメント検索の試み

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 }
}