バイトダンスモデル大規模展開実戦

class Text2Ids:
def __init__(self) -> None:
self.table: Dict[str, int] = {
"hello": 0,
"world": 1,
"[UNK]": 2,
}

def lookup(self, word: str)
return self.table.get(word, 2)

def__call__ (self, words: List[str])
return [self.lookup(w) for w in words]

import matx

class WorkFlow:
def __init__(self):
# 此处会进行代码编译，Python 代码自动编译封装为 Callable 对象
self.text2ids = matx.script(Text2Ids)()

def process(self, texts):
ids = self.text2ids(texts)
return ids

# test
handler = WorkFlow()
print(handler.process("hello world unknown"))
# output: [0, 1, 2]

# dump
mod = matx.trace(handler.process, "hello world")
print(mod.run({"texts": "hello world"}))
mod.save('./my_dir')
# load
mod = matx.load('./my_dir', -1)
print(mod.run({"texts": "hello world"}))

#include <string>
#include <vector>
#include <map>
#include <iostream>
#include <matxscript/pipeline/tx_session.h>
using namespace ::matxscript::runtime;
int main()
{
// test case
std::unordered_map<std::string, RTValue> feed_dict;
feed_dict.emplace("texts", Unicode(U"hello world"));
std::vector<std::pair<std::string, RTValue>> result;
const char* module_path = "./my_dir";
const char* module_name = "model.spec.json";
{
// -1 mean cpu
auto sess = TXSession::Load(module_path, module_name, -1);
auto result = sess->Run(feed_dict);
for (auto& r : result) {
std::cout << "key: " << r.first << ", value: " << r.second << std::endl;
}
}
return 0;
}

from typing import List, Dict, Tuple
import libcut
import matx
class Vocabulary:
...
def utf8_decoder(s: List[bytes]):
return [x.decode() for x in s]
class TextNDArrayBuilder:
...
class TextPipeline:
def __init__(self, mode: str = "eval"):
self.mode = mode
self.cut_engine = libcut.Cutter('/path/to/cut_models', ...)
self.vocab = matx.script(Vocabulary)('/path/to/vocab.txt')
self.decoder = matx.script(utf8_decoder)
self.input_builder = matx.script(TextNDArrayBuilder)(self.vocab)
def process(self, text: List[bytes]):
# List[bytes] 是对齐 C++ 的 vector<string>
text: List[str] = self.decoder(text)
words: List[List[str]] = self.cut_engine(text)
batch_ids: List[List[int]] = self.vocab(words)
input_ids, segment_ids, mask_ids = self.input_builder(batch_ids, 32)
if self.mode == "train":
return input_ids.torch(), segment_ids.torch(), mask_ids.torch()
return input_ids, segment_ids, mask_ids

from typing import List, Dict, Tuple
import matx
from matx import vision
class VisionPipeline:
def __init__(self,
 device_id: int = 0,
 mode: str = "eval",
 image_size: int = 224,):
self.is_training = mode == 'train'
self.mode = mode
...
def process(self, image,):
if self.is_training:
decode_nds = self.random_crop_decode(image)
flip_nds = self.random_flip(decode_nds)
resize_nds = self.resize(flip_nds)
transpose_nd = self.transpose_norm(resize_nds, vision.SYNC)
else:
decode_nds = self.decode(image)
resize_nds = self.resize(decode_nds)
crop_nds = self.center_crop(resize_nds)
transpose_nd = self.transpose_norm(crop_nds, vision.SYNC)
if self.mode == "trace":
return transpose_nd
return transpose_nd.torch()

class MultimodalEvalPipeline:
def __init__(self):
self.text_pipe = TextPipeline(mode="eval", ...)
self.vision_pipe = VisionPipeline(mode="eval", ...)
self.torch_model = torch.jit.load('/path/to/multimodal.jit', map_locatinotallow='cuda:0')
self.tx_model_op = matx.script(self.torch_model, device=0)

def eval(self, texts: List[bytes], images: List[bytes])
input_ids, segment_ids, mask_ids = self.text_pipe.process(texts)
images = self.vision_pipe.process(images)
scores = self.tx_model_op(input_ids, segment_ids, mask_ids, images)
return scores
# examples
example_batch_size = 8
text_examples = ['hello, world'.encode()] * example_batch_size
with open('/path/image.jpg', 'rb') as f:
image_example = f.read()
image_examples = [image_example] * example_batch_size
# pipeline instance
pipe = MultimodalEvalPipeline(...)
mod = matx.trace(pipe.eval, text_examples, image_examples)
# test
print(mod.run({"texts": text_examples, "images": image_examples}))
# save
mod.save('/path/to/my_multimodal')

#include <string>
#include <vector>
#include <map>
#include <iostream>
#include <matxscript/pipeline/tx_session.h>
using namespace ::matxscript::runtime;
int main()
{
// test case
std::unordered_map<std::string, RTValue> feed_dict;
feed_dict.emplace("texts", List({String("hello world")}));
feed_dict.emplace("images", List({String("......")}));
std::vector<std::pair<std::string, RTValue>> result;
const char* module_path = "/path/to/my_multimodal";
const char* module_name = "model.spec.json";
{
// cuda:0
auto sess = TXSession::Load(module_path, module_name, 0);
auto result = sess->Run(feed_dict);
for (auto& r : result) {
std::cout << "key: " << r.first << ", value: " << r.second << std::endl;
}
}
return 0;
}