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):
        # 此處會(huì)進(jìn)行代碼編譯，Python 代碼自動(dòng)編譯封裝為 Callable 對(duì)象
        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] 是對(duì)齊 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;
}