追記
こっちの方が楽に実現できました。
導入
CTranslate2でEmbedding用モデルも扱いたいなーと考えていたところ、だいぶ前にIssueに上がっていることに気づきました。
Issue内で、作者自らsentence-transformersで使うサンプルコードまで作成されていましたので、少し変更しながら写経してみました。
検証はいつも通りDatabricksで行っています。とはいえ、Databricks外でも動くと思います。
また、GPUクラスタを利用しています。
Step1. モデルを変換
Embeddingのモデルはintfloat/multilingual-e5-largeを使います。
Huggingfaceからモデルスナップショットをダウンロードし、Unity Catalogのボリュームに保管してある状態から開始します。
今回は量子化をせずに変換することにします。
# 必要なモジュールのインストール
%pip install -U -qq transformers accelerate ctranslate2 sentence-transformers
dbutils.library.restartPython()
UC_VOLUME = "/Volumes/モデルのスナップショット保管ボリューム"
uc_dir = "/models--intfloat--multilingual-e5-large"
src_path = f"{UC_VOLUME}{uc_dir}"
# 変換後モデルの保存先ローカルパス
local_output_dir = "/tmp/ctranslate2/models--intfloat--multilingual-e5-large"
# 今回は量子化しない
!ct2-transformers-converter --model {src_path} --output_dir {local_output_dir} --low_cpu_mem_usage
最後に、変換したモデルをUnity Catalogのボリュームに保管します。
このとき、モデルスナップショットのサブフォルダct2内にコピーしてください(大事)
dbutils.fs.cp("file:"+local_output_dir, f"{UC_VOLUME}{uc_dir}/ct2", recurse=True)
Step2. Embeddingを計算
作者のサンプルコードをベースに写経。
変換済みモデルを利用できるように変更しています。
WARNINGコメントに記載されているように、PoCとして作られたもののようなので、モデルによっては動作しなかったりすると思われます。利用は自己責任で。
%pip install -U -qq transformers accelerate ctranslate2 sentence-transformers langchain
dbutils.library.restartPython()
"""Run SentenceTransformer.encode with CTranslate2.
WARNING! This script is a proof of concept using LaBSE and can easily break for other models:
* the selected model does not have a registered converter in CTranslate2
* the model requires outputs that are not returned by CTranslate2
* the sentence-transformers library changed the loading logic or code structure
(this example was tested with sentence-transformers==2.2.2 and transformers==4.29.2)
"""
"""邦訳
WARNING! このスクリプトはLaBSEを使用した概念実証であり、他のモデルでは簡単に壊れる可能性があります：
* 選択されたモデルはCTranslate2に登録されたコンバータを持っていません。
* モデルはCTranslate2によって返されない出力結果を必要とします。
* sentence-transformersライブラリは、ローディングロジックまたはコード構造を変更しました。
(この例は、sentence-transformers==2.2.2とtransformers==4.29.2でテストされました)
"""
import os
import ctranslate2
import numpy as np
import sentence_transformers
import torch
def main():
sentences = ["This is an example sentence", "Each sentence is converted"]
model = CT2SentenceTransformer("/Volumes/モデルのスナップショット保管ボリューム/models--intfloat--multilingual-e5-large")
embeddings = model.encode(sentences)
print(embeddings)
class CT2SentenceTransformer(sentence_transformers.SentenceTransformer):
"""Extension of sentence_transformers.SentenceTransformer using a CTranslate2 model."""
def __init__(self, *args, compute_type="default", **kwargs):
super().__init__(*args, **kwargs)
self[0] = CT2Transformer(self[0], compute_type=compute_type)
class CT2Transformer(torch.nn.Module):
"""Wrapper around a sentence_transformers.models.Transformer which routes the forward
call to a CTranslate2 encoder model.
"""
def __init__(self, transformer, compute_type="default"):
super().__init__()
self.transformer = transformer
self.compute_type = compute_type
self.encoder = None
# 指定されたパスのct2サブフォルダを指定
self.ct2_model_dir = transformer.auto_model.config.name_or_path + "/ct2"
def children(self):
# Do not consider the "transformer" attribute as a child module so that it will stay on the CPU.
return []
def forward(self, features):
device = features["input_ids"].device
print(device)
if self.encoder is None:
# The encoder is lazy-loaded to correctly resolve the target device.
self.encoder = ctranslate2.Encoder(
self.ct2_model_dir,
device=device.type,
device_index=device.index or 0,
intra_threads=torch.get_num_threads(),
compute_type=self.compute_type,
)
input_ids = features["input_ids"].to(torch.int32)
length = features["attention_mask"].sum(1, dtype=torch.int32)
if device.type == "cpu":
# PyTorch CPU tensors do not implement the Array interface so a roundtrip to Numpy
# is required for both the input and output.
input_ids = input_ids.numpy()
length = length.numpy()
input_ids = ctranslate2.StorageView.from_array(input_ids)
length = ctranslate2.StorageView.from_array(length)
outputs = self.encoder.forward_batch(input_ids, length)
last_hidden_state = outputs.last_hidden_state
if device.type == "cpu":
last_hidden_state = np.array(last_hidden_state)
features["token_embeddings"] = torch.as_tensor(
last_hidden_state, device=device
).to(torch.float32)
return features
def tokenize(self, *args, **kwargs):
return self.transformer.tokenize(*args, **kwargs)
if __name__ == "__main__":
main()
出力
[[ 0.02693515 -0.01583192 -0.02195659 ... -0.02472507 -0.02929422
0.01645531]
[ 0.02027525 -0.01175944 -0.00298364 ... -0.01923108 -0.03131905
0.03152921]]
sentence-transformersで変換前モデルを読み込んだ結果とも比較してみましたが、ほぼ同じ結果でした（小数部7桁以降が若干異なるぐらいの差）
動作原理としては、通常のsentence_transformersのモデル読込処理の後、forwardメソッド部分でCT2モデルを読み込み＆利用するというハックをしているようです。
最初のモデル読込処理に変換前モデルファイル一式が必要なため、変換前モデルのパスを与えたのち、CT2モデルはそのサブフォルダct2を読み込むようにしています。
おまけ
langchainのHuggingfaceEmbeddingsの代わりにCT2変換Embeddingモデルを使う場合、以下のような感じのカスタムEmbeddingsを作成すればいけそうです。
コードの中身はHuggingfaceEmbeddingsクラスのコードを流用しました。
※ ざっと作っただけなので、テスト不十分です。
from typing import Any, Dict, List, Optional
from langchain.pydantic_v1 import BaseModel, Extra, Field
from langchain.schema.embeddings import Embeddings
class CT2Embeddings(BaseModel, Embeddings):
"""CTranslate2 and sentence_transformers embedding models.
To use, you should have the ``ctranslate2`` and ``sentence_transformers`` python package installed.
"""
client: Any #: :meta private:
model_name: str
"""Model name to use."""
cache_folder: Optional[str] = None
"""Path to store models.
Can be also set by SENTENCE_TRANSFORMERS_HOME environment variable."""
model_kwargs: Dict[str, Any] = Field(default_factory=dict)
"""Key word arguments to pass to the model."""
encode_kwargs: Dict[str, Any] = Field(default_factory=dict)
"""Key word arguments to pass when calling the `encode` method of the model."""
multi_process: bool = False
"""Run encode() on multiple GPUs."""
def __init__(self, **kwargs: Any):
"""Initialize the sentence_transformer."""
super().__init__(**kwargs)
try:
import sentence_transformers
import ctranslate2
except ImportError as exc:
raise ImportError(
"Could not import sentence_transformers or ctranslate2 python package. "
"Please install it with `pip install sentence-transformers ctranslate2`."
) from exc
self.client = CT2SentenceTransformer(
self.model_name, cache_folder=self.cache_folder, **self.model_kwargs
)
class Config:
"""Configuration for this pydantic object."""
extra = Extra.forbid
def embed_documents(self, texts: List[str]) -> List[List[float]]:
"""Compute doc embeddings using a HuggingFace transformer model.
Args:
texts: The list of texts to embed.
Returns:
List of embeddings, one for each text.
"""
import sentence_transformers
texts = list(map(lambda x: x.replace("\n", " "), texts))
if self.multi_process:
pool = self.client.start_multi_process_pool()
embeddings = self.client.encode_multi_process(texts, pool)
CT2SentenceTransformer.stop_multi_process_pool(pool)
else:
embeddings = self.client.encode(texts, **self.encode_kwargs)
return embeddings.tolist()
def embed_query(self, text: str) -> List[float]:
"""Compute query embeddings using a HuggingFace transformer model.
Args:
text: The text to embed.
Returns:
Embeddings for the text.
"""
return self.embed_documents([text])[0]
まとめ
CTranslate2を使ってEmbeddingする処理の紹介でした。
sentence-transformersの変更に大きく影響を受けそうな実装なので、現状試験利用ぐらいがいいのかもしれません。
動作速度についてですが、もともとEmbeddingsモデルは(GPU利用の場合)高速に動作するので、CTranslate2を使わなくても十分早いかなあという印象です。きちんと速度比較はしてないので、多量データのEmbedding処理を行う場合は結構違う可能性があります。
VRAM使用量については、量子化してないので返還前後でほぼ同じ量のメモリを使用していました。
量子化するとまた違うと思うのですが、量子化した場合の精度減少がEmbeddingの場合どれだけ発生するのかがわからずイマイチ踏み込めていません。
時間があるときに検証してみようと思います。
いろいろ注意点がありそうなのですが、CTranslate2が好きという方は試してみてください。
（そして情報をください）