Transformers (ONNX) 嵌入

TransformersEmbeddingModel 是一個 EmbeddingModel 實作,它在本機計算 句子嵌入,使用選定的 句子轉換器

您可以使用任何 HuggingFace 嵌入模型

它使用 預先訓練 的轉換器模型,序列化為 開放神經網路交換 (ONNX) 格式。

使用 Deep Java Library 和 Microsoft ONNX Java Runtime 函式庫來執行 ONNX 模型,並在 Java 中計算嵌入。

先決條件

為了在 Java 中執行,我們需要將Tokenizer 和 Transformer 模型序列化ONNX 格式。

使用 optimum-cli 序列化 - 一個快速達成此目的的方法是使用 optimum-cli 命令列工具。以下程式碼片段準備了 Python 虛擬環境,安裝必要的套件,並使用 optimum-cli 序列化(例如,匯出)指定的模型

python3 -m venv venv
source ./venv/bin/activate
(venv) pip install --upgrade pip
(venv) pip install optimum onnx onnxruntime sentence-transformers
(venv) optimum-cli export onnx --model sentence-transformers/all-MiniLM-L6-v2 onnx-output-folder

此程式碼片段將 sentence-transformers/all-MiniLM-L6-v2 轉換器匯出到 onnx-output-folder 資料夾。後者包含嵌入模型使用的 tokenizer.jsonmodel.onnx 檔案。

您可以選擇任何 huggingface 轉換器識別碼,或提供直接檔案路徑來取代 all-MiniLM-L6-v2。

自動配置

Spring AI 為 ONNX Transformer 嵌入模型提供 Spring Boot 自動配置。若要啟用它,請將以下依賴項新增至專案的 Maven pom.xml 檔案

<dependency>
    <groupId>org.springframework.ai</groupId>
    <artifactId>spring-ai-transformers-spring-boot-starter</artifactId>
</dependency>

或 Gradle build.gradle 建置檔案。

dependencies {
    implementation 'org.springframework.ai:spring-ai-transformers-spring-boot-starter'
}
請參閱 相依性管理 章節,將 Spring AI BOM 新增至您的建置檔案。請參閱 儲存庫 章節,將這些儲存庫新增至您的建置系統。

若要配置它,請使用 spring.ai.embedding.transformer.* 屬性。

例如,將此新增至您的 application.properties 檔案,以使用 intfloat/e5-small-v2 文字嵌入模型來配置用戶端

spring.ai.embedding.transformer.onnx.modelUri=https://huggingface.co/intfloat/e5-small-v2/resolve/main/model.onnx
spring.ai.embedding.transformer.tokenizer.uri=https://huggingface.co/intfloat/e5-small-v2/raw/main/tokenizer.json

完整的支援屬性清單如下

嵌入屬性

屬性 描述 預設值

spring.ai.embedding.transformer.enabled

啟用 Transformer 嵌入模型。

true

spring.ai.embedding.transformer.tokenizer.uri

由 ONNX 引擎建立的預先訓練 HuggingFaceTokenizer 的 URI (例如 tokenizer.json)。

onnx/all-MiniLM-L6-v2/tokenizer.json

spring.ai.embedding.transformer.tokenizer.options

HuggingFaceTokenizer 選項,例如 'addSpecialTokens'、'modelMaxLength'、'truncation'、'padding'、'maxLength'、'stride'、'padToMultipleOf'。留空以回復為預設值。

空值

spring.ai.embedding.transformer.cache.enabled

啟用遠端資源快取。

true

spring.ai.embedding.transformer.cache.directory

快取遠端資源 (例如 ONNX 模型) 的目錄路徑

${java.io.tmpdir}/spring-ai-onnx-model

spring.ai.embedding.transformer.onnx.modelUri

現有的預先訓練 ONNX 模型。

onnx/all-MiniLM-L6-v2/model.onnx

spring.ai.embedding.transformer.onnx.modelOutputName

ONNX 模型的輸出節點名稱,我們將用於嵌入計算。

last_hidden_state

spring.ai.embedding.transformer.onnx.gpuDeviceId

要執行的 GPU 裝置 ID。僅適用於 >= 0。否則會忽略。(需要額外的 onnxruntime_gpu 相依性)

-1

spring.ai.embedding.transformer.metadataMode

指定文件內容和中繼資料的哪些部分將用於計算嵌入。

NONE

錯誤和特殊情況

如果您看到類似 Caused by: ai.onnxruntime.OrtException: Supplied array is ragged,.. 的錯誤,您也需要在 application.properties 中啟用 tokenizer padding,如下所示

spring.ai.embedding.transformer.tokenizer.options.padding=true

如果您收到類似 The generative output names don’t contain expected: last_hidden_state. Consider one of the available model outputs: token_embeddings, …​. 的錯誤,您需要將模型輸出名稱設定為模型正確的值。請考慮錯誤訊息中列出的名稱。例如

spring.ai.embedding.transformer.onnx.modelOutputName=token_embeddings

如果您收到類似 ai.onnxruntime.OrtException: Error code - ORT_FAIL - message: Deserialize tensor onnx::MatMul_10319 failed.GetFileLength for ./model.onnx_data failed:Invalid fd was supplied: -1 的錯誤,表示您的模型大於 2GB,並且序列化為兩個檔案:model.onnxmodel.onnx_data

model.onnx_data 稱為 外部資料,預期與 model.onnx 位於相同的目錄下。

目前唯一的解決方法是將大型 model.onnx_data 複製到您執行 Boot 應用程式的資料夾中。

如果您收到類似 ai.onnxruntime.OrtException: Error code - ORT_EP_FAIL - message: Failed to find CUDA shared provider 的錯誤,表示您正在使用 GPU 參數 spring.ai.embedding.transformer.onnx.gpuDeviceId,但缺少 onnxruntime_gpu 相依性。

<dependency>
    <groupId>com.microsoft.onnxruntime</groupId>
    <artifactId>onnxruntime_gpu</artifactId>
</dependency>

請根據 CUDA 版本 (ONNX Java Runtime) 選擇適當的 onnxruntime_gpu 版本。

手動配置

如果您未使用 Spring Boot,您可以手動配置 Onnx Transformers 嵌入模型。為此,請將 spring-ai-transformers 相依性新增至專案的 Maven pom.xml 檔案

<dependency>
  <groupId>org.springframework.ai</groupId>
  <artifactId>spring-ai-transformers</artifactId>
</dependency>
請參閱 相依性管理 章節,將 Spring AI BOM 新增至您的建置檔案。

然後建立新的 TransformersEmbeddingModel 實例,並使用 setTokenizerResource(tokenizerJsonUri)setModelResource(modelOnnxUri) 方法來設定匯出的 tokenizer.jsonmodel.onnx 檔案的 URI。(支援 classpath:file:https: URI 結構描述)。

如果模型未明確設定,TransformersEmbeddingModel 預設為 sentence-transformers/all-MiniLM-L6-v2

維度

384

平均效能

58.80

速度

14200 句/秒

大小

80MB

以下程式碼片段說明如何手動使用 TransformersEmbeddingModel

TransformersEmbeddingModel embeddingModel = new TransformersEmbeddingModel();

// (optional) defaults to classpath:/onnx/all-MiniLM-L6-v2/tokenizer.json
embeddingModel.setTokenizerResource("classpath:/onnx/all-MiniLM-L6-v2/tokenizer.json");

// (optional) defaults to classpath:/onnx/all-MiniLM-L6-v2/model.onnx
embeddingModel.setModelResource("classpath:/onnx/all-MiniLM-L6-v2/model.onnx");

// (optional) defaults to ${java.io.tmpdir}/spring-ai-onnx-model
// Only the http/https resources are cached by default.
embeddingModel.setResourceCacheDirectory("/tmp/onnx-zoo");

// (optional) Set the tokenizer padding if you see an errors like:
// "ai.onnxruntime.OrtException: Supplied array is ragged, ..."
embeddingModel.setTokenizerOptions(Map.of("padding", "true"));

embeddingModel.afterPropertiesSet();

List<List<Double>> embeddings = this.embeddingModel.embed(List.of("Hello world", "World is big"));
如果您手動建立 TransformersEmbeddingModel 的實例,則必須在設定屬性之後和使用用戶端之前呼叫 afterPropertiesSet() 方法。

第一次呼叫 embed() 會下載大型 ONNX 模型,並將其快取在本機檔案系統上。因此,第一次呼叫可能比平常花費更長的時間。使用 #setResourceCacheDirectory(<path>) 方法來設定儲存 ONNX 模型的本機資料夾。預設快取資料夾為 ${java.io.tmpdir}/spring-ai-onnx-model

將 TransformersEmbeddingModel 建立為 Bean 會更方便(也更受歡迎)。這樣您就不必手動呼叫 afterPropertiesSet()

@Bean
public EmbeddingModel embeddingModel() {
   return new TransformersEmbeddingModel();
}