これは、なにをしたくて書いたもの？
ElasticsearchやApache Solrではベクトル検索ができるようです。となると、Apache Luceneにその基礎があるはずなので、ちょっと
見ておこうかなと思いまして。
今回はkNN検索を対象としているのですが、ANNに切り替わることについては別エントリーで書いています。
Apache LuceneでkNN検索とANN＋HNSWを使い分ける（Codecを使ったHNSWのパラメーター設定付き） - CLOVER🍀
KnnByteVectorField、KnnFloatVectorField、VectorSimilarityFunction、KnnByteVectorQuery、KnnFloatVectorQuery
Apache Luceneでベクトル検索（kNN検索）をするために使うクラスを紹介していきます。
まずはFieldから。保存するベクトルの型がbyteの配列なのかfloatの配列７なのかでKnnByteVectorField、KnnFloatVectorFieldの
2つのFieldクラスを使い分けます。
これらのクラスは、フィールド名、ベクトル、ベクトル類似度関数を指定してインスタンスを生成します。
また実験的APIのようです。
WARNING: This API is experimental and might change in incompatible ways in the next release.
またベクトルの次元数は1024までです。
IllegalArgumentException - if any parameter is null, or the vector is empty or has dimension > 1024.
ベクトル類似度関数はVectorSimilarityFunction列挙型で指定します。
VectorSimilarityFunction (Lucene 9.10.0 core API)
指定できる定数は以下の4つです。
COSINE… コサイン類似度- すべてのベクトルを単位長に正規化し、代わりに
DOT_PRODUCTを用いることを推奨 - 元のベクトルを保存する必要があり、事前に正規化できない場合にのみ使用すること
- すべてのベクトルを単位長に正規化し、代わりに
DOT_PRODUCT… ドット積- コサイン類似度を最適化したものとして意図されている
- ドキュメントとクエリーで使うベクトルの両方を正規化する必要がある
EUCLIDEAN… ユークリッド距離MAXIMUM_INNER_PRODUCT… 最大内積DOT_PRODUCTに似ているが、入力の正規化は必要ない- ベクトルの大きさに関する情報が、埋め込みベクトルストアにとって有用になることがある
ちなみに、KnnByteVectorField、KnnFloatVectorFieldのどちらもベクトル類似度関数の指定は省略することができるのですが、
その場合はユークリッド距離がデフォルトで設定されます。
最後にクエリーです。以下の4つがあるようです。
- KnnByteVectorQuery (Lucene 9.10.0 core API)
- KnnFloatVectorQuery (Lucene 9.10.0 core API)
- ByteVectorSimilarityQuery (Lucene 9.10.0 core API)
- FloatVectorSimilarityQuery (Lucene 9.10.0 core API)
KnnByteVectorQueryとKnnFloatVectorQueryはkNN検索を実行するクエリーです。
検索対象のフィールド名、ベクトル、検索するドキュメントの数（k）を指定します。追加でフィルターを行うQueryを指定することも
できます。
動作としては以下のようです。
※フィルターのコストによって検索方法が変わるのですが、これがkNN検索なのかANNなのかの分岐になっているように思います
- フィルターコストがk未満の場合は、完全一致検索を実行する
- それ以外の場合は、フィルターに従ってkNN検索を実行する
- kNN検索が完了せずに探索するベクトルが多すぎる場合は、停止して完全一致検索を実行する
2つのクエリーの違いは、対象となるFieldがKnnByteVectorFieldなのかKnnFloatVectorFieldなのか、byte配列なのかfloat配列なのか
ですね。
ByteVectorSimilarityQuery、FloatVectorSimilarityQueryは類似性のしきい値を超えるすべての（およその、approximate）ベクトルを
検索するクエリーです。この2つは実験的APIのようです。
WARNING: This API is experimental and might change in incompatible ways in the next release.
これまで何度かベクトル化をやってきた時に見てきたのはすべて小数のベクトルだったので、今回はKnnFloatVectorFieldと
KnnFloatVectorQueryを使ってみたいと思います。
HNSW
今回は扱いませんが、ANNで使う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)
Apache LuceneにおけるANNやHNSWについては、こちらのエントリーで書きました。
Apache LuceneでkNN検索とANN＋HNSWを使い分ける（Codecを使ったHNSWのパラメーター設定付き） - CLOVER🍀
環境
今回の環境はこちら。
$ java --version openjdk 21.0.2 2024-01-16 OpenJDK Runtime Environment (build 21.0.2+13-Ubuntu-122.04.1) OpenJDK 64-Bit Server VM (build 21.0.2+13-Ubuntu-122.04.1, mixed mode, sharing) $ mvn --version Apache Maven 3.9.6 (bc0240f3c744dd6b6ec2920b3cd08dcc295161ae) Maven home: $HOME/.sdkman/candidates/maven/current Java version: 21.0.2, vendor: Private Build, 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で行うことにします。FastAPIを使って、ベクトル化を行う簡単なAPIを作ります。
$ pip3 install sentence-transformers fastapi uvicorn[standard]
インストールされたライブラリー。
$ pip3 list Package Version ------------------------ ---------- annotated-types 0.6.0 anyio 4.3.0 certifi 2024.2.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.3 filelock 3.14.0 fsspec 2024.5.0 h11 0.14.0 httpcore 1.0.5 httptools 0.6.1 httpx 0.27.0 huggingface-hub 0.23.0 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.4.127 nvidia-nvtx-cu12 12.1.105 orjson 3.10.3 packaging 24.0 pillow 10.3.0 pip 22.0.2 pydantic 2.7.1 pydantic_core 2.18.2 Pygments 2.18.0 python-dotenv 1.0.1 python-multipart 0.0.9 PyYAML 6.0.1 regex 2024.5.15 requests 2.31.0 rich 13.7.1 safetensors 0.4.3 scikit-learn 1.4.2 scipy 1.13.0 sentence-transformers 2.7.0 setuptools 59.6.0 shellingham 1.5.4 sniffio 1.3.1 starlette 0.37.2 sympy 1.12 threadpoolctl 3.5.0 tokenizers 0.19.1 torch 2.3.0 tqdm 4.66.4 transformers 4.40.2 triton 2.3.0 typer 0.12.3 typing_extensions 4.11.0 ujson 5.10.0 urllib3 2.2.1 uvicorn 0.29.0 uvloop 0.19.0 watchfiles 0.21.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
確認。
$ curl -s -XPOST -H 'Content-Type: application/json' localhost:8000/embeddings/encode -d '{"model": "intfloat/multilingual-e5-base", "text": "query: Hello World"}' | jq
{
"model": "intfloat/multilingual-e5-base",
"embedding": [
0.03324141725897789,
0.04988044500350952,
0.00241446984000504,
〜省略〜
-0.04364151135087013,
-0.04888230562210083,
0.03604992479085922
],
"dimension": 768
}
こんな感じです。ベクトル化はこちらを呼び出して行うことにします。
準備
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.25.3</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> </dependencies>
Apache Luceneのベクトル検索を使うには、lucene-coreがあればOKです。
Jacksonは先ほど作成したAPIを呼び出す際に、オブジェクトをJSONに変換するために使用します。
次に、テキストをベクトル化するAPIを呼び出すクライアントを作成。
src/test/java/org/littlewings/lucene/knn/EmbeddingClient.java
package org.littlewings.lucene.knn; 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/knn/KnnSearchTest.java
package org.littlewings.lucene.knn; import org.apache.lucene.document.*; import org.apache.lucene.index.*; import org.apache.lucene.search.*; import org.apache.lucene.store.ByteBuffersDirectory; import org.apache.lucene.store.Directory; 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; class KnnSearchTest { 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; } // ここに、データ作成とテストを書く }
こちらはテキストをベクトル化するメソッドです。
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
データを作成する
まずはデータを作成しましょう。いい例がないかなと思ったのですが、Qdrantのチュートリアルを元にすることにしました。
こんな感じですね。
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
)
);
}
データは全部で13個です。
createDocumentメソッドの中身はこちら。
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; }
descriptionに対してベクトル化し、KnnFloatVectorFieldでFieldとして登録しています。ベクトル類似度関数はユークリッド距離としました。
float[] vector = textToVector("passage: " + description); document.add(new KnnFloatVectorField("description_vector", vector, VectorSimilarityFunction.EUCLIDEAN));
ちなみにpassage:というのはintfloat/multilingual-e5でドキュメントに指定する接頭辞です。
ベクトル検索（knn検索）を試す
では、ベクトル検索（knn検索）を試してみましょう。
まずはこちら。
@Test void gettingStarted() throws IOException { try (Directory directory = new ByteBuffersDirectory()) { IndexWriterConfig config = new IndexWriterConfig(); try (IndexWriter writer = new IndexWriter(directory, config)) { writer.addDocuments(createDocuments()); } try (DirectoryReader reader = DirectoryReader.open(directory)) { IndexSearcher searcher = new IndexSearcher(reader); float[] vector = textToVector("query: alien invasion"); KnnFloatVectorQuery query = new KnnFloatVectorQuery("description_vector", vector, 15); // KnnFloatVectorQuery query = KnnFloatVectorField.newVectorQuery("description_vector", vector, 15); 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); } } }
Apache LuceneのDirectoryを作成して、データを登録します。
try (Directory directory = new ByteBuffersDirectory()) { IndexWriterConfig config = new IndexWriterConfig(); try (IndexWriter writer = new IndexWriter(directory, config)) { writer.addDocuments(createDocuments()); }
KnnFloatVectorQueryを使って検索します。KnnFloatVectorQueryはインスタンスを直接作成してもいいですし、
KnnFloatVectorField#newVectorQueryで作成することもできます。
try (DirectoryReader reader = DirectoryReader.open(directory)) { IndexSearcher searcher = new IndexSearcher(reader); float[] vector = textToVector("query: alien invasion"); KnnFloatVectorQuery query = new KnnFloatVectorQuery("description_vector", vector, 15); // KnnFloatVectorQuery query = KnnFloatVectorField.newVectorQuery("description_vector", vector, 15);
3つ目の引数は、ベクトル検索で検索対象となるドキュメントの数です。今回は全部のドキュメントが入るようにしました。
ここの数を大きくすると検索精度は上がるはずですが、トレードオフとして速度が落ちるという結果になるはずです。
ちなみにquery:というのはintfloat/multilingual-e5で検索時に使う接頭辞です。
float[] vector = textToVector("query: alien invasion");
検索して、上位3件を取得してアサーション。
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); }
Qdrantのチュートリアルと近い結果になっています。なお、Qdrantのチュートリアルではベクトル類似度関数はコサイン類似度に
なっています。
最後に、KnnFloatVectorQueryにQueryを追加できるようなので、こちらも試しておきます。
@Test void withQueryAsFilter() throws IOException { try (Directory directory = new ByteBuffersDirectory()) { IndexWriterConfig config = new IndexWriterConfig(); try (IndexWriter writer = new IndexWriter(directory, config)) { writer.addDocuments(createDocuments()); } try (DirectoryReader reader = DirectoryReader.open(directory)) { IndexSearcher searcher = new IndexSearcher(reader); float[] vector = textToVector("query: alien invasion"); Query filter = IntPoint.newRangeQuery("year", 2000, Integer.MAX_VALUE); KnnFloatVectorQuery query = new KnnFloatVectorQuery( "description_vector", vector, 15, filter ); 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(2); assertThat(resultDocuments.get(0).getField("name").stringValue()) .isEqualTo("The Three-Body Problem"); assertThat(resultDocuments.get(0).getField("year").numericValue().intValue()) .isEqualTo(2008); assertThat(resultDocuments.get(1).getField("name").stringValue()) .isEqualTo("The Hunger Games"); assertThat(resultDocuments.get(1).getField("year").numericValue().intValue()) .isEqualTo(2008); } } }
先ほどとの違いはこちらですね。
Query filter = IntPoint.newRangeQuery("year", 2000, Integer.MAX_VALUE); KnnFloatVectorQuery query = new KnnFloatVectorQuery( "description_vector", vector, 15, filter );
Queryを作成し、KnnFloatVectorQueryのコンストラクターの第4引数に指定します。こちらは「フィルター」という位置づけになり、
ベクトル検索の前に適用されるようです。
KnnFloatVectorQuery (Lucene 9.10.0 core API)
今回は2000年以上のデータに絞り込んでいます。
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(2); assertThat(resultDocuments.get(0).getField("name").stringValue()) .isEqualTo("The Three-Body Problem"); assertThat(resultDocuments.get(0).getField("year").numericValue().intValue()) .isEqualTo(2008); assertThat(resultDocuments.get(1).getField("name").stringValue()) .isEqualTo("The Hunger Games"); assertThat(resultDocuments.get(1).getField("year").numericValue().intValue()) .isEqualTo(2008);
このテストコード全体も載せておきます。
src/test/java/org/littlewings/lucene/knn/KnnSearchTest.java
package org.littlewings.lucene.knn; import org.apache.lucene.document.*; import org.apache.lucene.index.*; import org.apache.lucene.search.*; import org.apache.lucene.store.ByteBuffersDirectory; import org.apache.lucene.store.Directory; 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; class KnnSearchTest { 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; } 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; } @Test void gettingStarted() throws IOException { try (Directory directory = new ByteBuffersDirectory()) { IndexWriterConfig config = new IndexWriterConfig(); try (IndexWriter writer = new IndexWriter(directory, config)) { writer.addDocuments(createDocuments()); } try (DirectoryReader reader = DirectoryReader.open(directory)) { IndexSearcher searcher = new IndexSearcher(reader); float[] vector = textToVector("query: alien invasion"); KnnFloatVectorQuery query = new KnnFloatVectorQuery("description_vector", vector, 15); // KnnFloatVectorQuery query = KnnFloatVectorField.newVectorQuery("description_vector", vector, 15); 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 withQueryAsFilter() throws IOException { try (Directory directory = new ByteBuffersDirectory()) { IndexWriterConfig config = new IndexWriterConfig(); try (IndexWriter writer = new IndexWriter(directory, config)) { writer.addDocuments(createDocuments()); } try (DirectoryReader reader = DirectoryReader.open(directory)) { IndexSearcher searcher = new IndexSearcher(reader); float[] vector = textToVector("query: alien invasion"); Query filter = IntPoint.newRangeQuery("year", 2000, Integer.MAX_VALUE); KnnFloatVectorQuery query = new KnnFloatVectorQuery( "description_vector", vector, 15, filter ); 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(2); assertThat(resultDocuments.get(0).getField("name").stringValue()) .isEqualTo("The Three-Body Problem"); assertThat(resultDocuments.get(0).getField("year").numericValue().intValue()) .isEqualTo(2008); assertThat(resultDocuments.get(1).getField("name").stringValue()) .isEqualTo("The Hunger Games"); assertThat(resultDocuments.get(1).getField("year").numericValue().intValue()) .isEqualTo(2008); } } } }
全部で13個のデータの登録と1回のクエリーを実行するテストですが、このテストを実行するとまあまあ時間がかかります。
$ mvn test
約2分ですね。
[INFO] Tests run: 2, Failures: 0, Errors: 0, Skipped: 0, Time elapsed: 101.9 s -- in org.littlewings.lucene.knn.KnnSearchTest [INFO] [INFO] Results: [INFO] [INFO] Tests run: 2, Failures: 0, Errors: 0, Skipped: 0 [INFO] [INFO] ------------------------------------------------------------------------ [INFO] BUILD SUCCESS [INFO] ------------------------------------------------------------------------ [INFO] Total time: 01:43 min [INFO] Finished at: 2024-05-19T00:19:10+09:00 [INFO] ------------------------------------------------------------------------
CPU環境ということもあって、ベクトル化が遅いのです…。あと、まとめてベクトルに変換した方がいいとかもあると思いますが。
ところで、このテストを実行するとこういう警告がApache Luceneから出力されます。
5月 19, 2024 12:10:39 午前 org.apache.lucene.internal.vectorization.VectorizationProvider lookup 警告: Java vector incubator module is not readable. For optimal vector performance, pass '--add-modules jdk.incubator.vector' to enable Vector API.
JEP 448: Vector API (Sixth Incubator)
こうしてみましょう。
$ mvn test -DargLine='--add-modules jdk.incubator.vector'
これで、Vector API（jdk.incubator.vectorモジュール）をApache Luceneが使うようです。
WARNING: Using incubator modules: jdk.incubator.vector [INFO] Running org.littlewings.lucene.knn.KnnSearchTest 5月 19, 2024 12:23:33 午前 org.apache.lucene.internal.vectorization.PanamaVectorizationProvider <init> 情報: Java vector incubator API enabled; uses preferredBitSize=256; FMA enabled
どうなんでしょう？誤差のような気もします。ふつうに逆転したりもしましたし。
[INFO] Tests run: 2, Failures: 0, Errors: 0, Skipped: 0, Time elapsed: 96.68 s -- in org.littlewings.lucene.knn.KnnSearchTest [INFO] [INFO] Results: [INFO] [INFO] Tests run: 2, Failures: 0, Errors: 0, Skipped: 0 [INFO] [INFO] ------------------------------------------------------------------------ [INFO] BUILD SUCCESS [INFO] ------------------------------------------------------------------------ [INFO] Total time: 01:38 min [INFO] Finished at: 2024-05-19T00:25:05+09:00 [INFO] ------------------------------------------------------------------------
とりあえず、今回はこんなところで。
おわりに
Apache Luceneでベクトル検索（kNN検索）を試してみました。
ほとんど情報がなかったので、どれを使ったらいいのかを調べたりなどいろいろ苦労しましたが、最終的にはなんとかなったので
まあいいかなと。
それにしても、ベクトル化やベクトル検索に対する知識がだいぶ足りないので、こういうものに触れつつ学んでいかないとですね。
