samples/py_custom_module/decode_secret_message.py

# Copyright 2022 The IREE Authors
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

import numpy as np
import os
import iree.compiler as compiler
import iree.runtime as rt

TOKEN_TABLE = [
    b"hi",
    b"thanks",
    b"bye",
    b"for",
    b"all",
    b"so",
    b"the",
    b"fish",
    b"the",
    b"end",
    b"long",
    b"and",
    b".",
    b"now",
    b"there",
]


def create_tokenizer_module():
  """Creates a module which defines some custom methods for decoding."""

  class Detokenizer:

    def __init__(self, iface):
      # Any class state here is maintained per-context.
      self.start_of_text = True
      self.start_of_sentence = True

    def reset(self):
      self.start_of_text = True
      self.start_of_sentence = True

    def accumtokens(self, ids_tensor_ref, token_list_ref):
      # TODO: This little dance to turn BufferView refs into real arrays... is not good.
      ids_bv = ids_tensor_ref.deref(rt.HalBufferView)
      ids_array = ids_bv.map().asarray(
          ids_bv.shape, rt.HalElementType.map_to_dtype(ids_bv.element_type))
      token_list = token_list_ref.deref(rt.VmVariantList)
      for index in range(ids_array.shape[0]):
        token_id = ids_array[index]
        token = TOKEN_TABLE[token_id]

        # And this dance to make a buffer... is also not good.
        # A real implementation would just map the constant memory, etc.
        buffer = rt.VmBuffer(len(token))
        buffer_view = memoryview(buffer)
        buffer_view[:] = token
        token_list.push_ref(buffer)
      return ids_array.shape[0]

    def jointokens(self, token_list_ref):
      # The world's dumbest detokenizer. Ideally, the state tracking
      # would be in a module private type that got retained and passed
      # back through.
      token_list = token_list_ref.deref(rt.VmVariantList)
      text = bytearray()
      for i in range(len(token_list)):
        item = bytes(token_list.get_as_object(i, rt.VmBuffer))
        if item == b".":
          text.extend(b".")
          self.start_of_sentence = True
        else:
          if not self.start_of_text:
            text.extend(b" ")
          else:
            self.start_of_text = False
          if self.start_of_sentence:
            text.extend(item[0:1].decode("utf-8").upper().encode("utf-8"))
            text.extend(item[1:])
            self.start_of_sentence = False
          else:
            text.extend(item)

      # TODO: This dance to make a buffer is still bad.
      results = rt.VmBuffer(len(text))
      memoryview(results)[:] = text
      return results.ref

  iface = rt.PyModuleInterface("detokenizer", Detokenizer)
  iface.export("accumtokens", "0rr_i", Detokenizer.accumtokens)
  iface.export("jointokens", "0r_r", Detokenizer.jointokens)
  iface.export("reset", "0v_v", Detokenizer.reset)
  return iface.create()


def compile():
  return compiler.tools.compile_file(os.path.join(os.path.dirname(__file__),
                                                  "main.mlir"),
                                     target_backends=["vmvx"])


def main():
  print("Compiling...")
  vmfb_contents = compile()
  print("Decoding secret message...")
  config = rt.Config("local-sync")
  main_module = rt.VmModule.from_flatbuffer(config.vm_instance, vmfb_contents)
  modules = config.default_vm_modules + (
      create_tokenizer_module(),
      main_module,
  )
  context = rt.SystemContext(vm_modules=modules, config=config)

  # First message.
  count = context.modules.main.add_tokens(
      np.asarray([5, 10, 11, 1, 3, 4, 5, 7, 12], dtype=np.int32))
  print(f"ADDED {count} tokens")

  # Second message.
  count = context.modules.main.add_tokens(np.asarray([2, 13], dtype=np.int32))
  print(f"ADDED {count} tokens")

  text = bytes(context.modules.main.get_results().deref(rt.VmBuffer))
  print(f"RESULTS: {text}")

  assert text == b"So long and thanks for all so fish. Bye now"

  # Reset and decode some more.
  context.modules.main.reset()
  count = context.modules.main.add_tokens(
      np.asarray([0, 14, 12], dtype=np.int32))
  print(f"ADDED {count} tokens")
  text = bytes(context.modules.main.get_results().deref(rt.VmBuffer))
  print(f"RESULTS: {text}")
  assert text == b"Hi there."


if __name__ == "__main__":
  main()