sentence-transformersでEmbeddingモデルを実行

1. 概要

Embeddingは文を数ベクトルに変換する方法です。意味が似ている2つの文は、似ているベクトルを生成します。これはsemantic search、RAG、テキストクラスタリング、コンテンツ推奨の基盤です。

このチュートリアルでは、sentence-transformersライブラリを使用して、Hugging Faceのembeddingモデルをローカルマシンで実行します。読み終わったら、embeddingの作成方法、類似度の計算、semantic searchの実装、テストの作成方法、およびSpring Bootアプリケーションから呼び出すHTTPサービスとしてのモデルの公開方法を理解できるようになります。

2. インストール

sentence-transformersはPythonライブラリで、PyTorch上で実行されます。専用のvirtual environmentを作成することをお勧めします:

python -m venv .venv
source .venv/bin/activate
pip install sentence-transformers

このコマンドはtorch、transformers、huggingface_hubをインストールします。マシンにNVIDIA GPUがある場合は、まずCUDA対応のtorchをインストールしてから、sentence-transformersをインストールしてください。

モデルを初めてロードするとき、ライブラリはそれをHugging Face Hubからキャッシュディレクトリ(デフォルトは~/.cache/huggingface)にダウンロードします。その後、モデルはキャッシュから直接読み込まれます。

3. 最初のembeddingを作成

all-MiniLM-L6-v2から始めます。これは小さなモデル(約90 MB)で、高速であり、英語テキストの実験に適しています:

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

sentences = [
    "The order was shipped yesterday.",
    "Your package left our warehouse last night.",
    "I would like to change my password.",
]

embeddings = model.encode(sentences)

print(embeddings.shape)
print(model.get_sentence_embedding_dimension())

説明:

  • SentenceTransformer(...)はHubからモデルとトークナイザーをロードするか、すでにある場合はキャッシュからロードします。
  • model.encode(sentences)は文字列のリストを受け取り、numpy.ndarrayを返します。各文は1行に対応します。
  • get_sentence_embedding_dimension()はベクトルの次元数を示します。

結果:

(3, 384)
384

3つの文があり、各文は384次元のベクトルになります。異なるモデルは異なる次元数を持ちます。例えば、768または1024などです。

3.1. ベトナム語テキスト

all-MiniLM-L6-v2は主に英語で学習されています。ベトナム語では、多言語モデルを使用する必要があります:

model = SentenceTransformer(
    "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
)

vectors = model.encode([
    "Đơn hàng đã được giao cho đơn vị vận chuyển.",
    "Gói hàng của bạn đã rời kho tối qua.",
])
print(vectors.shape)
(2, 384)

このモデルも384次元のベクトルを生成します。ベトナム語を含む約50の言語を理解しています。

4. 類似度の計算

ベクトルは、それらを相互に比較してこそ有用です。最も一般的な尺度はcosine similarityです。バージョン3以降、ライブラリにはmodel.similarityが組み込まれています:

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

sentences = [
    "The order was shipped yesterday.",
    "Your package left our warehouse last night.",
    "I would like to change my password.",
]
embeddings = model.encode(sentences)

scores = model.similarity(embeddings, embeddings)
print(scores)

model.similarity(a, b)はtorch.Tensorマトリックスを返します。セル[i][j]は文iと文jの間の類似度です。出力は以下のようになります(バージョンによって数値がわずかに異なる場合があります):

tensor([[1.0000, 0.6123, 0.0412],
        [0.6123, 1.0000, 0.0587],
        [0.0412, 0.0587, 1.0000]])

対角線は常に1です。各文はそれ自身と完全に一致するためです。配送に関する2つの文は高いスコアを持ちますが、ほとんど単語が一致していません。パスワード変更に関する文は、ほぼ無関係です。

4.1. ベクトルの正規化

ベクトルをベクトルデータベースに保存する場合は、長さ1に正規化してください:

embeddings = model.encode(sentences, normalize_embeddings=True)
score = embeddings[0] @ embeddings[1]
print(round(float(score), 4))
0.6123

ベクトルが正規化されると、内積(@)はcosine similarityの正確な値を与えます。この操作はより効率的であり、pgvector、Qdrant、Elasticsearchなど多くのデータベースはこのために特別に最適化されています。

実際の問題は、顧客の質問に対して、FAQストレージから最適な回答を見つけることです。util.semantic_search関数がこれを実行します:

from sentence_transformers import SentenceTransformer, util

model = SentenceTransformer(
    "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
)

faq = [
    "Làm sao để theo dõi đơn hàng?",
    "Chính sách hoàn tiền như thế nào?",
    "Tôi quên mật khẩu đăng nhập.",
    "Phí vận chuyển ra nước ngoài là bao nhiêu?",
]
faq_embeddings = model.encode(faq, convert_to_tensor=True, normalize_embeddings=True)

query = "Hàng của tôi đang ở đâu rồi?"
query_embedding = model.encode(query, convert_to_tensor=True, normalize_embeddings=True)

hits = util.semantic_search(query_embedding, faq_embeddings, top_k=2)[0]
for hit in hits:
    print(f"{hit['score']:.3f}  {faq[hit['corpus_id']]}")

説明:

  • FAQ全体を1回エンコードします。実際のアプリケーションでは、このステップは起動時またはデータが変更されたときに実行されます。
  • convert_to_tensor=Trueは結果をテンソルの形で保持するため、numpyとの間の変換が不要になります。
  • semantic_searchは各クエリに対してリストを返します。各要素はcorpus_idとscoreを含むdictで、降順でソートされています。

サンプル出力:

0.721  Làm sao để theo dõi đơn hàng?
0.388  Phí vận chuyển ra nước ngoài là bao nhiêu?

このアプローチは、メモリ内に数十万のドキュメントがある場合に適しています。数百万のドキュメントがある場合は、ANN indexを備えたベクトルデータベースにしてください。

**注:**E5やBGEのような一部のモデルは追加のプレフィックスを要求します。例えば、質問には"query: "、ドキュメントには"passage: "です。プレフィックスを省略すると、検索品質が著しく低下する可能性があります。使用する前にモデルカードを読んでください。encodeにprompt="query: "を渡すことができます。

6. パフォーマンスの最適化

encodeにはパフォーマンスに直接影響するいくつかのパラメータがあります:

import torch
from sentence_transformers import SentenceTransformer

device = "cuda" if torch.cuda.is_available() else "cpu"
model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2", device=device)

documents = [f"Customer review number {i}" for i in range(10_000)]

embeddings = model.encode(
    documents,
    batch_size=64,
    show_progress_bar=True,
    normalize_embeddings=True,
)
print(embeddings.shape, device)
Batches: 100%|██████████| 157/157 [00:09<00:00, 16.8it/s]
(10000, 384) cpu
  • deviceはCPUまたはGPUを選択します。GPUでは、速度が数十倍高速になることがあります。
  • batch_sizeのデフォルトは32です。GPUでは、64~256に増やすと通常、VRAMが不足しない限り、より高速に実行されます。
  • 常にリスト全体をencodeに送信してください。forループで1文ずつencodeを呼び出すと、数倍遅くなります。

各モデルには入力長の制限があります。model.max_seq_lengthを参照してください。all-MiniLM-L6-v2の場合、この値は256トークンです。超過したテキスト部分は警告なく自動的にカットされます。したがって、長いドキュメントはエンコードする前にチャンク化する必要があります。

7. Embeddingのテストを作成

Embeddingは実数なので、特定の値と1対1で比較することはできません。具体的な値ではなく、特性をテストしてください。pytestをインストールします:

pip install pytest
# test_embeddings.py
import numpy as np
import pytest
from sentence_transformers import SentenceTransformer


@pytest.fixture(scope="module")
def model():
    return SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")


def test_embedding_has_expected_dimension(model):
    vectors = model.encode(["Order #42 is delayed"])
    assert vectors.shape == (1, 384)


def test_normalized_embedding_has_unit_length(model):
    vector = model.encode("Refund request", normalize_embeddings=True)
    assert np.linalg.norm(vector) == pytest.approx(1.0, abs=1e-5)


def test_similar_sentences_score_higher(model):
    shipping = model.encode("The order was shipped yesterday.")
    package = model.encode("Your package left our warehouse last night.")
    password = model.encode("I would like to change my password.")

    related = model.similarity(shipping, package).item()
    unrelated = model.similarity(shipping, password).item()
    assert related > unrelated

Fixtureはscope="module"を持つため、モデルはファイル全体で1回だけロードされます。3番目のテストは、誰かが不適切なモデルに変更した場合のエラーをキャッチします。

8. Spring Bootアプリケーションからモデルを呼び出す

バックエンドのほとんどはJavaで書かれているかもしれません。シンプルな方法は、FastAPIを使用してモデルを小さなHTTPサービスでラップし、Spring Bootから呼び出すことです。

8.1. FastAPIでのEmbeddingサービス

pip install fastapi uvicorn
# app.py
from fastapi import FastAPI
from pydantic import BaseModel
from sentence_transformers import SentenceTransformer

app = FastAPI()
model = SentenceTransformer(
    "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
)


class EmbedRequest(BaseModel):
    texts: list[str]


@app.post("/embed")
def embed(request: EmbedRequest):
    vectors = model.encode(request.texts, normalize_embeddings=True)
    return {
        "embeddings": vectors.tolist(),
        "dimension": model.get_sentence_embedding_dimension(),
    }
uvicorn app:app --port 8000

モデルはモジュール起動時に1回ロードされ、リクエストごとに再ロードされません。エンドポイントはasync defではなくdefを使用します。理由は、encodeがブロッキング操作であり、FastAPIがdef関数を別のスレッドプール内で実行するため、イベントループがブロックされないからです。

8.2. Spring BootのクライアントDependency Maven:

<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-web</artifactId>
</dependency>
<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-test</artifactId>
    <scope>test</scope>
</dependency>
embedding.base-url=http://localhost:8000
public record EmbedRequest(List<String> texts) {}

public record EmbedResponse(List<List<Float>> embeddings, int dimension) {}

@Service
public class EmbeddingClient {

    private final RestClient restClient;

    public EmbeddingClient(RestClient.Builder builder,
                           @Value("${embedding.base-url}") String baseUrl) {
        this.restClient = builder.baseUrl(baseUrl).build();
    }

    public List<List<Float>> embed(List<String> texts) {
        EmbedResponse response = restClient.post()
                .uri("/embed")
                .body(new EmbedRequest(texts))
                .retrieve()
                .body(EmbedResponse.class);
        return response.embeddings();
    }
}

RestClient(Spring Framework 6.1以降で利用可能)は、EmbedRequestをJSONとして送信し、JacksonでレスポンスをレコードEmbedResponseにマップします。結果を使用してpgvectorに保存したり、他のベクトルと比較したりできます。

8.3. MockRestServiceServerでのクライアントテスト

import static org.assertj.core.api.Assertions.assertThat;
import static org.springframework.test.web.client.match.MockRestRequestMatchers.*;
import static org.springframework.test.web.client.response.MockRestResponseCreators.withSuccess;

class EmbeddingClientTest {

    @Test
    void shouldReturnEmbeddingsFromService() {
        RestClient.Builder builder = RestClient.builder();
        MockRestServiceServer server = MockRestServiceServer.bindTo(builder).build();

        server.expect(requestTo("http://localhost:8000/embed"))
              .andExpect(method(HttpMethod.POST))
              .andExpect(jsonPath("$.texts[0]").value("Đơn hàng bị trễ"))
              .andRespond(withSuccess("""
                      {"embeddings": [[0.1, 0.2, 0.3]], "dimension": 3}
                      """, MediaType.APPLICATION_JSON));

        EmbeddingClient client = new EmbeddingClient(builder, "http://localhost:8000");

        List<List<Float>> result = client.embed(List.of("Đơn hàng bị trễ"));

        assertThat(result).hasSize(1);
        assertThat(result.getFirst()).containsExactly(0.1f, 0.2f, 0.3f);
        server.verify();
    }
}

MockRestServiceServer.bindTo(builder)は、builderがRestClientを作成する前に呼び出す必要があります。このテストでは、クライアントがコンストラクタ内でのみbuild()を呼び出すため、これは常に当てはまります。これにより、テストは実際のPythonサービスなしで実行できます。

9. 一般的なエラー

  • 複数のモデルのベクトルを混ぜること。 異なる2つのモデルのベクトルは、次元数が同じであっても、異なる空間にあります。モデルを変更することは、すべてのデータを再度エンコードすることを意味します。
  • 一方の正規化を忘れること。 ドキュメントは正規化されているがクエリは正規化されていない場合、内積のスコアが偏ります。両側で同じencode設定を使用してください。
  • テキストが長すぎること。 max_seq_lengthを超える部分は警告なくカットされます。エンコード前にドキュメントをチャンク化してください。
  • 各リクエストでモデルをロードすること。 この操作には数秒とメモリが必要です。起動時に1回だけロードしてください。
  • コンテナにネットワークがないこと。 本番環境ではHubへのアクセスが通常ブロックされます。モデルをあらかじめイメージにダウンロードしてから、HF_HUB_OFFLINE=1を設定してください。

注:Embeddingは正確な検索を置き換えることはできません。注文番号、電話番号、またはSKUの場合、フルテキスト検索または通常のSQLクエリがより良い結果を提供します。実際には、両方を組み合わせた(ハイブリッド検索)方法が通常最も効果的です。

10. 結論

わずか数行のコードでsentence-transformersを使用してembeddingを作成しました。また、cosine similarityを計算し、FAQストレージのsemantic searchを構築し、batch_sizeとdeviceを調整してより高速に実行しました。さらに、特性によるembeddingのテスト方法を学び、モデルをHTTPサービスとしてラップしてSpring BootからRestClientで呼び出す方法を習得しました。

次の論理的なステップは、ベクトルをpgvectorなどのベクトルデータベースに保存し、RAGパイプラインのretrieval部分として使用することです。

huggingfacesentence-transformersembeddingpythonspring-boot