Sentence Transformers 在 AWS Inferentia 上使用 Optimum Neuron
文本模型
本教程有一个笔记本版本，请点击此处查看。
本指南介绍如何使用 Optimum Neuron 在 AWS Inferentia2 上编译、加载和使用 Sentence Transformers (SBERT) 模型，从而实现高效的嵌入计算。 Sentence Transformers 是用于生成句子嵌入的强大模型。您可以使用 Sentence Transformers 计算 100 多种语言的句子/文本嵌入。然后可以使用余弦相似度等方法比较这些嵌入，以查找具有相似含义的句子。这对语义文本相似度、语义搜索或释义挖掘很有用。
将 Sentence Transformers 模型转换为 AWS Inferentia2
首先，您需要将 Sentence Transformers 模型转换为与 AWS Inferentia2 兼容的格式。您可以使用 optimum-cli 或 NeuronModelForSentenceTransformers 类使用 Optimum Neuron 编译 Sentence Transformers 模型。以下是两种方法的示例。我们必须确保 sentence-transformers 已安装。这仅用于导出模型。
pip install sentence-transformers
这里我们将使用 NeuronModelForSentenceTransformers，它可以用于将任何 Sntence Transformers 模型转换为与 AWS Inferentia2 兼容的格式，或者加载已转换的模型。使用 NeuronModelForSentenceTransformers 导出模型时，您需要设置 export=True 并定义输入形状和批次大小。输入形状由 sequence_length 定义，批次大小由 batch_size 定义。
from optimum.neuron import NeuronModelForSentenceTransformers
# Sentence Transformers model from HuggingFace
model_id = "BAAI/bge-small-en-v1.5"
input_shapes = {"batch_size": 1, "sequence_length": 384} # mandatory shapes
# Load Transformers model and export it to AWS Inferentia2
model = NeuronModelForSentenceTransformers.from_pretrained(model_id, export=True, **input_shapes)
# Save model to disk
model.save_pretrained("bge_emb_inf2/")这里我们将使用 optimum-cli 转换模型。与 NeuronModelForSentenceTransformers 类似，我们需要定义输入形状和批次大小。输入形状由 sequence_length 定义，批次大小由 batch_size 定义。optimum-cli 将自动将模型转换为与 AWS Inferentia2 兼容的格式，并将其保存到指定的输出目录。
optimum-cli export neuron -m BAAI/bge-small-en-v1.5 --sequence_length 384 --batch_size 1 --task feature-extraction bge_emb_inf2/加载编译后的 Sentence Transformers 模型并运行推断
一旦我们拥有一个编译后的 Sentence Transformers 模型，无论是我们自己导出的还是在 Hugging Face Hub 上可用的，我们都可以加载它并运行推断。为了加载模型，我们可以使用 NeuronModelForSentenceTransformers 类，它是 SentenceTransformer 类的抽象层。NeuronModelForSentenceTransformers 类将自动将输入填充到指定的 sequence_length 并使用 AWS Inferentia2 运行推断。
from optimum.neuron import NeuronModelForSentenceTransformers
from transformers import AutoTokenizer
model_id_or_path = "bge_emb_inf2/"
tokenizer_id = "BAAI/bge-small-en-v1.5"
# Load model and tokenizer
model = NeuronModelForSentenceTransformers.from_pretrained(model_id_or_path)
tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
# Run inference
prompt = "I like to eat apples"
encoded_input = tokenizer(prompt, return_tensors='pt')
outputs = model(**encoded_input)
token_embeddings = outputs.token_embeddings
sentence_embedding = outputs.sentence_embedding
print(f"token embeddings: {token_embeddings.shape}") # torch.Size([1, 7, 384])
print(f"sentence_embedding: {sentence_embedding.shape}") # torch.Size([1, 384])生产环境使用
要将这些模型部署到生产环境，请参考 Amazon SageMaker 博客.
CLIP
为 AWS Inferentia2 编译 CLIP
您可以使用 Optimum Neuron 编译 CLIP 模型，方法是使用 optimum-cli 或 NeuronModelForSentenceTransformers 类。选择您喜欢的其中一种方法。
- 使用 Optimum CLI
optimum-cli export neuron -m sentence-transformers/clip-ViT-B-32 --sequence_length 64 --text_batch_size 3 --image_batch_size 1 --num_channels 3 --height 224 --width 224 --task feature-extraction --subfolder 0_CLIPModel clip_emb/- 使用
NeuronModelForSentenceTransformers类
from optimum.neuron import NeuronModelForSentenceTransformers
model_id = "sentence-transformers/clip-ViT-B-32"
# configs for compiling model
input_shapes = {
"num_channels": 3,
"height": 224,
"width": 224,
"text_batch_size": 3,
"image_batch_size": 1,
"sequence_length": 64,
}
emb_model = NeuronModelForSentenceTransformers.from_pretrained(
model_id, subfolder="0_CLIPModel", export=True, library_name="sentence_transformers", dynamic_batch_size=False, **input_shapes
)
# Save locally or upload to the HuggingFace Hub
save_directory = "clip_emb/"
emb_model.save_pretrained(save_directory)加载编译后的 Sentence Transformers 模型并运行推断
from PIL import Image
from sentence_transformers import util
from transformers import CLIPProcessor
from optimum.neuron import NeuronModelForSentenceTransformers
save_directory = "clip_emb"
emb_model = NeuronModelForSentenceTransformers.from_pretrained(save_directory)
processor = CLIPProcessor.from_pretrained(save_directory)
inputs = processor(
text=["Two dogs in the snow", 'A cat on a table', 'A picture of London at night'], images=Image.open("two_dogs_in_snow.jpg"), return_tensors="pt", padding=True
)
outputs = emb_model(**inputs)
# Compute cosine similarities
cos_scores = util.cos_sim(outputs.image_embeds, outputs.text_embeds)
print(cos_scores)
# tensor([[0.3072, 0.1016, 0.1095]])警告
由于启用了动态批处理的编译模型只接受具有相同批次大小的输入张量，因此如果输入文本和图像具有不同的批次大小，我们就不能设置 dynamic_batch_size=True。并且由于 NeuronModelForSentenceTransformers 类将输入填充到编译期间使用的批次大小（text_batch_size 和 image_batch_size），您可以在编译期间使用相对较大的批次大小来提高灵活性，但代价是计算量会增加。
例如，如果您想要编码 3 或 4 或 5 个文本和 1 个图像，您可以在编译期间设置 text_batch_size = 5 = max(3, 4, 5) 和 image_batch_size = 1。