ベクターストアを使って類似文章の検索

RAGを実現するための前段階として、テキスト埋め込みを使って類似の文章を引っ張ってきてそれを基に文章を作成する、という方法があります。
これをすることで、

社内用のデータだったり、
専門的な知識だったり、
最新の情報だったり、

を用いるような文章を作成することができるようになります。

なお、 RAGって何？ とか テキスト埋め込みの理屈を教えて？ というのはすっ飛ばしますので、どこかで調べてきてください。
（実はふんわりとしか理解していないので詳しく説明することができないことは内緒です）

テキスト埋め込みと類似文章の検索

「似ている文章を探してきて」ということを実現するために、どういうことをしているかというと、

各々の文章を単語単位に分解
検索される各々の文章で、文章内の単語同士を距離感に変換
単語同士の距離感を覚えておく
検索したい文章も単語同士の距離感（検索用）に変換
検索される文章の単語の距離感と検索したい文章の距離感が似ているものをピックアップ

ということをやっている、はず。
で、この 距離感への変換 というのをテキスト埋め込みと呼んでいる、と。
自分の理解はこんなところです。
（距離感やなくて距離やで、とかいうのも含めて）

まー、ふんわりとしか理解していないのがよくわかりますね。

使用するテキスト埋め込みについて

少し前に ruri-v3をOpenAI互換embedding APIで使うという記事を書きました。
今回はローカルに展開しているこのruri-v3をSemantic Kernelで使ってみます。

注意

この記事は Semantic Kernel が 1.49.0 の際に書いたものです。
Semantic Kernel は更新速度が非常に早く、記載の表記から変更されていることが非常によくあります。
実際にコーディングする際には確認の上、作成してください。

類似文章を検索するプログラム

サンプルはこちら。

比較的短めなので、まずは全文を掲載。

Program.cs


using Microsoft.Extensions.VectorData;
using Microsoft.SemanticKernel;
using Microsoft.SemanticKernel.Connectors.InMemory;
using Microsoft.SemanticKernel.Embeddings;
using System.Diagnostics;
using System.Text.Json;
using System.Text.Json.Serialization;

internal class Program
{
    private static async Task Main(string[] args)
    {
        #region 初期化
        const string SrcBase = "test"; // test / train / valid
        const string SrcFileName = $"./{SrcBase}.json";
        const string SrcUri = $"https://raw.githubusercontent.com/yahoojapan/JGLUE/refs/heads/main/datasets/jsts-v1.3/{SrcBase}-v1.3.json";
        const string CollectionName = $"{SrcBase}_datas";

#pragma warning disable SKEXP0010 // 種類は、評価の目的でのみ提供されています。将来の更新で変更または削除されることがあります。続行するには、この診断を非表示にします。
        var kernel = Kernel
            .CreateBuilder()
            .AddInMemoryVectorStore()
            .AddOpenAITextEmbeddingGeneration("cl-nagoya/ruri-v3-310m", "-", httpClient: new HttpClient()
            {
                BaseAddress = new Uri("http://localhost:7997")
            })
            .Build();
#pragma warning restore SKEXP0010 // 種類は、評価の目的でのみ提供されています。将来の更新で変更または削除されることがあります。続行するには、この診断を非表示にします。

#pragma warning disable SKEXP0001 // 種類は、評価の目的でのみ提供されています。将来の更新で変更または削除されることがあります。続行するには、この診断を非表示にします。
        var embedding = kernel.GetRequiredService<ITextEmbeddingGenerationService>();
#pragma warning restore SKEXP0001 // 種類は、評価の目的でのみ提供されています。将来の更新で変更または削除されることがあります。続行するには、この診断を非表示にします。

        var vectorStore = kernel.GetRequiredService<IVectorStore>();
        var collection = vectorStore.GetCollection<string, TrainData>(CollectionName);
        await collection.CreateCollectionIfNotExistsAsync();
        #endregion

        #region 元データダウンロード（無いとき）
        if (!File.Exists(SrcFileName))
        {
            using (var client = new HttpClient())
            {
                var res = await client.GetAsync(SrcUri);
                if (res.IsSuccessStatusCode)
                {
                    using (var dstf = File.OpenWrite(SrcFileName))
                    {
                        await res.Content.CopyToAsync(dstf);
                        await dstf.FlushAsync();
                    }
                }
            }
        }
        #endregion

        #region ベクターストア作成
        if (File.Exists(SrcFileName))
        {
            using (var fs = File.OpenRead(SrcFileName))
            {
                using (var sr = new StreamReader(fs))
                {
                    Console.WriteLine("start embeddings");
                    var locker = new SemaphoreSlim(5, 5);
                    var lines = sr.ReadToEnd();
                    var json = await Task.WhenAll(lines
                        .Split("\n")
                        .Where(l => !string.IsNullOrEmpty(l))
                        .Select(l => JsonSerializer.Deserialize<TrainData>(l))
                        .Select(async j =>
                        {
                            await locker.WaitAsync();
                            try
                            {
                                j.SentenceVector = await embedding.GenerateEmbeddingAsync(j.Sentence);
                            }
                            finally
                            {
                                locker.Release();
                            }
                            return j;
                        }));
                    Console.WriteLine("add collection");
                    await collection.UpsertAsync(json);
                    Console.WriteLine("end of initialize");
                }
            }
        }
        #endregion

        #region 検索のテスト
        string line = string.Empty;
        while (true)
        {
            Console.WriteLine();
            Console.Write("Vector検索 > ");
            line = Console.ReadLine() ?? string.Empty;
            if (string.IsNullOrEmpty(line))
            {
                Console.WriteLine(">>> finish");
                break;
            }

            var searchVector = await embedding.GenerateEmbeddingAsync($"検索クエリ: {line}");
            var searchResult = collection.SearchEmbeddingAsync(searchVector, top: 5);
            await foreach (var item in searchResult)
            {
                Console.WriteLine($"'{item.Record.Sentence}' (score: {item.Score})");
            }
        }
        #endregion
    }
}

public class TrainData
{
    [JsonPropertyName("sentence_pair_id")]
    [VectorStoreRecordKey]
    public string Id { get; set; } = string.Empty;

    [JsonPropertyName("sentence1")]
    [VectorStoreRecordData]
    public string Sentence { get; set; } = string.Empty;

    [VectorStoreRecordVector(Dimensions: 768, DistanceFunction = DistanceFunction.CosineSimilarity)]
    public ReadOnlyMemory<float> SentenceVector { get; set; }
}

元データは jsts を使わせてもらいます。
ソース内見てもらえばわかるのですが、データ量の多さからtrainを使うように最初は組んでたんですが、そのデータ量の多さに変換時間を待ちきれませんで…。
testにしてあります。

ベクターストア初期設定関連

データ形式の宣言

末尾に記載しているクラス TrainData がベクターストアのアクセス用のクラスです。
例によって例の如く、クラスがrow、プロパティがcolumn的なもの、カラムに対する設定はAttributeで、となっています。
Attributeの設定の仕方に関しては、こちらを確認してください。

最後のベクターデータの設定 VectorStoreRecordVector に関しては、 ruri-v3のHuggingFaceのページにある、 Model Details の欄を参考に設定してあります。
モデルにより記載データが違うのでデプロイしているモデルにより変更してください。
（サンプルコードは cl-nagoya/ruri-v3-310m 用）

C#宣言	モデル記載欄
Dimensions	Output Dimensionality
DistanceFunction	Similarity Function

使用するための初期設定

使用するベクターストアを登録・取得しているのはこちら。
インメモリを使用、コレクション名設定、コレクションが無い時の作成、を行っています。
コレクション取得 GetCollection の際に指定しているのは、キーの型とデータの型です。

.AddInMemoryVectorStore()

var vectorStore = kernel.GetRequiredService<IVectorStore>();
var collection = vectorStore.GetCollection<string, TrainData>(CollectionName);
await collection.CreateCollectionIfNotExistsAsync();

ローカルのruri-v3を使う初期設定関連

以前の記事通り、ローカルでOpenAIのEmbedding互換のAPIを使えるアプリ、Infinityを使用し、それに接続するように設定します。
登録の際にはモデル名、API接続キー、接続するURLのベースを登録しています。

.AddOpenAITextEmbeddingGeneration("cl-nagoya/ruri-v3-310m", "-", httpClient: new HttpClient()
{
    BaseAddress = new Uri("http://localhost:7997")
})

var embedding = kernel.GetRequiredService<ITextEmbeddingGenerationService>();

データの登録関連

テキストのベクターへの変換

ベクターストアに登録する前には、テキストをベクターデータに変換しておきます。

j.SentenceVector = await embedding.GenerateEmbeddingAsync(j.Sentence);

ベクターストアへデータを登録

データがそろったら、ベクターストアへデータを登録します。

await collection.UpsertAsync(json);

データの検索

ベクターストアから類似テキストを検索するのは、以下になります。
ruri-v3では、検索用のベクターを作成する際には、 検索クエリ: をつけます。

var searchVector = await embedding.GenerateEmbeddingAsync($"検索クエリ: {line}");
var searchResult = collection.SearchEmbeddingAsync(searchVector, top: 5);

その他のベクターストア

現在Semantic Kernelが対応しているベクターストアの一覧はこちらになります。

今からどのベクターストアがいいか、と考えるなら、MSのドキュメントのサンプルに書かれているベクターストア関連の量から考えて、Qdrantが一番ではないかな、という感じです。

まとめ

ベクターストアに関しては、なかなか奥が深いです。
Semantic Kernelでは、のちにに紹介するPluginと組み合わせてRAGを構成するように考えられていて、その際にもまた出てくるでしょう。

少し前に MSの太田さんがRAGの記事を書かれてました。
実は、このシリーズ、空き時間にサンプルプログラム組んで、少し寝かせてから公開してます。
寝かせてる間に記事がかぶってしまいました…。

Discussion