对于特定用例，我必须将自定义模型部署到 OpenSearch 中，首先，我将 HuggingFace 模型导出到 TorchScript 中并已成功注册，然后我尝试部署它，但不幸的是部署操作失败并返回以下错误：
{
"model_id": "Guem7I4BM4PGcXAkFYKV",
"task_type": "DEPLOY_MODEL",
"function_name": "TEXT_EMBEDDING",
"state": "FAILED",
"worker_node": [
"d7zdcxaASDyFKGX7AhlG0g",
"5fc7S9yJQQCWarLbAnVESg",
"TPu-1KznRzCwj8tQAgHOMQ"
],
"create_time": 1713367889874,
"last_update_time": 1713368035088,
"error": """
{"d7zdcxaASDyFKGX7AhlG0g":"The following operation failed in the TorchScript interpreter.
\nTraceback of TorchScript, serialized code (most recent call last):
\n File \"code/__torch__.py\", line 12, in forward
\n input_ids = inputs[\"input_ids\"]
\n attention_mask = inputs[\"attention_mask\"]
\n input = inputs[\"token_type_ids\"]
\n ~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
\n _0 = (model).forward(input_ids, attention_mask, input, )
\n return {\"sentence_embedding\": _0}
\n
\nTraceback of TorchScript, original code (most recent call last):
\n/Users/Library/Python/3.9/lib/python/site-packages/torch/jit/_trace.py(1074): trace_module
\n/Users/Library/Python/3.9/lib/python/site-packages/torch/jit/_trace.py(806): trace
\n/Users/opensearch/CustomModel/TEST/import_torch.py(43): export_to_torchscript
\n/Users/opensearch/CustomModel/TEST/import_torch.py(52): <module>
\nRuntimeError: KeyError: token_type_ids
\n","5fc7S9yJQQCWarLbAnVESg":"The following operation failed in the TorchScript interpreter.
\nTraceback of TorchScript, serialized code (most recent call last):
\n File \"code/__torch__.py\", line 12, in forward
\n input_ids = inputs[\"input_ids\"]
\n attention_mask = inputs[\"attention_mask\"]
\n input = inputs[\"token_type_ids\"]
\n ~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
\n _0 = (model).forward(input_ids, attention_mask, input, )
\n return {\"sentence_embedding\": _0}
\n
\nTraceback of TorchScript, original code (most recent call last):
\n/Users/Library/Python/3.9/lib/python/site-packages/torch/jit/_trace.py(1074): trace_module
\n/Users/Library/Python/3.9/lib/python/site-packages/torch/jit/_trace.py(806): trace
\n/Users/opensearch/CustomModel/TEST/import_torch.py(43): export_to_torchscript
\n/Users/opensearch/CustomModel/TEST/import_torch.py(52): <module>
\nRuntimeError: KeyError: token_type_ids
\n","TPu-1KznRzCwj8tQAgHOMQ":"The following operation failed in the TorchScript interpreter.
\nTraceback of TorchScript, serialized code (most recent call last):
\n File \"code/__torch__.py\", line 12, in forward
\n input_ids = inputs[\"input_ids\"]
\n attention_mask = inputs[\"attention_mask\"]
\n input = inputs[\"token_type_ids\"]
\n ~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
\n _0 = (model).forward(input_ids, attention_mask, input, )
\n return {\"sentence_embedding\": _0}
\n
\nTraceback of TorchScript, original code (most recent call last):
\n/Users/Library/Python/3.9/lib/python/site-packages/torch/jit/_trace.py(1074): trace_module
\n/Users/Library/Python/3.9/lib/python/site-packages/torch/jit/_trace.py(806): trace
\n/Users/opensearch/CustomModel/TEST/import_torch.py(43): export_to_torchscript
\n/Users/opensearch/CustomModel/TEST/import_torch.py(52): <module>
\nRuntimeError: KeyError: token_type_ids\n"}
""",
"is_async": true
}
这是我用来将模型导出到 TorchScript 的 python 脚本
import torch
from transformers import AutoModel, AutoTokenizer, PreTrainedTokenizer
from transformers.utils import PaddingStrategy
from sentence_transformers import SentenceTransformer
class TorchScriptWrapper(torch.nn.Module):
def __init__(self, model):
super(TorchScriptWrapper, self).__init__()
self.model = model
def forward(self, inputs: dict):
with torch.no_grad():
outputs = self.model(inputs)
return {"sentence_embedding": outputs['sentence_embedding']}
def export_to_torchscript(model_name: str, is_sentence_transformer: bool, output_path: str, max_seq_length: int = 128):
tokenizer: PreTrainedTokenizer = AutoTokenizer.from_pretrained(model_name)
if is_sentence_transformer:
model = SentenceTransformer(model_name, device="cpu")
else:
model = AutoModel.from_pretrained(model_name)
model.eval()
# Define example text
text = "This is a test string"
# Create inputs
inputs = tokenizer(text, padding=PaddingStrategy.MAX_LENGTH, return_tensors="pt", max_length=max_seq_length)
# Instantiate the wrapper class
model_wrapper = TorchScriptWrapper(model)
# Unpack HF batch encoding into a regular dict
new_inputs = {}
new_inputs["input_ids"] = inputs["input_ids"]
new_inputs["attention_mask"] = inputs["attention_mask"]
if inputs.get("token_type_ids", None) is not None:
new_inputs["token_type_ids"] = inputs["token_type_ids"]
# Trace the wrapper class
traced_model = torch.jit.trace(model_wrapper, new_inputs, strict=False)
# Save traced model to file
traced_model.save(output_path)
if __name__ == "__main__":
# Load pre-trained model and tokenizer
model_name = "sentence-transformers/LaBSE"
export_to_torchscript(model_name, True, output_path="torchscript_labse.pt")
PS：即使 token_type_ids 键是输入字典的一部分，也会发生错误
问题已解决！
这是解决问题的更新脚本
from transformers import AutoModel, AutoTokenizer, PreTrainedTokenizer
from transformers.utils import PaddingStrategy
from sentence_transformers import SentenceTransformer
class TorchScriptWrapper(torch.nn.Module):
def __init__(self, model):
super(TorchScriptWrapper, self).__init__()
self.model = model
def forward(self, inputs: dict):
inputs['token_type_ids'] = torch.zeros_like(inputs['attention_mask'], dtype=torch.int)
with torch.no_grad():
outputs = self.model(inputs)
return {"sentence_embedding": outputs['sentence_embedding']}
def export_to_torchscript(model_name: str, is_sentence_transformer: bool, output_path: str, max_seq_length: int = 128):
tokenizer: PreTrainedTokenizer = AutoTokenizer.from_pretrained(model_name)
if is_sentence_transformer:
model = SentenceTransformer(model_name, device="cpu")
else:
model = AutoModel.from_pretrained(model_name)
model.eval()
# Define example text
text = "This is a test string"
# Create inputs
inputs = tokenizer(text, padding=PaddingStrategy.MAX_LENGTH, return_tensors="pt", max_length=max_seq_length)
# Instantiate the wrapper class
model_wrapper = TorchScriptWrapper(model)
# Unpack HF batch encoding into a regular dict
new_inputs = {}
new_inputs["input_ids"] = inputs["input_ids"]
new_inputs["attention_mask"] = inputs["attention_mask"]
#if inputs.get("token_type_ids", None) is not None:
# new_inputs["token_type_ids"] = inputs["token_type_ids"]
# Trace the wrapper class
traced_model = torch.jit.trace(model_wrapper, new_inputs, strict=False)
# Save traced model to file
traced_model.save(output_path)
if __name__ == "__main__":
# Load pre-trained model and tokenizer
model_name = "sentence-transformers/LaBSE"
export_to_torchscript(model_name, True, output_path="torchscript_labse.pt")