これは、なにをしたくて書いたもの？
前にApache Luceneでベクトル検索（kNN検索）を試してみました。
Apache Luceneでベクトル検索（kNN検索）を試す - CLOVER🍀
この後でANNに関係するHNSWのパッケージやクラスがあるのを見つけたので、今回はこちらを扱ってみたいなと思いまして。
kNN検索、ANNとHNSW
前回使ったのはkNN検索でした。これはQdrantの時にも扱いました。
Qdrantのチュートリアルから、「検索品質を測定する（Measure retrieval quality）」を試す - CLOVER🍀
あらためて内容を書いておきます。
k近傍法（ケイきんぼうほう、英: k-nearest neighbor algorithm, k-NN）は、入力との類似度が高い上位 k個の学習データで多数決/平均するアルゴリズムである。
k近傍法は以下の手順からなる：
1. 入力と全学習データとの類似度（距離）測定
2. 類似度上位 k 個の選出
3. 選出されたデータの多数決あるいは平均
すなわち「入力とよく似た k 個のデータで多数決/平均する」単純なアルゴリズムである
一方でANNは最近傍探索の派生で、近似最近傍探索と呼ばれるものです。
近似最近傍探索 (英: approximate nearest neighbor, ANN)
kNN検索に比べると、近似で検索することになります。
ANNを使うことによる利点は、検索対象のベクトル空間が大きくなった時に計算量が指数関数的に増加する（次元の呪い）のですが
この対策になっています。
つまりkNN検索に比べると、ANNは速度を犠牲にして高速化を狙っていることになります。
ではどうやって近似するのかというと、それにはアルゴリズムがあり、そのひとつがHNSW（Hierarchical navigable small world）です。
Hierarchical navigable small world - Wikipedia
これは多くのベクトルデータベースで採用されているグラフベースの近似最近傍検索手法です。
HNSWは階層グラフになっていて、次の2つのパラメーターを取ります。
- m（ノードごとのエッジの数＝各ノードが接続する隣接ノードの数） … この値を大きくするほど検索の精度は向上するが、より多くのメモリやディスクを消費する
- efConstruction（インデックス構築中に考慮する近傍の数） … この値を大きくするほど検索の精度は向上するが、インデックス作成速度は遅くなる
HNSWについては、Pineconeのドキュメントも参考になるでしょう。
Hierarchical Navigable Small Worlds (HNSW) | Pinecone
Apache LuceneのHNSW関連のパッケージ・クラス
Apache LuceneのHNSW関連のパッケージはこちらです。
org.apache.lucene.util.hnsw (Lucene 9.10.0 core API)
検索時には、このあたりを使うみたいですね。
- HnswGraphBuilder (Lucene 9.10.0 core API)
- HnswGraph (Lucene 9.10.0 core API)
- HnswGraphSearcher (Lucene 9.10.0 core API)
ただ、これらのクラスを直接使うことはないと思います。
その理由は以下のクラスなのですが、こちらは抽象クラスです。
KnnVectorsFormat (Lucene 9.10.0 core API)
具象クラスとしてはこれらがあります。
- Lucene99HnswVectorsFormat (Lucene 9.10.0 core API)
- Lucene99HnswScalarQuantizedVectorsFormat (Lucene 9.10.0 core API)
- PerFieldKnnVectorsFormat (Lucene 9.10.0 core API)
このうちPerFieldKnnVectorsFormatは、フィールドごとに設定するためのクラスのようです。
これらのクラスでmやefConstructionを設定するのですが、その方法を知るにはApache LuceneのCodecについて知っておく必要があります。
Apache LuceneのCodec
Apache LuceneのCodecクラスは、転置インデックスセグメントをエンコード／デコードするためのクラスです。
Codec (Lucene 9.10.0 core API)
このクラスから、各種データ形式へのFormatを取得できます。
各Formatクラスからはインデックスファイル（セグメント）を読み書きするためのクラスが提供されるので、これを介してインデックス
ファイル（セグメント）にアクセスします。なのでEncoder／Decoderなわけですね。
Codecについては、elastic社のブログが参考になるでしょう。
What is an Apache Lucene Codec? | Elastic Blog
なお、Codecクラスは抽象クラスであり、具象クラスとしては現時点のApache LuceneではLucene99Codecクラスが使われます。
Lucene99Codec (Lucene 9.10.0 core API)
Apache LuceneのCodecとベクトル検索
Codecクラスから各種データ形式へのFormatが取得できると書きましたが、ベクトルで使われるKnnVectorsFormatもその一種です。
KnnVectorsFormat (Lucene 9.10.0 core API)
KnnVectorsFormatからは、次の2つのクラスのインスタンスが得られます。
KnnVectorsReaderにはsearchメソッドが定義されていて、この中で呼び出し方に応じてkNN検索またはANNが実行されることに
なります。
KnnVectorsReaderクラスは抽象クラスで、現時点の実装クラスはLucene99HnswVectorsReaderですね。
Lucene99HnswVectorsReader (Lucene 9.10.0 core API)
具体的には、KnnByteVectorQueryまたはKnnFloatVectorQueryで指定するkの値（検索するドキュメント数）がインデックス内の
全ベクトル数（≒全ドキュメントの数）より小さければHNSWを使った（HnswGraphSearcherによる）ANNを、そうでなければkNN検索を
実行するようです。
if (knnCollector.k() < scorer.maxOrd()) { HnswGraphSearcher.search(scorer, collector, getGraph(fieldEntry), acceptedOrds); } else { // if k is larger than the number of vectors, we can just iterate over all vectors // and collect them for (int i = 0; i < scorer.maxOrd(); i++) { if (acceptedOrds == null || acceptedOrds.get(i)) { knnCollector.incVisitedCount(1); knnCollector.collect(scorer.ordToDoc(i), scorer.score(i)); } } }
HnswGraphSearcherなどを直接使うことがないだろうと書いたのは、このあたりが理由です。
これらのクエリーの説明には、以下のようなことが書かれています。
- If the filter cost is less than k, just execute an exact search
- Otherwise run a kNN search subject to the filter
- If the kNN search visits too many vectors without completing, stop and run an exact search
この中のフィルターのコストによっては動作が変わるという記述が、kNN検索が行われるのかANNが行われるのかの切り替えについて
触れている内容なのかなと思います。
なお、以下のクラスとこの文脈での関係はわかっていません…。
- ByteVectorSimilarityQuery (Lucene 9.10.0 core API)
- FloatVectorSimilarityQuery (Lucene 9.10.0 core API)
ところで、HNSWにはmとefConstructionという2つのパラメーターがあるという話を書きました。これらを設定するのは以下のクラスに
なります。
- Lucene99HnswVectorsFormat (Lucene 9.10.0 core API)
- Lucene99HnswScalarQuantizedVectorsFormat (Lucene 9.10.0 core API)
そしてデフォルト値はmが16、
efConstruction（Apache Lucene上ではbeamWidthという名前になっています）が100になっています。
ちなみにmの最大値は512、efConstructionの最大値は3200です。
Hibernate Searchのドキュメントでは、mは2〜100の範囲で指定することが推奨されていました。efConstructionについては特に推奨値は
ありませんでしたね。
Hibernate Search 7.1.1.Final: Reference Documentation
特にCodecの設定を行わない場合は、Lucene99Codecが作成したKnnVectorsFormatのインスタンスが使われます。
これはPerFieldKnnVectorsFormatとLucene99HnswVectorsFormatの組み合わせになっています。
※Lucene99Codec#getKnnVectorsFormatForFieldはdefaultKnnVectorsFormatを呼び出しているだけです。この時fieldの値は無視します
private final KnnVectorsFormat defaultKnnVectorsFormat; private final KnnVectorsFormat knnVectorsFormat = new PerFieldKnnVectorsFormat() { @Override public KnnVectorsFormat getKnnVectorsFormatForField(String field) { return Lucene99Codec.this.getKnnVectorsFormatForField(field); } };
this.defaultKnnVectorsFormat = new Lucene99HnswVectorsFormat();
Lucene99HnswVectorsFormatをデフォルトコンストラクタで作成した場合、mもefConstructionもデフォルト値になるため、
これをカスタマイズする場合はApache LuceneのCodecの設定を行うことになります。
CodecはIndexWriterConfigで設定することになります。
IndexWriterConfig (Lucene 9.10.0 core API)
そして、設定したm（maxConn）やefConstruction（beamWidth）は書き込み時に使われます。
@Override public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { return new Lucene99HnswVectorsWriter( state, maxConn, beamWidth, flatVectorsFormat.fieldsWriter(state), numMergeWorkers, mergeExec); }
このあたりが今回Apache LuceneのCodecの話を持ち出した事情ですね。
では、今回はCodecのカスタマイズをしたり、kNN検索、ANNをしたりしてみましょう。
準備
今回の環境はこちら。
$ java --version openjdk 21.0.3 2024-04-16 OpenJDK Runtime Environment (build 21.0.3+9-Ubuntu-1ubuntu122.04.1) OpenJDK 64-Bit Server VM (build 21.0.3+9-Ubuntu-1ubuntu122.04.1, mixed mode, sharing) $ mvn --version Apache Maven 3.9.7 (8b094c9513efc1b9ce2d952b3b9c8eaedaf8cbf0) Maven home: $HOME/.sdkman/candidates/maven/current Java version: 21.0.3, vendor: Ubuntu, runtime: /usr/lib/jvm/java-21-openjdk-amd64 Default locale: ja_JP, platform encoding: UTF-8 OS name: "linux", version: "5.15.0-107-generic", arch: "amd64", family: "unix"
また、Pythonも使います。
$ python3 --version Python 3.10.12 $ pip3 --version pip 22.0.2 from /usr/lib/python3/dist-packages/pip (python 3.10)
テキスト埋め込み
テキスト埋め込みを行う機能はApache Luceneにはないので、Sentence Transformersを使ったAPIを作ってこちらを利用することにします。
$ pip3 install sentence-transformers fastapi uvicorn[standard]
インストールされたライブラリー。
$ pip3 list Package Version ------------------------ ---------- annotated-types 0.7.0 anyio 4.4.0 certifi 2024.6.2 charset-normalizer 3.3.2 click 8.1.7 dnspython 2.6.1 email_validator 2.1.1 exceptiongroup 1.2.1 fastapi 0.111.0 fastapi-cli 0.0.4 filelock 3.14.0 fsspec 2024.6.0 h11 0.14.0 httpcore 1.0.5 httptools 0.6.1 httpx 0.27.0 huggingface-hub 0.23.2 idna 3.7 Jinja2 3.1.4 joblib 1.4.2 markdown-it-py 3.0.0 MarkupSafe 2.1.5 mdurl 0.1.2 mpmath 1.3.0 networkx 3.3 numpy 1.26.4 nvidia-cublas-cu12 12.1.3.1 nvidia-cuda-cupti-cu12 12.1.105 nvidia-cuda-nvrtc-cu12 12.1.105 nvidia-cuda-runtime-cu12 12.1.105 nvidia-cudnn-cu12 8.9.2.26 nvidia-cufft-cu12 11.0.2.54 nvidia-curand-cu12 10.3.2.106 nvidia-cusolver-cu12 11.4.5.107 nvidia-cusparse-cu12 12.1.0.106 nvidia-nccl-cu12 2.20.5 nvidia-nvjitlink-cu12 12.5.40 nvidia-nvtx-cu12 12.1.105 orjson 3.10.3 packaging 24.0 pillow 10.3.0 pip 22.0.2 pydantic 2.7.3 pydantic_core 2.18.4 Pygments 2.18.0 python-dotenv 1.0.1 python-multipart 0.0.9 PyYAML 6.0.1 regex 2024.5.15 requests 2.32.3 rich 13.7.1 safetensors 0.4.3 scikit-learn 1.5.0 scipy 1.13.1 sentence-transformers 3.0.0 setuptools 59.6.0 shellingham 1.5.4 sniffio 1.3.1 starlette 0.37.2 sympy 1.12.1 threadpoolctl 3.5.0 tokenizers 0.19.1 torch 2.3.0 tqdm 4.66.4 transformers 4.41.2 triton 2.3.0 typer 0.12.3 typing_extensions 4.12.1 ujson 5.10.0 urllib3 2.2.1 uvicorn 0.30.1 uvloop 0.19.0 watchfiles 0.22.0 websockets 12.0
api.py
from fastapi import FastAPI from pydantic import BaseModel import os from sentence_transformers import SentenceTransformer app = FastAPI() class EmbeddingRequest(BaseModel): model: str text: str normalize: bool = False class EmbeddingResponse(BaseModel): model: str embedding: list[float] dimension: int @app.post("/embeddings/encode") def encode(request: EmbeddingRequest) -> EmbeddingResponse: sentence_transformer_model = SentenceTransformer( request.model, device=os.getenv("EMBEDDING_API_DEVICE", "cpu") ) embeddings = sentence_transformer_model.encode(sentences=[request.text], normalize_embeddings=request.normalize) embedding = embeddings[0] # numpy array to float list embedding_as_float = embedding.tolist() return EmbeddingResponse( model=request.model, embedding=embedding_as_float, dimension=sentence_transformer_model.get_sentence_embedding_dimension() )
リクエストにはモデル、変換するテキストを指定します。
起動。
$ uvicorn api:app
こちらのAPIを使ってテキスト埋め込みを行います。
準備
Maven依存関係等はこちら。
<properties> <maven.compiler.release>21</maven.compiler.release> <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding> <project.reporting.outputEncoding>UTF-8</project.reporting.outputEncoding> </properties> <dependencies> <dependency> <groupId>org.apache.lucene</groupId> <artifactId>lucene-core</artifactId> <version>9.10.0</version> </dependency> <dependency> <groupId>org.junit.jupiter</groupId> <artifactId>junit-jupiter</artifactId> <version>5.10.2</version> <scope>test</scope> </dependency> <dependency> <groupId>org.assertj</groupId> <artifactId>assertj-core</artifactId> <version>3.26.0</version> <scope>test</scope> </dependency> <dependency> <groupId>com.fasterxml.jackson.core</groupId> <artifactId>jackson-databind</artifactId> <version>2.17.1</version> <scope>test</scope> </dependency> <dependency> <groupId>org.jboss.byteman</groupId> <artifactId>byteman-bmunit5</artifactId> <scope>test</scope> <version>4.0.23</version> </dependency> </dependencies> <build> <plugins> <plugin> <artifactId>maven-surefire-plugin</artifactId> <version>3.2.5</version> <configuration> <argLine>-Djdk.attach.allowAttachSelf -XX:+EnableDynamicAgentLoading</argLine> </configuration> </plugin> </plugins> </build>
Apache Luceneのベクトル検索を行うにはlucene-coreがあれば十分で、HNSWに関するパッケージもlucene-coreに含まれています。
確認はテストコードで行い、Jacksonは先ほど作成したAPIを呼び出す際に使います。
なぜかBytemanが入っていますが、これは検索に使ったのがkNN検索なのかANNなのか、
それからm（maxConn）やefConstruction（beamWidth）をカスタマイズした時に反映できているのか、といった確認に使います。
APIを呼び出すためのクライアントはこちら。
src/test/java/org/littlewings/lucene/ann/EmbeddingClient.java
package org.littlewings.lucene.ann; import com.fasterxml.jackson.databind.ObjectMapper; import java.io.IOException; import java.io.UncheckedIOException; import java.net.URI; import java.net.http.HttpClient; import java.net.http.HttpRequest; import java.net.http.HttpResponse; import java.nio.charset.StandardCharsets; import java.util.List; public class EmbeddingClient implements AutoCloseable { private String url; private HttpClient httpClient; private ObjectMapper objectMapper; EmbeddingClient(String host, int port, HttpClient httpClient) { this.url = String.format("http://%s:%d/embeddings/encode", host, port); this.httpClient = httpClient; this.objectMapper = new ObjectMapper(); } public static EmbeddingClient create(String host, int port) { return new EmbeddingClient( host, port, HttpClient .newBuilder() .version(HttpClient.Version.HTTP_1_1) .followRedirects(HttpClient.Redirect.ALWAYS) .build() ); } public EmbeddingResponse execute(EmbeddingRequest request) { try { String json = objectMapper.writeValueAsString(request); HttpRequest httpRequest = HttpRequest.newBuilder() .uri(URI.create(url)) .header("Content-Type", "application/json") .POST(HttpRequest.BodyPublishers.ofString(json)) .build(); HttpResponse<String> httpResponse = httpClient .send(httpRequest, HttpResponse.BodyHandlers.ofString(StandardCharsets.UTF_8)); return objectMapper.readValue(httpResponse.body(), EmbeddingResponse.class); } catch (IOException e) { throw new UncheckedIOException(e); } catch (InterruptedException e) { throw new RuntimeException(e); } } @Override public void close() { httpClient.close(); } public record EmbeddingRequest(String model, String text) { } public record EmbeddingResponse(String model, List<Float> embedding, int dimension) { } }
テストコードの雛形。
src/test/java/org/littlewings/lucene/ann/AnnSearchTest.java
package org.littlewings.lucene.ann; import org.apache.lucene.codecs.Codec; import org.apache.lucene.codecs.KnnVectorsFormat; import org.apache.lucene.codecs.lucene99.Lucene99Codec; import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat; import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat; import org.apache.lucene.document.*; import org.apache.lucene.index.*; import org.apache.lucene.search.IndexSearcher; import org.apache.lucene.search.KnnFloatVectorQuery; import org.apache.lucene.search.TopDocs; import org.apache.lucene.store.ByteBuffersDirectory; import org.apache.lucene.store.Directory; import org.jboss.byteman.contrib.bmunit.BMScript; import org.jboss.byteman.contrib.bmunit.BMScripts; import org.jboss.byteman.contrib.bmunit.BMUnitConfig; import org.jboss.byteman.contrib.bmunit.WithByteman; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import java.io.IOException; import java.io.UncheckedIOException; import java.util.Arrays; import java.util.List; import static org.assertj.core.api.Assertions.assertThat; @WithByteman @BMUnitConfig(loadDirectory = "target/test-classes") @BMScripts(scripts = { @BMScript("trace-ann-searcher.btm"), @BMScript("trace-hnsw-parameters.btm"), } ) class AnnSearchTest { EmbeddingClient embeddingClient; @BeforeEach void setUp() { embeddingClient = EmbeddingClient.create("localhost", 8000); } @AfterEach void tearDown() { embeddingClient.close(); } float[] textToVector(String text) { EmbeddingClient.EmbeddingRequest request = new EmbeddingClient.EmbeddingRequest( "intfloat/multilingual-e5-base", text ); EmbeddingClient.EmbeddingResponse response = embeddingClient.execute(request); float[] vector = new float[response.embedding().size()]; for (int i = 0; i < response.embedding().size(); i++) { vector[i] = response.embedding().get(i); } return vector; } // ここに、データ作成とテストを書く }
テキスト埋め込みに使うモデルは、intfloat/multilingual-e5-baseとします。
intfloat/multilingual-e5-base · Hugging Face
データの作成
データの作成は、以下のメソッドを使って行います。
List<Document> createDocuments() {
return List.of(
createDocument(
"The Time Machine",
"A man travels through time and witnesses the evolution of humanity.",
"H.G. Wells",
1895
),
createDocument(
"Ender's Game",
"A young boy is trained to become a military leader in a war against an alien race.",
"Orson Scott Card",
1985
),
createDocument(
"Brave New World",
"A dystopian society where people are genetically engineered and conditioned to conform to a strict social hierarchy.",
"Aldous Huxley",
1932
),
createDocument(
"The Hitchhiker's Guide to the Galaxy",
"A comedic science fiction series following the misadventures of an unwitting human and his alien friend.",
"Douglas Adams",
1979
),
createDocument(
"Dune",
"A desert planet is the site of political intrigue and power struggles.",
"Frank Herbert",
1965
),
createDocument(
"Foundation",
"A mathematician develops a science to predict the future of humanity and works to save civilization from collapse.",
"Isaac Asimov",
1951
),
createDocument(
"Snow Crash",
"A futuristic world where the internet has evolved into a virtual reality metaverse.",
"Neal Stephenson",
1992
),
createDocument(
"Neuromancer",
"A hacker is hired to pull off a near-impossible hack and gets pulled into a web of intrigue.",
"William Gibson",
1984
),
createDocument(
"The War of the Worlds",
"A Martian invasion of Earth throws humanity into chaos.",
"H.G. Wells",
1898
),
createDocument(
"The Hunger Games",
"A dystopian society where teenagers are forced to fight to the death in a televised spectacle.",
"Suzanne Collins",
2008
),
createDocument(
"The Andromeda Strain",
"A deadly virus from outer space threatens to wipe out humanity.",
"Michael Crichton",
1969
),
createDocument(
"The Left Hand of Darkness",
"A human ambassador is sent to a planet where the inhabitants are genderless and can change gender at will.",
"Ursula K. Le Guin",
1969
),
createDocument(
"The Three-Body Problem",
"Humans encounter an alien civilization that lives in a dying system.",
"Liu Cixin",
2008
)
);
}
Document createDocument(String name, String description, String author, int year) {
Document document = new Document();
document.add(new TextField("name", name, Field.Store.YES));
document.add(new TextField("description", description, Field.Store.YES));
document.add(new TextField("author", author, Field.Store.YES));
document.add(new IntField("year", year, Field.Store.YES));
float[] vector = textToVector("passage: " + description);
document.add(new KnnFloatVectorField("description_vector", vector, VectorSimilarityFunction.EUCLIDEAN));
return document;
}
float[] textToVector(String text) {
EmbeddingClient.EmbeddingRequest request =
new EmbeddingClient.EmbeddingRequest(
"intfloat/multilingual-e5-base",
text
);
EmbeddingClient.EmbeddingResponse response = embeddingClient.execute(request);
float[] vector = new float[response.embedding().size()];
for (int i = 0; i < response.embedding().size(); i++) {
vector[i] = response.embedding().get(i);
}
return vector;
}
これはこちらのエントリーで使ったものと同じ、13個のデータで
Apache Luceneでベクトル検索（kNN検索）を試す - CLOVER🍀
Qdrantのチュートリアルのものでもあります。
テキスト埋め込みを行い、ベクトルを保存するフィールドはこちらです。ベクトル類似度関数はユークリッド距離とします。
float[] vector = textToVector("passage: " + description); document.add(new KnnFloatVectorField("description_vector", vector, VectorSimilarityFunction.EUCLIDEAN));
passage:はintfloat/multilingual-e5でドキュメントに指定する接頭辞です。
トレースの考え方
検索時にkNN検索とANNのどちらになっているのかを確認しようと、Explanationを使ってみたのですが。
Explanation (Lucene 9.10.0 core API)
得られた結果がこんな感じで、ほとんど情報がありませんでした。
0.75919974 = within top 13
これでは困るのでどう確認するか迷ったところ、Bytemanを使うことにしました。
BMUnit : Using Byteman with JUnit or TestNG fro...| JBoss.org Content Archive (Read Only)
スクリプトの指定はこちら。
@WithByteman @BMUnitConfig(loadDirectory = "target/test-classes") @BMScripts(scripts = { @BMScript("trace-ann-searcher.btm"), @BMScript("trace-hnsw-parameters.btm"), } ) class AnnSearchTest {
作成したBytemanのスクリプトはこちらです。
kNN検索かANNのどちらを行うかの条件であるKnnCollector#kとRandomVectorScorer#maxOrdの記録と、ANNを実行する際に使われる
HnswGraphSearcher#searchのトレース。HnswGraphSearcher#searchが呼び出されていない場合は、kNN検索が実行されていることに
なります。
src/test/resources/trace-ann-searcher.btm
RULE print KnnCollector k value
INTERFACE org.apache.lucene.search.KnnCollector
METHOD k
AT EXIT
IF TRUE
DO traceln("KnnCollector#k = " + $!)
ENDRULE
RULE print RandomVectorScorer maxOrd value
INTERFACE org.apache.lucene.util.hnsw.RandomVectorScorer
METHOD maxOrd
AT EXIT
IF TRUE
DO traceln("RandomVectorScorer#maxOrd = " + $!)
ENDRULE
RULE print HnswGraphSearcher called
CLASS org.apache.lucene.util.hnsw.HnswGraphSearcher
METHOD search(org.apache.lucene.util.hnsw.RandomVectorScorer, org.apache.lucene.search.KnnCollector, org.apache.lucene.util.hnsw.HnswGraph, org.apache.lucene.util.hnsw.HnswGraphSearcher, org.apache.lucene.util.Bits)
AT ENTRY
IF TRUE
DO traceln("execute ANN(HNSW) HnswGraphSearcher#search")
ENDRULE
また、インデックスの書き込み時に使われるm（maxConn）、efConstruction（beamWidth）の値も記録するようにしましょう。
src/test/resources/trace-hnsw-parameters.btm
RULE print HNSW parameters
CLASS org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsWriter
METHOD <init>
AT ENTRY
IF TRUE
DO traceln("m = " + $2 + ", beamWith = " + $3)
ENDRULE
こちらを使って確認していきたいと思います。
kNN検索を行う
まず最初はkNN検索を行ってみましょう。
@Test void knnSearch() throws IOException { try (Directory directory = new ByteBuffersDirectory()) { IndexWriterConfig config = new IndexWriterConfig(); List<Document> documents = createDocuments(); try (IndexWriter writer = new IndexWriter(directory, config)) { writer.addDocuments(documents); } try (DirectoryReader reader = DirectoryReader.open(directory)) { IndexSearcher searcher = new IndexSearcher(reader); int k = documents.size(); // all document float[] vector = textToVector("query: alien invasion"); KnnFloatVectorQuery query = new KnnFloatVectorQuery("description_vector", vector, k); TopDocs topDocs = searcher.search(query, 3); StoredFields storedFields = searcher.storedFields(); List<Document> resultDocuments = Arrays.stream(topDocs.scoreDocs) .map(scoreDoc -> { try { return storedFields.document(scoreDoc.doc); } catch (IOException e) { throw new UncheckedIOException(e); } }) .toList(); assertThat(resultDocuments).hasSize(3); assertThat(resultDocuments.get(0).getField("name").stringValue()) .isEqualTo("The Hitchhiker's Guide to the Galaxy"); assertThat(resultDocuments.get(0).getField("year").numericValue().intValue()) .isEqualTo(1979); assertThat(resultDocuments.get(1).getField("name").stringValue()) .isEqualTo("The Three-Body Problem"); assertThat(resultDocuments.get(1).getField("year").numericValue().intValue()) .isEqualTo(2008); assertThat(resultDocuments.get(2).getField("name").stringValue()) .isEqualTo("The Andromeda Strain"); assertThat(resultDocuments.get(2).getField("year").numericValue().intValue()) .isEqualTo(1969); } } }
KnnFloatVectorQueryに指定するkはドキュメントの数と同じにします。kはベクトル検索で検索対象とするドキュメントの数ですね。
厳密なことを言うと、比較対象はインデックス内の全ドキュメントの数というよりはその中のベクトルの数なのでえすが、今回は
すべてのベクトル用のフィールドに値を設定しているので全ベクトル数＝全ドキュメントの数になっています。
int k = documents.size(); // all document float[] vector = textToVector("query: alien invasion"); KnnFloatVectorQuery query = new KnnFloatVectorQuery("description_vector", vector, k);
query:というのは、intfloat/multilingual-e5で検索時に使う接頭辞です。
Bytemanを使って差し込んだスクリプトで、実行時に得られたログはこちら。
KnnCollector#kとRandomVectorScorer#maxOrdの値。
※たくさん出力されます
KnnCollector#k = 13 RandomVectorScorer#maxOrd = 13
そして、HnswGraphSearcher#searchを呼び出したログはありませんでした。
m（maxConn）、efConstruction（beamWidth）の値はこちら。デフォルト値ですね。
m = 16, beamWith = 100
ANNを行う
次はANNを行ってみます。
@Test void annSearch() throws IOException { try (Directory directory = new ByteBuffersDirectory()) { IndexWriterConfig config = new IndexWriterConfig(); List<Document> documents = createDocuments(); try (IndexWriter writer = new IndexWriter(directory, config)) { writer.addDocuments(documents); } try (DirectoryReader reader = DirectoryReader.open(directory)) { IndexSearcher searcher = new IndexSearcher(reader); int k = documents.size() - 1; // all documents - 1 float[] vector = textToVector("query: alien invasion"); KnnFloatVectorQuery query = new KnnFloatVectorQuery("description_vector", vector, k); TopDocs topDocs = searcher.search(query, 3); StoredFields storedFields = searcher.storedFields(); List<Document> resultDocuments = Arrays.stream(topDocs.scoreDocs) .map(scoreDoc -> { try { return storedFields.document(scoreDoc.doc); } catch (IOException e) { throw new UncheckedIOException(e); } }) .toList(); assertThat(resultDocuments).hasSize(3); assertThat(resultDocuments.get(0).getField("name").stringValue()) .isEqualTo("The Hitchhiker's Guide to the Galaxy"); assertThat(resultDocuments.get(0).getField("year").numericValue().intValue()) .isEqualTo(1979); assertThat(resultDocuments.get(1).getField("name").stringValue()) .isEqualTo("The Three-Body Problem"); assertThat(resultDocuments.get(1).getField("year").numericValue().intValue()) .isEqualTo(2008); assertThat(resultDocuments.get(2).getField("name").stringValue()) .isEqualTo("The Andromeda Strain"); assertThat(resultDocuments.get(2).getField("year").numericValue().intValue()) .isEqualTo(1969); } } }
kNN検索との違いは、この1点です。
int k = documents.size() - 1; // all documents - 1 float[] vector = textToVector("query: alien invasion"); KnnFloatVectorQuery query = new KnnFloatVectorQuery("description_vector", vector, k);
KnnFloatVectorQueryに指定するkはドキュメントの数から1引いています。ここでもすべてのドキュメントでベクトルに値を保存して
いるので、インデックス内の全ベクトル数は全ドキュメントの数と同じになります。
Bytemanを使って差し込んだスクリプトで、実行時に得られたログはこちら。
KnnCollector#kとRandomVectorScorer#maxOrdの値。
KnnCollector#k = 12 RandomVectorScorer#maxOrd = 13
そして、HnswGraphSearcher#searchを呼び出したことを表すログが記録されました。
execute ANN(HNSW) HnswGraphSearcher#search
m（maxConn）、efConstruction（beamWidth）の値はデフォルト値で同じ値ですね。
m = 16, beamWith = 100
というわけで、検索対象とするドキュメント数を、インデックスに保存している全ベクトル数（≒全ドキュメントの数）よりも
小さい値を指定すると近似となりANNに切り替わることがわかりました。
デフォルトのCodecを確認する
次は、HNSWのパラメーターであるmとefConstructionの確認に移ります。
まずはデフォルトのCodecを確認してみましょう。
@Test void defaultCodec() throws ReflectiveOperationException { Codec codec = Codec.getDefault(); assertThat(codec).isExactlyInstanceOf(Lucene99Codec.class); assertThat(codec.knnVectorsFormat()).isInstanceOf(PerFieldKnnVectorsFormat.class); Lucene99Codec lucene99Codec = (Lucene99Codec) codec; java.lang.reflect.Field defaultKnnVectorsFormatField = Lucene99Codec.class.getDeclaredField("defaultKnnVectorsFormat"); defaultKnnVectorsFormatField.setAccessible(true); KnnVectorsFormat knnVectorsFormat = (KnnVectorsFormat) defaultKnnVectorsFormatField.get(lucene99Codec); assertThat(knnVectorsFormat).isExactlyInstanceOf(Lucene99HnswVectorsFormat.class); }
このCodec#getDefaultがIndexWriterConfigで使われるCodecのデフォルトのインスタンスとなっています。
Apache Lucene 9.10.0ではLucene99Codecになります。
Lucene99Codecが使用するKnnVectorsFormatの実装は、Lucene99HnswVectorsFormatになっていることも確認できました。
独自のCodecを作成する
それでは、HNSWのパラメーターであるm（maxConn）、efConstruction（beamWidth）をカスタマイズするために独自の
Codecを作成します。
最初はKnnVectorsFormatのインスタンスだけ作ればいいのかと思っていたのですが、Codecは変更できないのでCodec自体を作成する
必要があります。
となると、Codecが要求するメソッドをすべて実装する必要があるのかなと思ったりするのですが
Codec (Lucene 9.10.0 core API)
そんなことはなく、FilterCodecを使うと既存のCodecの具象クラスを使って差分を実装することができます。
FilterCodec (Lucene 9.10.0 core API)
今回はLucene99Codecをベースにします。
Lucene99Codec (Lucene 9.10.0 core API)
そんなわけで、作成した独自のCodecはこちら。
src/test/java/org/littlewings/lucene/ann/MyCustomCodec.java
package org.littlewings.lucene.ann; import org.apache.lucene.codecs.FilterCodec; import org.apache.lucene.codecs.KnnVectorsFormat; import org.apache.lucene.codecs.lucene99.Lucene99Codec; import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat; import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat; public class MyCustomCodec extends FilterCodec { private KnnVectorsFormat knnVectorsFormat = new MyKnnVectorsFormat(); public MyCustomCodec() { super("MyCustom", new Lucene99Codec()); } @Override public KnnVectorsFormat knnVectorsFormat() { return knnVectorsFormat; } static class MyKnnVectorsFormat extends PerFieldKnnVectorsFormat { private KnnVectorsFormat defaultKnnVectorsFormat = new Lucene99HnswVectorsFormat(); private KnnVectorsFormat forDescriptionVectorField = new Lucene99HnswVectorsFormat(32, 150); @Override public KnnVectorsFormat getKnnVectorsFormatForField(String field) { System.out.printf("KnnVectorsFormat#getKnnVectorsFormatForField = %s%n", field); if ("description_vector".equals(field)) { return forDescriptionVectorField; } else { return defaultKnnVectorsFormat; } } } }
親コンストラクタには、Codecの名前と委譲先になるCodecの実装のインスタンスを渡します。オーバーライドしたメソッド以外は、
このインスタンスに処理が転送されることになります。
public MyCustomCodec() { super("MyCustom", new Lucene99Codec()); }
今回はdescription_vectorフィールドのみカスタマイズしたLucene99HnswVectorsFormatのインスタンスを割り当て、こちらの
m（maxConn）を32に、efConstruction（beamWidth）を150にします。
private KnnVectorsFormat knnVectorsFormat = new MyKnnVectorsFormat(); 〜省略〜 @Override public KnnVectorsFormat knnVectorsFormat() { return knnVectorsFormat; } static class MyKnnVectorsFormat extends PerFieldKnnVectorsFormat { private KnnVectorsFormat defaultKnnVectorsFormat = new Lucene99HnswVectorsFormat(); private KnnVectorsFormat forDescriptionVectorField = new Lucene99HnswVectorsFormat(32, 150); @Override public KnnVectorsFormat getKnnVectorsFormatForField(String field) { System.out.printf("KnnVectorsFormat#getKnnVectorsFormatForField = %s%n", field); if ("description_vector".equals(field)) { return forDescriptionVectorField; } else { return defaultKnnVectorsFormat; } } }
呼び出された時にはフィールド名を出力するようにしています。
description_vector以外のフィールドにはデフォルトのLucene99HnswVectorsFormatのインスタンスを使います。
このCodecはServiceLoaderの仕組みで呼び出すので、以下のファイルを作成して実装クラスの名前を書いておきます。
src/test/resources/META-INF/services/org.apache.lucene.codecs.Codec
org.littlewings.lucene.ann.MyCustomCodec
こうすることで、Codec#forNameでこのインスタンスを取得できます。Codec#forNameで指定する値は、Codecを作成する時に
親コンストラクタに指定した名前です。
@Test void customizeHsnwParamersCodec1() { Codec codec = Codec.forName("MyCustom"); // load from ServiceLoader assertThat(codec).isExactlyInstanceOf(MyCustomCodec.class); }
もっとも、ServiceLoaderの仕組みに従わずにふつうにnewしてもいいのですが、お作法に則ると前述のやり方ということになるでしょう。
@Test void customizeHsnwParamersCodec2() { Codec codec = new MyCustomCodec(); }
では、このCodecを使ってインデックスにドキュメントを登録、検索してみます。
その前に、利用するドキュメントを少し変えたいと思います。
Document createDocument2(String name, String description, String author, int year) { Document document = new Document(); document.add(new TextField("name", name, Field.Store.YES)); document.add(new TextField("description", description, Field.Store.YES)); document.add(new TextField("author", author, Field.Store.YES)); document.add(new IntField("year", year, Field.Store.YES)); float[] vector = textToVector("passage: " + description); document.add(new KnnFloatVectorField("description_vector", vector, VectorSimilarityFunction.EUCLIDEAN)); document.add(new KnnFloatVectorField("description_vector2", vector, VectorSimilarityFunction.EUCLIDEAN)); return document; }
ベクトル用のフィールドを2つにしました。片方はデフォルトのm（maxConn）とefConstruction（beamWidth）を使います。
確認。
@Test void customizeHnswParameters() throws IOException { try (Directory directory = new ByteBuffersDirectory()) { IndexWriterConfig config = new IndexWriterConfig(); Codec codec = Codec.forName("MyCustom"); config.setCodec(codec); List<Document> documents = List.of( createDocument2( "The Time Machine", "A man travels through time and witnesses the evolution of humanity.", "H.G. Wells", 1895 ), createDocument2( "Ender's Game", "A young boy is trained to become a military leader in a war against an alien race.", "Orson Scott Card", 1985 ), createDocument2( "Brave New World", "A dystopian society where people are genetically engineered and conditioned to conform to a strict social hierarchy.", "Aldous Huxley", 1932 ) ); try (IndexWriter writer = new IndexWriter(directory, config)) { writer.addDocuments(documents); } try (DirectoryReader reader = DirectoryReader.open(directory)) { IndexSearcher searcher = new IndexSearcher(reader); int k = 1; // ANN float[] vector = textToVector("query: alien invasion"); KnnFloatVectorQuery query = new KnnFloatVectorQuery("description_vector", vector, k); TopDocs topDocs = searcher.search(query, 1); StoredFields storedFields = searcher.storedFields(); List<Document> resultDocuments = Arrays.stream(topDocs.scoreDocs) .map(scoreDoc -> { try { return storedFields.document(scoreDoc.doc); } catch (IOException e) { throw new UncheckedIOException(e); } }) .toList(); assertThat(resultDocuments).hasSize(1); assertThat(resultDocuments.get(0).getField("name").stringValue()) .isEqualTo("Ender's Game"); assertThat(resultDocuments.get(0).getField("year").numericValue().intValue()) .isEqualTo(1985); } } }
Codecを設定しているのはこちら。
IndexWriterConfig config = new IndexWriterConfig(); Codec codec = Codec.forName("MyCustom"); config.setCodec(codec);
検索はANNにしました。
int k = 1; // ANN float[] vector = textToVector("query: alien invasion"); KnnFloatVectorQuery query = new KnnFloatVectorQuery("description_vector", vector, k);
ログを確認しましょう。
KnnVectorsFormat#getKnnVectorsFormatForField = description_vector KnnVectorsFormat#getKnnVectorsFormatForField = description_vector m = 32, beamWith = 150 KnnVectorsFormat#getKnnVectorsFormatForField = description_vector2 KnnVectorsFormat#getKnnVectorsFormatForField = description_vector2 m = 16, beamWith = 100
description_vectorの方は、カスタマイズしたm（maxConn）とefConstruction（beamWidth）の値が使われていることが確認できます。
全体のドキュメント数が3なのに対して、検索対象としたドキュメント数は1なのでANNになりました。
KnnCollector#k = 1 RandomVectorScorer#maxOrd = 3 KnnCollector#k = 1 KnnCollector#k = 1 execute ANN(HNSW) HnswGraphSearcher#search
これで、確認したいことはできたかなと。
オマケ
最後に、今回作成したテストコード全体を載せておきます。
src/test/java/org/littlewings/lucene/ann/AnnSearchTest.java
package org.littlewings.lucene.ann; import org.apache.lucene.codecs.Codec; import org.apache.lucene.codecs.KnnVectorsFormat; import org.apache.lucene.codecs.lucene99.Lucene99Codec; import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat; import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat; import org.apache.lucene.document.*; import org.apache.lucene.index.*; import org.apache.lucene.search.IndexSearcher; import org.apache.lucene.search.KnnFloatVectorQuery; import org.apache.lucene.search.TopDocs; import org.apache.lucene.store.ByteBuffersDirectory; import org.apache.lucene.store.Directory; import org.jboss.byteman.contrib.bmunit.BMScript; import org.jboss.byteman.contrib.bmunit.BMScripts; import org.jboss.byteman.contrib.bmunit.BMUnitConfig; import org.jboss.byteman.contrib.bmunit.WithByteman; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import java.io.IOException; import java.io.UncheckedIOException; import java.util.Arrays; import java.util.List; import static org.assertj.core.api.Assertions.assertThat; @WithByteman @BMUnitConfig(loadDirectory = "target/test-classes") @BMScripts(scripts = { @BMScript("trace-ann-searcher.btm"), @BMScript("trace-hnsw-parameters.btm"), } ) class AnnSearchTest { EmbeddingClient embeddingClient; @BeforeEach void setUp() { embeddingClient = EmbeddingClient.create("localhost", 8000); } @AfterEach void tearDown() { embeddingClient.close(); } List<Document> createDocuments() { return List.of( createDocument( "The Time Machine", "A man travels through time and witnesses the evolution of humanity.", "H.G. Wells", 1895 ), createDocument( "Ender's Game", "A young boy is trained to become a military leader in a war against an alien race.", "Orson Scott Card", 1985 ), createDocument( "Brave New World", "A dystopian society where people are genetically engineered and conditioned to conform to a strict social hierarchy.", "Aldous Huxley", 1932 ), createDocument( "The Hitchhiker's Guide to the Galaxy", "A comedic science fiction series following the misadventures of an unwitting human and his alien friend.", "Douglas Adams", 1979 ), createDocument( "Dune", "A desert planet is the site of political intrigue and power struggles.", "Frank Herbert", 1965 ), createDocument( "Foundation", "A mathematician develops a science to predict the future of humanity and works to save civilization from collapse.", "Isaac Asimov", 1951 ), createDocument( "Snow Crash", "A futuristic world where the internet has evolved into a virtual reality metaverse.", "Neal Stephenson", 1992 ), createDocument( "Neuromancer", "A hacker is hired to pull off a near-impossible hack and gets pulled into a web of intrigue.", "William Gibson", 1984 ), createDocument( "The War of the Worlds", "A Martian invasion of Earth throws humanity into chaos.", "H.G. Wells", 1898 ), createDocument( "The Hunger Games", "A dystopian society where teenagers are forced to fight to the death in a televised spectacle.", "Suzanne Collins", 2008 ), createDocument( "The Andromeda Strain", "A deadly virus from outer space threatens to wipe out humanity.", "Michael Crichton", 1969 ), createDocument( "The Left Hand of Darkness", "A human ambassador is sent to a planet where the inhabitants are genderless and can change gender at will.", "Ursula K. Le Guin", 1969 ), createDocument( "The Three-Body Problem", "Humans encounter an alien civilization that lives in a dying system.", "Liu Cixin", 2008 ) ); } Document createDocument(String name, String description, String author, int year) { Document document = new Document(); document.add(new TextField("name", name, Field.Store.YES)); document.add(new TextField("description", description, Field.Store.YES)); document.add(new TextField("author", author, Field.Store.YES)); document.add(new IntField("year", year, Field.Store.YES)); float[] vector = textToVector("passage: " + description); document.add(new KnnFloatVectorField("description_vector", vector, VectorSimilarityFunction.EUCLIDEAN)); return document; } float[] textToVector(String text) { EmbeddingClient.EmbeddingRequest request = new EmbeddingClient.EmbeddingRequest( "intfloat/multilingual-e5-base", text ); EmbeddingClient.EmbeddingResponse response = embeddingClient.execute(request); float[] vector = new float[response.embedding().size()]; for (int i = 0; i < response.embedding().size(); i++) { vector[i] = response.embedding().get(i); } return vector; } @Test void knnSearch() throws IOException { try (Directory directory = new ByteBuffersDirectory()) { IndexWriterConfig config = new IndexWriterConfig(); List<Document> documents = createDocuments(); try (IndexWriter writer = new IndexWriter(directory, config)) { writer.addDocuments(documents); } try (DirectoryReader reader = DirectoryReader.open(directory)) { IndexSearcher searcher = new IndexSearcher(reader); int k = documents.size(); // all document float[] vector = textToVector("query: alien invasion"); KnnFloatVectorQuery query = new KnnFloatVectorQuery("description_vector", vector, k); TopDocs topDocs = searcher.search(query, 3); StoredFields storedFields = searcher.storedFields(); List<Document> resultDocuments = Arrays.stream(topDocs.scoreDocs) .map(scoreDoc -> { try { return storedFields.document(scoreDoc.doc); } catch (IOException e) { throw new UncheckedIOException(e); } }) .toList(); assertThat(resultDocuments).hasSize(3); assertThat(resultDocuments.get(0).getField("name").stringValue()) .isEqualTo("The Hitchhiker's Guide to the Galaxy"); assertThat(resultDocuments.get(0).getField("year").numericValue().intValue()) .isEqualTo(1979); assertThat(resultDocuments.get(1).getField("name").stringValue()) .isEqualTo("The Three-Body Problem"); assertThat(resultDocuments.get(1).getField("year").numericValue().intValue()) .isEqualTo(2008); assertThat(resultDocuments.get(2).getField("name").stringValue()) .isEqualTo("The Andromeda Strain"); assertThat(resultDocuments.get(2).getField("year").numericValue().intValue()) .isEqualTo(1969); } } } @Test void annSearch() throws IOException { try (Directory directory = new ByteBuffersDirectory()) { IndexWriterConfig config = new IndexWriterConfig(); List<Document> documents = createDocuments(); try (IndexWriter writer = new IndexWriter(directory, config)) { writer.addDocuments(documents); } try (DirectoryReader reader = DirectoryReader.open(directory)) { IndexSearcher searcher = new IndexSearcher(reader); int k = documents.size() - 1; // all documents - 1 float[] vector = textToVector("query: alien invasion"); KnnFloatVectorQuery query = new KnnFloatVectorQuery("description_vector", vector, k); TopDocs topDocs = searcher.search(query, 3); StoredFields storedFields = searcher.storedFields(); List<Document> resultDocuments = Arrays.stream(topDocs.scoreDocs) .map(scoreDoc -> { try { return storedFields.document(scoreDoc.doc); } catch (IOException e) { throw new UncheckedIOException(e); } }) .toList(); assertThat(resultDocuments).hasSize(3); assertThat(resultDocuments.get(0).getField("name").stringValue()) .isEqualTo("The Hitchhiker's Guide to the Galaxy"); assertThat(resultDocuments.get(0).getField("year").numericValue().intValue()) .isEqualTo(1979); assertThat(resultDocuments.get(1).getField("name").stringValue()) .isEqualTo("The Three-Body Problem"); assertThat(resultDocuments.get(1).getField("year").numericValue().intValue()) .isEqualTo(2008); assertThat(resultDocuments.get(2).getField("name").stringValue()) .isEqualTo("The Andromeda Strain"); assertThat(resultDocuments.get(2).getField("year").numericValue().intValue()) .isEqualTo(1969); } } } @Test void defaultCodec() throws ReflectiveOperationException { Codec codec = Codec.getDefault(); assertThat(codec).isExactlyInstanceOf(Lucene99Codec.class); assertThat(codec.knnVectorsFormat()).isInstanceOf(PerFieldKnnVectorsFormat.class); Lucene99Codec lucene99Codec = (Lucene99Codec) codec; java.lang.reflect.Field defaultKnnVectorsFormatField = Lucene99Codec.class.getDeclaredField("defaultKnnVectorsFormat"); defaultKnnVectorsFormatField.setAccessible(true); KnnVectorsFormat knnVectorsFormat = (KnnVectorsFormat) defaultKnnVectorsFormatField.get(lucene99Codec); assertThat(knnVectorsFormat).isExactlyInstanceOf(Lucene99HnswVectorsFormat.class); } @Test void customizeHsnwParamersCodec1() { Codec codec = Codec.forName("MyCustom"); // load from ServiceLoader assertThat(codec).isExactlyInstanceOf(MyCustomCodec.class); } @Test void customizeHsnwParamersCodec2() { Codec codec = new MyCustomCodec(); } Document createDocument2(String name, String description, String author, int year) { Document document = new Document(); document.add(new TextField("name", name, Field.Store.YES)); document.add(new TextField("description", description, Field.Store.YES)); document.add(new TextField("author", author, Field.Store.YES)); document.add(new IntField("year", year, Field.Store.YES)); float[] vector = textToVector("passage: " + description); document.add(new KnnFloatVectorField("description_vector", vector, VectorSimilarityFunction.EUCLIDEAN)); document.add(new KnnFloatVectorField("description_vector2", vector, VectorSimilarityFunction.EUCLIDEAN)); return document; } @Test void customizeHnswParameters() throws IOException { try (Directory directory = new ByteBuffersDirectory()) { IndexWriterConfig config = new IndexWriterConfig(); Codec codec = Codec.forName("MyCustom"); config.setCodec(codec); List<Document> documents = List.of( createDocument2( "The Time Machine", "A man travels through time and witnesses the evolution of humanity.", "H.G. Wells", 1895 ), createDocument2( "Ender's Game", "A young boy is trained to become a military leader in a war against an alien race.", "Orson Scott Card", 1985 ), createDocument2( "Brave New World", "A dystopian society where people are genetically engineered and conditioned to conform to a strict social hierarchy.", "Aldous Huxley", 1932 ) ); try (IndexWriter writer = new IndexWriter(directory, config)) { writer.addDocuments(documents); } try (DirectoryReader reader = DirectoryReader.open(directory)) { IndexSearcher searcher = new IndexSearcher(reader); int k = 1; // ANN float[] vector = textToVector("query: alien invasion"); KnnFloatVectorQuery query = new KnnFloatVectorQuery("description_vector", vector, k); TopDocs topDocs = searcher.search(query, 1); StoredFields storedFields = searcher.storedFields(); List<Document> resultDocuments = Arrays.stream(topDocs.scoreDocs) .map(scoreDoc -> { try { return storedFields.document(scoreDoc.doc); } catch (IOException e) { throw new UncheckedIOException(e); } }) .toList(); assertThat(resultDocuments).hasSize(1); assertThat(resultDocuments.get(0).getField("name").stringValue()) .isEqualTo("Ender's Game"); assertThat(resultDocuments.get(0).getField("year").numericValue().intValue()) .isEqualTo(1985); } } } }
おわりに
Apache LuceneでkNN検索とANNの切り替え方および、ANNで使われるHNSWのパラメーターの設定方法を見てみました。
けっこう大変だったのですが、使っているのがkNN検索なのかANNなのかぼやっとしていることが多かったので、こうやって切り替え条件が
わかるとだいぶクリアになりました。
また、mやefConstructionのカスタマイズ方法やどのように反映されるのかが確認できたのもよかったです。
