{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"Thanks for the public notebooks below:  \nhttps://www.kaggle.com/code/hoyso48/1st-place-solution-training  \nhttps://www.kaggle.com/code/irohith/aslfr-ctc-based-on-prev-comp-1st-place  \nhttps://www.kaggle.com/code/markwijkhuizen/aslfr-transformer-training-inference  \nThis is the 3rd place solution training code, you could refer the solution here:  \nhttps://www.kaggle.com/competitions/asl-fingerspelling/discussion/434393  ","metadata":{}},{"cell_type":"markdown","source":"# Install libs","metadata":{}},{"cell_type":"code","source":"try:\n  from icecream import ic\nexcept Exception:\n  !pip install -q icecream --no-index --find-links=file:///kaggle/input/icecream","metadata":{"execution":{"iopub.status.busy":"2023-08-27T08:04:10.271437Z","iopub.execute_input":"2023-08-27T08:04:10.271816Z","iopub.status.idle":"2023-08-27T08:04:25.060108Z","shell.execute_reply.started":"2023-08-27T08:04:10.271788Z","shell.execute_reply":"2023-08-27T08:04:25.058794Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Import libs","metadata":{}},{"cell_type":"code","source":"import sys, os\nimport numpy as np\nimport pandas as pd\nimport json\nimport re\nimport six\nimport glob\nimport traceback\nimport inspect\nfrom typing import Union\nfrom collections import Counter, OrderedDict, defaultdict\nfrom collections.abc import Iterable\nfrom multiprocessing import cpu_count\nfrom IPython.display import display\nfrom tqdm.notebook import tqdm\nfrom icecream import ic\nimport transformers\nimport tensorflow as tf\nimport torch\nfrom torch import nn, einsum\nimport torch.nn.functional as F\nic(tf.__version__, torch.__version__)","metadata":{"execution":{"iopub.status.busy":"2023-08-27T08:04:25.06231Z","iopub.execute_input":"2023-08-27T08:04:25.062663Z","iopub.status.idle":"2023-08-27T08:04:37.423591Z","shell.execute_reply.started":"2023-08-27T08:04:25.062628Z","shell.execute_reply":"2023-08-27T08:04:37.422245Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Force tf not use gpu memory, for only use tf for tfrecord reading,preprocessing and postprocessing, not for train(using torch)","metadata":{}},{"cell_type":"code","source":"tf.config.set_visible_devices([], 'GPU')","metadata":{"execution":{"iopub.status.busy":"2023-08-27T08:04:37.425403Z","iopub.execute_input":"2023-08-27T08:04:37.4263Z","iopub.status.idle":"2023-08-27T08:04:37.439226Z","shell.execute_reply.started":"2023-08-27T08:04:37.426255Z","shell.execute_reply":"2023-08-27T08:04:37.438045Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Flags","metadata":{}},{"cell_type":"code","source":"class FLAGS(object):\n  # for tfrecords args, you could ignore\n  seed = 1024\n  batch_parse = False\n  sparse_to_dense = True\n  eval_keys = []\n  incl_keys = []\n  excl_keys = []\n  recount_tfrecords = False  \n  batch_sizes = []\n  buffer_size = 1024\n  buckets = None\n  drop_remainder = None\n  shard_by_files = True\n  shuffle_batch = None\n  shuffle_files = None\n  num_dataset_threads = 0\n  num_prefetch_batches = 1024\n  repeat_then_shuffle = False\n  length_index = 1\n  length_key = None\n  dynamic_pad = True\n  cache = False\n  cache_after_map = False\n  fixed_random = False\n  parallel_read_files = True\n  padding_idx = 0\n  dataset_keys = []\n  dataset_excl_keys = []\n  exclude_varlen_keys = False\n  prefetch = None\n  dataset_ordered = False\n    \n  torch = True\n  keras = False\n    \n  # online==False means using n-fold split and train on fold 1,2, folds-1 while valid on fold 0\n  # online==True means using all train data but still will valid on fold 0\n  online = False  \n  folds = 4\n  fold = 0\n  fold_seed = 1229\n  root = '../input/asl-fingerspelling'\n  working = '/kaggle/working'\n  use_z = True  # use x,y,z if True\n  norm_frames = True # norm frames using x - mean / std\n  concat_frames = True # concat original and normalized frames\n  add_pos = True # add abs frame pos, like 1/1000., 2/1000.\n  sup_weight = 0.1 # for supplement dataset assigin weight 0.1\n  \n  train_files = []\n  valid_files = []\n      \n  mix_sup = True # train & sup dataset\n  vie = 5 # valid interval epochs \n  lr = 2e-3\n  epochs = 400 \n  batch_size = 128\n  eval_batch_size = 256\n  awp = True\n  adv_start_epoch = None\n  adv_lr = 0.2\n  adv_eps = 0\n  fp16 = False # notice fp16 could not be set True if using awp here, otherwise nan\n  optimizer = 'Adam'\n  opt_eps = 1e-6 \n  scheduler = 'linear'\n  # for model related configs\n  encoder_layers = 17\n  encoder_units = 200 \n  n_frames = 320  \n  distributed = False\n\n  # for preprocess\n  trunct_method = 'resize'\n\n  # for aug\n  flip_rate = 0.25\n  resample_rate = 0.8\n  temporal_mask_rate = 0.8\n  temporal_mask_prob = 0.5\n  temporal_mask_range = [0.1, 0.5]\n  spatio_mask_rate = 0.\n  spatio_mask_prob = 0.\n  add_pos_before_resample = True\n  shift_rate = 0.\n  shift_range = [-0.05, 0.05]\n  shift_method = 1\n  temporal_seq_mask_rate = 0.5\n  temporal_seq_mask_max = 2\n  temporal_seq_mask_range = [0.1, 0.2]\n  shift_rate = 0.75\n  scale_method = 0\n  scale_rate = 0.75\n  rotate_rate = 0.75\n  \n  # for model\n  ksize_vals = [15]\n  skip_factor = 0.5\n  mhatt_heads = 8\n  mhatt_dimhead = 32\n\nFLAGS.adv_start_epoch = int(FLAGS.epochs * 0.15)\n    \ndef load_json(filename):\n  with open(filename) as fh:\n    obj = json.load(fh)\n  return obj","metadata":{"execution":{"iopub.status.busy":"2023-08-27T08:04:37.443034Z","iopub.execute_input":"2023-08-27T08:04:37.443466Z","iopub.status.idle":"2023-08-27T08:04:37.460607Z","shell.execute_reply.started":"2023-08-27T08:04:37.443428Z","shell.execute_reply":"2023-08-27T08:04:37.459249Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"FLAGS.batch_size = 128\nFLAGS.acc_steps = 4 # gradient acc\nFLAGS.batch_size = FLAGS.batch_size // FLAGS.acc_steps","metadata":{"execution":{"iopub.status.busy":"2023-08-27T08:04:37.462604Z","iopub.execute_input":"2023-08-27T08:04:37.463036Z","iopub.status.idle":"2023-08-27T08:04:37.474786Z","shell.execute_reply.started":"2023-08-27T08:04:37.462996Z","shell.execute_reply":"2023-08-27T08:04:37.473717Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Common configs","metadata":{}},{"cell_type":"code","source":"LPOSE = [13, 15, 17, 19, 21]\nRPOSE = [14, 16, 18, 20, 22]\nPOSE = LPOSE + RPOSE\n\nLIP = [\n    61, 185, 40, 39, 37, 0, 267, 269, 270, 409,\n    291, 146, 91, 181, 84, 17, 314, 405, 321, 375,\n    78, 191, 80, 81, 82, 13, 312, 311, 310, 415,\n    95, 88, 178, 87, 14, 317, 402, 318, 324, 308,\n]\nic(len(LIP))\nLLIP = [84,181,91,146,61,185,40,39,37,87,178,88,95,78,191,80,81,82]\nRLIP = [314,405,321,375,291,409,270,269,267,317,402,318,324,308,415,310,311,312]\nMID_LIP = [i for i in LIP if i not in LLIP + RLIP]\nic(len(LLIP), len(RLIP), len(MID_LIP))\n\nNOSE=[\n    1,2,98,327\n]\nLNOSE = [98]\nRNOSE = [327]\nMID_NOSE = [i for i in NOSE if i not in LNOSE + RNOSE]\n\nLEYE = [\n    263, 249, 390, 373, 374, 380, 381, 382, 362,\n    466, 388, 387, 386, 385, 384, 398,\n]\nREYE = [\n    33, 7, 163, 144, 145, 153, 154, 155, 133,\n    246, 161, 160, 159, 158, 157, 173,\n]\n\nN_HAND_POINTS = 21\nN_POSE_POINTS = len(LPOSE)\nN_LIP_POINTS = len(LLIP)\nN_EYE_POINTS = len(LEYE)\nN_NOSE_POINTS = len(LNOSE)\nN_MID_POINTS = len(MID_LIP + MID_NOSE)\n\nSEL_COLS = []\nfor i in range(N_HAND_POINTS):\n  SEL_COLS.extend([f'x_left_hand_{i}', f'y_left_hand_{i}', f'z_left_hand_{i}'])\nfor i in range(N_HAND_POINTS):\n  SEL_COLS.extend([f'x_right_hand_{i}', f'y_right_hand_{i}', f'z_right_hand_{i}'])\nfor i in LPOSE:\n  SEL_COLS.extend([f'x_pose_{i}', f'y_pose_{i}', f'z_pose_{i}'])\nfor i in RPOSE:\n  SEL_COLS.extend([f'x_pose_{i}', f'y_pose_{i}', f'z_pose_{i}'])\nfor i in LLIP:\n  SEL_COLS.extend([f'x_face_{i}', f'y_face_{i}', f'z_face_{i}'])\nfor i in RLIP:\n  SEL_COLS.extend([f'x_face_{i}', f'y_face_{i}', f'z_face_{i}'])\n\nfor i in LEYE:\n  SEL_COLS.extend([f'x_face_{i}', f'y_face_{i}', f'z_face_{i}'])\nfor i in REYE:\n  SEL_COLS.extend([f'x_face_{i}', f'y_face_{i}', f'z_face_{i}'])\n  \nfor i in LNOSE:\n  SEL_COLS.extend([f'x_face_{i}', f'y_face_{i}', f'z_face_{i}'])\nfor i in RNOSE:\n  SEL_COLS.extend([f'x_face_{i}', f'y_face_{i}', f'z_face_{i}'])\n  \nfor i in MID_LIP:\n  SEL_COLS.extend([f'x_face_{i}', f'y_face_{i}', f'z_face_{i}'])\nfor i in MID_NOSE:\n  SEL_COLS.extend([f'x_face_{i}', f'y_face_{i}', f'z_face_{i}'])\n    \nN_COLS = len(SEL_COLS)\nic(N_COLS)\n    \nCHAR2IDX = load_json(f'../input/asl-fingerspelling/character_to_prediction_index.json')\nCHAR2IDX = {k: v + 1 for k, v in CHAR2IDX.items()}\nN_CHARS = len(CHAR2IDX)\nic(N_CHARS)\n\nPAD_IDX = 0\nSOS_IDX = PAD_IDX # Start Of Sentence\nEOS_IDX = N_CHARS + 1 # End Of Sentence\nic(PAD_IDX, SOS_IDX, EOS_IDX)\n\nPAD_TOKEN = '<PAD>'\nSOS_TOKEN = PAD_TOKEN\nEOS_TOKEN = '<EOS>'\n\nCHAR2IDX[PAD_TOKEN] = PAD_IDX\nCHAR2IDX[EOS_TOKEN] = EOS_IDX \n\nADDRESS_TOKEN = '<ADDRESS>'\nURL_TOKEN = '<URL>'\nPHONE_TOKEN = '<PHONE>'\nSUP_TOKEN = '<SUP>'\n\nVOCAB_SIZE = len(CHAR2IDX)\nIDX2CHAR = {v: k for k, v in CHAR2IDX.items()}\nic(VOCAB_SIZE)\nic(len(IDX2CHAR))\n\nSTATS = {}\nCLASSES = [\n  'address', \n  'url', \n  'phone', \n  'sup',\n  ]\nPHRASE_TYPES = dict(zip(CLASSES, range(len(CLASSES))))\nN_TYPES = len(CLASSES)\nMAX_PHRASE_LEN = 32\n\ndef get_phrase_type(phrase):\n  # Phone Number\n  if re.match(r'^[\\d+-]+$', phrase):\n    return 'phone'\n  # url\n  elif any([substr in phrase for substr in ['www', '.', '/']\n           ]) and ' ' not in phrase:\n    return 'url'\n  # Address\n  else:\n    return 'address'\n\ndef get_vocab_size():\n  vocab_size = VOCAB_SIZE\n  return vocab_size\n\ndef get_n_cols(no_motion=False, use_z=None):\n  n_cols = N_COLS\n  if use_z is None:\n    use_z = FLAGS.use_z\n  \n  if FLAGS.concat_frames:\n    assert FLAGS.norm_frames\n    n_cols += N_COLS\n  \n  if not use_z:\n    n_cols = n_cols // 3 * 2\n    \n  if FLAGS.add_pos:\n    n_cols += 1\n  \n  return n_cols","metadata":{"execution":{"iopub.status.busy":"2023-08-27T08:04:37.477037Z","iopub.execute_input":"2023-08-27T08:04:37.47752Z","iopub.status.idle":"2023-08-27T08:04:37.948382Z","shell.execute_reply.started":"2023-08-27T08:04:37.477473Z","shell.execute_reply":"2023-08-27T08:04:37.947319Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\"\"\"\n    Tensorflow layer to process data in TFLite\n    Data needs to be processed in the model itself, so we can not use Python\n    [None, N_COLS] -> [1, NONE, N_COLS] -> [1, N_FRAMES, N_COLS]\n\"\"\"\nclass PreprocessLayer(tf.keras.layers.Layer):\n\n  def __init__(self, n_frames=128, training=False):\n    super(PreprocessLayer, self).__init__()\n    self.n_frames = n_frames\n    self.n_cols = get_n_cols(no_motion=True, use_z=True)\n    ic(self.n_cols)\n    self.means = np.load(f'../input/3rd-place-step-3-gen-mean-and-std/means.npy')\n    self.stds = np.load(f'../input/3rd-place-step-3-gen-mean-and-std/stds.npy')\n    ic(self.means.shape, self.stds.shape)\n    self.training = training\n\n  ## seems not needed as training perf is similar\n  @tf.function(input_signature=(tf.TensorSpec(shape=[None, N_COLS],\n                                              dtype=tf.float32),),)\n  ## this will cause error\n  # @tf.function()\n  def call(self, data):\n    trunct_method = FLAGS.trunct_method\n    training = self.training\n\n    # Hacky (add dim in front line shape [None,N_COLS] to shape [1, None, N_COLS])\n    data = data[None]\n    dtype = data.dtype\n    \n    # [1, None, N_COLS//3, 3]\n    data = reshape(data)\n    \n    if training and FLAGS.use_aug:\n      data = Apply(scale, FLAGS.scale_rate)(data)\n      data = Apply(rotate, FLAGS.rotate_rate)(data)\n      # shift not affect much..\n      data = Apply(shift, FLAGS.shift_rate)(data)\n      data = OneOf([scale, rotate, shift], FLAGS.affine_rate)(data)\n      data = Apply(temporal_seq_mask, FLAGS.temporal_seq_mask_rate)(data)\n    \n    data_ = data\n    data = reshape_back(data)\n      \n    # seems norm before resize produce better results\n    # -------norm frames\n    l = []\n    if FLAGS.norm_frames:\n      if FLAGS.concat_frames:\n        l.append(data_)\n      l.append(\n          (data - self.means) / self.stds,\n      )\n      l[-1] = reshape(l[-1])\n    else:\n      l.append(data_)\n            \n    # [1, None, feats, 3]\n    data = tf.concat(l, axis=-2)\n\n    n_cols = self.n_cols if not FLAGS.add_pos else self.n_cols - 1\n    data = tf.reshape(data, [1, -1, n_cols])\n    # Fill NaN Values With 0\n    data = tf.where(tf.math.is_nan(data), tf.zeros_like(data), data)\n                       \n    # add_pos helps a lot and add_pos_before_resample also +0.002 compare to add_pos_after_resample\n    if FLAGS.add_pos and FLAGS.add_pos_before_resample:\n      # n_frames = len(data[0])\n      n_frames = tf.shape(data)[1]\n      pos = get_pos(n_frames, dtype)\n      data = tf.concat([data, pos], axis=-1)\n      \n    #  now basic features done, we resample for frames with basic features\n    if training and FLAGS.use_aug:\n      # resample for time range, like 120 frames to 135 or 90 frames\n      # resample helps a lot\n      data = Apply(resample, FLAGS.resample_rate)(data)\n      data = tf.cast(data, dtype)\n      \n    # ------add addional info like add pos after resample, actually similar results before or after resample\n    # n_frames = len(data[0])\n    n_frames = tf.shape(data)[1]\n    \n    # do it after expand dim 0, and do not use expand_dims before...\n    if FLAGS.add_pos and (not FLAGS.add_pos_before_resample):\n      pos = get_pos(n_frames, dtype)\n      data = tf.concat([data, pos], axis=-1)\n\n    if not FLAGS.always_resize:\n      # Pad Zeros  TODO dynamic type so not pad? dynanmic seq2seq input\n      # n_frames = len(data[0])\n      n_frames = tf.shape(data)[1]\n      if n_frames < self.n_frames:\n        if FLAGS.pad_frames:\n          if FLAGS.pad_method == 'zero':\n            data = tf.concat((data,\n                              tf.zeros([1, self.n_frames - n_frames, self.n_cols],\n                                        dtype=dtype)),\n                            axis=1)\n          else:\n            data = tf.image.resize(\n              data,\n              [1, self.n_frames],\n              method=FLAGS.pad_resize_method,\n            )\n        elif n_frames < FLAGS.encode_pool_size:\n          data = tf.concat(\n              (data,\n               tf.zeros([1, FLAGS.encode_pool_size - n_frames, self.n_cols],\n                        dtype=dtype)),\n              axis=1)\n        elif n_frames < FLAGS.min_frames:\n          data = tf.concat(\n              (data,\n               tf.zeros([1, FLAGS.min_frames - n_frames, self.n_cols],\n                        dtype=dtype)),\n              axis=1)\n\n      # n_frames = len(data[0])\n      n_frames = tf.shape(data)[1]\n      if n_frames > self.n_frames:\n        if trunct_method == 'resize':\n          # Downsample\n          data = tf.image.resize(\n              data,\n              [1, self.n_frames],\n              method=FLAGS.resize_method,\n          )\n          # For resize The return value has type float32, unless the method is ResizeMethod.NEAREST_NEIGHBOR\n          data = tf.cast(data, dtype)\n        else:\n          data = data[:, :self.n_frames, :]\n    else:\n      data = tf.image.resize(\n          data,\n          [1, self.n_frames],\n          method=FLAGS.resize_method,\n      )\n      data = tf.cast(data, dtype)\n\n    if FLAGS.pad_frames:\n      data = tf.reshape(data, [1, self.n_frames, self.n_cols])\n      \n    # 这个mask是否放在最后最好 是否放在前面 然后filter 全0的帧？那样比较麻烦 感觉这样还行...\n    if training and FLAGS.use_aug:\n      # temperal mask help a lot\n      data = Apply(temporal_mask, FLAGS.temporal_mask_rate)(data)\n      data = Apply(spatio_mask, FLAGS.spatio_mask_rate)(data)\n      \n    if not FLAGS.use_z:\n      if FLAGS.add_pos:\n        pos_data = data[...,-2:-1]\n        data = data[...,:-1]\n      data = tf.reshape(data, [1, self.n_frames, -1, 3])\n      data = data[...,:2]\n      data = tf.reshape(data, [1, self.n_frames, -1])\n      if FLAGS.add_pos:\n        data = tf.concat([data, pos_data], axis=-1)\n             \n    # Squeeze Batch Dimension\n    data = tf.squeeze(data, axis=[0])\n    # tf.print('----', data.shape)\n    return data\n\nclass PreProcssor(object):\n  def __init__(self, subset='valid', squeeze=False):\n    training = subset == 'train'\n    self.prepocess = PreprocessLayer(FLAGS.n_frames, training=training)\n    self.squeeze = squeeze\n\n  def __call__(self, fe):\n    # # TODO HACK for hug datasets.. ds = ds.to_tf_dataset(batch_size=1)\n    if self.squeeze:\n      for key in fe:\n        fe[key] = tf.squeeze(fe[key], axis=0)\n    x = fe\n    weights = fe['weight']\n    fe['frames'] = tf.reshape(fe['frames'], [fe['n_frames'], N_COLS])\n    fe['frames'] = self.prepocess(fe['frames']) \n    fe['phrase_len'] = tf.cond(tf.greater(fe['phrase_len'], MAX_PHRASE_LEN), lambda: tf.constant(MAX_PHRASE_LEN, fe['phrase_len'].dtype), lambda: fe['phrase_len'])\n    if FLAGS.no_eos:\n      # as in tfrecords we have eos, but in training if we don't have eos we change it to pad\n      mask = tf.logical_or(fe['phrase_'] == PAD_IDX, fe['phrase_'] == EOS_IDX)\n      mask = tf.cast(mask, fe['phrase_'].dtype)\n      fe['phrase_'] = fe['phrase_'] * (1 - mask) + mask * PAD_IDX\n        \n    if FLAGS.mix_sup:\n      weights_mask = tf.cast(weights != 1, tf.float32)\n      weights = weights_mask * FLAGS.sup_weight + (1 - weights_mask)\n      x['weight'] = weights\n    \n    y = fe['phrase_']\n    return x, y\n  ","metadata":{"execution":{"iopub.status.busy":"2023-08-27T08:05:01.859581Z","iopub.execute_input":"2023-08-27T08:05:01.860017Z","iopub.status.idle":"2023-08-27T08:05:01.898898Z","shell.execute_reply.started":"2023-08-27T08:05:01.859985Z","shell.execute_reply":"2023-08-27T08:05:01.897664Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"FLAGS.use_aug = True\nFLAGS.always_resize = False\nFLAGS.affine_rate = 0\nFLAGS.force_flip = False\nFLAGS.scale_rate = 0.75\nFLAGS.scale_range = [0.8, 1.2]\nFLAGS.rotate_range = [-15.0, 15.0]\nFLAGS.shift_range = [-0.05, 0.05]\nFLAGS.temporal_seq_mask_range = [0.1, 0.2]\nFLAGS.pad_frames = True\nFLAGS.pad_method = 'zero'\nFLAGS.resize_method = 'bilinear'\nFLAGS.no_eos = False","metadata":{"execution":{"iopub.status.busy":"2023-08-27T08:05:02.243789Z","iopub.execute_input":"2023-08-27T08:05:02.244233Z","iopub.status.idle":"2023-08-27T08:05:02.251284Z","shell.execute_reply.started":"2023-08-27T08:05:02.244196Z","shell.execute_reply":"2023-08-27T08:05:02.250059Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from torch.utils.data import IterableDataset\n\nclass TorchDataset(IterableDataset):\n  def __init__(self, subset='eval'):\n    self.subset = subset\n\n    if subset == 'train':\n      dataset = TfDataset('train', files=FLAGS.train_files)\n      datas = dataset.make_batch(FLAGS.batch_size, \n                             shuffle=True,\n                             repeat=True,\n                             drop_remainder=True,\n                             return_numpy=True)\n    else:\n      dataset = TfDataset('valid', files=FLAGS.valid_files)\n      datas = dataset.make_batch(FLAGS.eval_batch_size,  \n                             shuffle=False,\n                             repeat=False, \n                             drop_remainder=False,\n                             return_numpy=False) # if set return numpy = True then you could only visit 1 time..\n    \n    self.num_instances = dataset.num_instances\n    ic(self.num_instances)\n    self.num_steps = dataset.num_steps\n    ic(self.num_steps)\n    \n    self.datas = datas\n    self.data_iter = iter(datas)\n         \n  def __iter__(self):\n    if self.subset == 'train':\n      while True:\n        x, y = next(self.data_iter)\n        del x['phrase_type']\n        del x['phrase']\n        if FLAGS.distributed:\n          rank = gezi.get('RANK', 0)\n          start = rank * FLAGS.batch_size\n          end = start + FLAGS.batch_size\n          for key in x:\n            x[key] = x[key][start:end]\n          y = y[start:end]\n        yield x, y\n    else:\n      for x, y in self.datas:\n        for key in x:\n          x[key] = x[key].numpy()\n        y = y.numpy()\n        yield x, y\n    \n  def __len__(self):\n    return self.num_instances\n  \ndef get_dataloaders():\n  kwargs = {\n      'num_workers': 0, # >0 will hang...\n      'pin_memory': True,\n      'persistent_workers': False,\n  }\n  \n  train_ds = TorchDataset('train')\n  eval_ds = TorchDataset('eval')\n  \n  train_dl = torch.utils.data.DataLoader(train_ds,\n                                         batch_size=None,\n                                         sampler=None,\n                                         **kwargs)\n  eval_dl = torch.utils.data.DataLoader(eval_ds,\n                                        batch_size=None,\n                                        sampler=None,\n                                        **kwargs)\n\n  return train_dl, eval_dl  \n","metadata":{"execution":{"iopub.status.busy":"2023-08-27T08:05:02.947344Z","iopub.execute_input":"2023-08-27T08:05:02.947785Z","iopub.status.idle":"2023-08-27T08:05:02.971308Z","shell.execute_reply.started":"2023-08-27T08:05:02.947751Z","shell.execute_reply":"2023-08-27T08:05:02.969648Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model(17 layers squeezeformer)","metadata":{}},{"cell_type":"code","source":"FLAGS.encode_pool_size = 2\nFLAGS.emb_batchnorm = True\nFLAGS.attn_drop = 0\nFLAGS.ff_drop = 0\nFLAGS.conv_drop = 0\nFLAGS.inst_drop_rate = 0.2\nFLAGS.cls_drop = 0.1\nFLAGS.conv1d_ksize_vals = [15]\nFLAGS.conv1d_expansion_factor = 2\nFLAGS.ff_mult = 4\nFLAGS.encoder_units = 200 \nFLAGS.encoder_layers = 17\nFLAGS.mhatt_dimhead = 32\nFLAGS.mhatt_heads = 8\nFLAGS.skip_factor = 0.5\n# time reduce from layer 8, means 7 + 10 = 17 layers while the first 7 layers using 320 frames and the last 10 layers using 160 frames\nFLAGS.time_reduce = True\nFLAGS.time_reduce_idx = 7\nFLAGS.time_kernel_size = 5 \nFLAGS.time_stride = 2\n# for ROPE\nFLAGS.scaling_type = 'dynamic'\nFLAGS.scaling_factor = 0.5","metadata":{"execution":{"iopub.status.busy":"2023-08-27T08:05:03.845404Z","iopub.execute_input":"2023-08-27T08:05:03.845839Z","iopub.status.idle":"2023-08-27T08:05:03.854065Z","shell.execute_reply.started":"2023-08-27T08:05:03.8458Z","shell.execute_reply":"2023-08-27T08:05:03.852812Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def reshape(data):\n  data = tf.reshape(data, (1, -1, N_COLS // 3, 3))\n  return data\n\ndef reshape_back(data):\n  data = tf.reshape(data, (1, -1, N_COLS))\n  return data\n\ndef get_pos(n_frames, dtype):\n  pos = tf.range(n_frames, dtype=dtype)\n  pos /= 1000.\n  pos = tf.reshape(pos, [1, n_frames, 1])\n  return pos\n\ndef simple_decode(pred):\n  x = tf.argmax(pred, axis=-1)\n  return x\n\ndef adjust_pad(pred):\n  pred = tf.nn.softmax(pred, axis=-1)\n  pred0, pred1 = pred[..., 0:1], pred[..., 1:]\n  pred0 *= tf.cast(pred0 > FLAGS.pad_thre, pred.dtype)\n  pred = tf.concat([pred0, pred1], axis=-1)\n  return pred\n\ndef ctc_decode(x):\n  x = tf.concat([x, tf.zeros((1), dtype=x.dtype)], axis=0)\n  diff = tf.not_equal(x[:-1], x[1:])\n  adjacent_indices = tf.where(diff)[:, 0]\n  x = tf.gather(x, adjacent_indices)\n  mask = x != PAD_IDX\n  x = tf.boolean_mask(x, mask, axis=0)\n  return x\n\n@tf.function()\ndef decode_phrase(pred):\n  pred = adjust_pad(pred)  \n  x = tf.argmax(pred, axis=-1)\n  x = ctc_decode(x)\n  return x  ","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class BatchNorm(nn.Module):\n  def __init__(self, num_features, momentum=0.1, eps=1e-5):\n    super().__init__()\n    self.bn = nn.BatchNorm1d(num_features, momentum=momentum, eps=eps)\n  \n  def forward(self, x):\n    x = x.permute(0, 2, 1)\n    x = self.bn(x)\n    x = x.permute(0, 2, 1)\n    return x\n\nclass InstanceDropout(nn.Module):\n  def __init__(self, p=0.5):\n    super().__init__()\n    self.p = p\n    self.dropout = nn.Dropout(p)\n\n  def forward(self, x):\n    mask = torch.ones_like(x[:,:1,:1])\n    mask = self.dropout(mask)\n    return x * mask","metadata":{"execution":{"iopub.status.busy":"2023-08-27T08:05:04.324329Z","iopub.execute_input":"2023-08-27T08:05:04.324713Z","iopub.status.idle":"2023-08-27T08:05:04.336064Z","shell.execute_reply.started":"2023-08-27T08:05:04.324683Z","shell.execute_reply":"2023-08-27T08:05:04.334885Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class SimpleEmbedding(nn.Module):\n\n  def __init__(self):\n    super().__init__()\n\n    self.emb_batchnorm = True if FLAGS.emb_batchnorm is None else FLAGS.emb_batchnorm\n    self.embedding = nn.Linear(get_n_cols(), FLAGS.encoder_units, bias=False)\n    if self.emb_batchnorm:\n      self.batch_norm = BatchNorm(FLAGS.encoder_units, momentum=0.05, eps=1e-3)\n\n  def forward(self, x):\n    x = self.embedding(x)\n    if self.emb_batchnorm:\n      x = self.batch_norm(x)\n    return x","metadata":{"execution":{"iopub.status.busy":"2023-08-27T08:05:04.960393Z","iopub.execute_input":"2023-08-27T08:05:04.960778Z","iopub.status.idle":"2023-08-27T08:05:04.968609Z","shell.execute_reply.started":"2023-08-27T08:05:04.960738Z","shell.execute_reply":"2023-08-27T08:05:04.967354Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class SqueezeformerEncoder(nn.Module):\n\n  def __init__(self):\n    super().__init__()\n    self.embedding = SimpleEmbedding()\n    attn_dropout, ff_dropout, conv_dropout = FLAGS.attn_drop, FLAGS.ff_drop, FLAGS.conv_drop\n    self.encoder = Squeezeformer(\n        dim=FLAGS.encoder_units,\n        depth=FLAGS.encoder_layers,\n        dim_head=FLAGS.mhatt_dimhead,\n        heads=FLAGS.mhatt_heads,\n        ff_mult=FLAGS.ff_mult,\n        conv_expansion_factor=FLAGS.conv1d_expansion_factor,  # 2\n        conv_kernel_size=FLAGS.conv1d_ksize_vals[0],\n        attn_dropout=attn_dropout,\n        ff_dropout=ff_dropout,\n        conv_dropout=conv_dropout)\n\n  def forward(self, x):\n    x = self.embedding(x)\n    x = self.encoder(x)\n    return x","metadata":{"execution":{"iopub.status.busy":"2023-08-27T08:05:05.398453Z","iopub.execute_input":"2023-08-27T08:05:05.399558Z","iopub.status.idle":"2023-08-27T08:05:05.40743Z","shell.execute_reply.started":"2023-08-27T08:05:05.399519Z","shell.execute_reply":"2023-08-27T08:05:05.406245Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class AvgPoolingModule(nn.Module):\n  def __init__(self, pool_size, *args, **kwargs) -> None:\n    super().__init__(*args, **kwargs)\n    self.pooling = nn.AvgPool1d(pool_size)\n    \n  def forward(self, x):\n    x = x.permute(0, 2, 1)\n    x = self.pooling(x)\n    x = x.permute(0, 2, 1)\n    return x","metadata":{"execution":{"iopub.status.busy":"2023-08-27T08:05:05.922187Z","iopub.execute_input":"2023-08-27T08:05:05.923338Z","iopub.status.idle":"2023-08-27T08:05:05.930271Z","shell.execute_reply.started":"2023-08-27T08:05:05.923286Z","shell.execute_reply":"2023-08-27T08:05:05.929132Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Encoder(nn.Module):\n\n  def __init__(self):\n    super().__init__()\n    self.encoder = SqueezeformerEncoder()\n    if FLAGS.encode_pool_size > 1:\n      self.pooling = AvgPoolingModule(FLAGS.encode_pool_size)\n\n  def forward(self, frames):\n    x = self.encoder(frames)  \n    if FLAGS.encode_pool_size > 1:\n      x = self.pooling(x)\n    return x","metadata":{"execution":{"iopub.status.busy":"2023-08-27T08:05:06.362163Z","iopub.execute_input":"2023-08-27T08:05:06.362599Z","iopub.status.idle":"2023-08-27T08:05:06.369966Z","shell.execute_reply.started":"2023-08-27T08:05:06.362568Z","shell.execute_reply":"2023-08-27T08:05:06.368905Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class InferModel(nn.Module):\n  def __init__(self, model, **kwargs):\n    super().__init__(**kwargs)\n    self.model = model\n  \n  def forward(self, frames):\n    res = self.model.infer(frames)\n    return res\n  \n  def infer(self, frames):\n    return self.forward(frames)","metadata":{"execution":{"iopub.status.busy":"2023-08-27T08:05:06.861258Z","iopub.execute_input":"2023-08-27T08:05:06.861647Z","iopub.status.idle":"2023-08-27T08:05:06.86888Z","shell.execute_reply.started":"2023-08-27T08:05:06.861618Z","shell.execute_reply":"2023-08-27T08:05:06.867431Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Model(nn.Module):\n\n  def __init__(self, **kwargs):\n    super().__init__(**kwargs)\n    self.encoder = Encoder()\n    self.classifer = nn.Sequential(\n            nn.Dropout(FLAGS.cls_drop),\n            nn.Linear(FLAGS.encoder_units, get_vocab_size()),\n        )\n      \n  # TODO check training flag ok\n  def encode(self, frames):\n    return self.encoder(frames)\n\n  def forward_(self, frames):\n    x = self.encode(frames)\n    x = self.classifer(x)\n    return x\n\n  def forward(self, inputs):\n    if self.training:\n      self.input_ = inputs\n    x = self.forward_(inputs['frames'])\n    res = {\n      'pred': x,\n    }\n    return res\n  \n  def infer(self, frames):\n    return self.forward_(frames)\n\n  def get_loss_fn(self):  \n    def ctc_loss(loss_obj, preds, labels, labels_lengths, weights=None):\n      preds = F.log_softmax(preds, dim=-1)\n      preds_lengths = torch.sum(torch.ones_like(preds[:,:,0]).long(), dim=-1)\n      loss = loss_obj(preds.transpose(0, 1), labels, preds_lengths, labels_lengths)\n      if weights is not None:\n        loss = torch.mean(loss * weights)\n      return loss\n      \n    def loss_fn(res, labels, x, step=None, epoch=None, training=None):\n      scalars = {}\n      weights = None\n      reduction = 'mean'\n      if FLAGS.mix_sup:\n        weights = x['weight']\n        reduction = 'none'\n      loss_obj = nn.CTCLoss(zero_infinity=True, reduction=reduction)\n      preds = res['pred'].float()\n      labels = labels.float()\n      labels_lengths = torch.sum((labels != PAD_IDX).long(), dim=-1)\n      loss = ctc_loss(loss_obj, preds, labels, labels_lengths, weights)      \n      return loss\n\n    return loss_fn\n  \n  def get_infer_model(self):\n    return InferModel(self)","metadata":{"execution":{"iopub.status.busy":"2023-08-27T08:05:07.129188Z","iopub.execute_input":"2023-08-27T08:05:07.129625Z","iopub.status.idle":"2023-08-27T08:05:07.146184Z","shell.execute_reply.started":"2023-08-27T08:05:07.129593Z","shell.execute_reply":"2023-08-27T08:05:07.144801Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Code for squeezeformer, notice most code copy/modify from NEMO, and for ROPE encoding copy from huggingface llma part","metadata":{}},{"cell_type":"code","source":"# simplfiy from NEMO\nclass TimeReductionModule(nn.Module):\n    \"\"\"\n    Squeezeformer Time Reduction procedure. Downsamples the audio by `stride` in the time dimension.\n\n    Args:\n        d_model (int): input dimension of MultiheadAttentionMechanism and PositionwiseFeedForward\n        out_dim (int): Output dimension of the module.\n        kernel_size (int): Conv kernel size for depthwise convolution in convolution module\n        stride (int): Downsampling factor in time dimension.\n    \"\"\"\n\n    def __init__(self, d_model: int, out_dim: int, kernel_size: int = 5, stride: int = 2):\n        super().__init__()\n\n        self.d_model = d_model\n        self.out_dim = out_dim\n        self.kernel_size = kernel_size\n        self.stride = stride\n        ## NOTICE modify here so can dived by 2...\n        # self.padding = max(0, self.kernel_size - self.stride) \n        ##  # like k=5, stride=2 here padding is 2 which make 320 -> 160 -> 80\n        self.padding = (self.kernel_size + 1) // self.stride - 1\n\n        self.dw_conv = nn.Conv1d(\n            in_channels=d_model,\n            out_channels=d_model,\n            kernel_size=kernel_size,\n            stride=stride,\n            padding=self.padding,\n            groups=d_model,\n        )\n\n        self.pw_conv = nn.Conv1d(\n            in_channels=d_model, out_channels=out_dim, kernel_size=1, stride=1, padding=0, groups=1,\n        )\n\n        self.reset_parameters()\n\n    def forward(self, x):\n        x = x.transpose(1, 2)  # [B, C, T]\n        x = self.dw_conv(x)\n        x = self.pw_conv(x)\n        x = x.transpose(1, 2)  # [B, T, C]\n        return x\n\n    def reset_parameters(self):\n        dw_max = self.kernel_size ** -0.5\n        pw_max = self.d_model ** -0.5\n\n        with torch.no_grad():\n            torch.nn.init.uniform_(self.dw_conv.weight, -dw_max, dw_max)\n            torch.nn.init.uniform_(self.dw_conv.bias, -dw_max, dw_max)\n            torch.nn.init.uniform_(self.pw_conv.weight, -pw_max, pw_max)\n            torch.nn.init.uniform_(self.pw_conv.bias, -pw_max, pw_max)","metadata":{"execution":{"iopub.status.busy":"2023-08-27T08:05:07.751411Z","iopub.execute_input":"2023-08-27T08:05:07.751846Z","iopub.status.idle":"2023-08-27T08:05:07.765889Z","shell.execute_reply.started":"2023-08-27T08:05:07.75181Z","shell.execute_reply":"2023-08-27T08:05:07.764445Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Swish(nn.SiLU):\n    \"\"\"\n    Swish activation function introduced in 'https://arxiv.org/abs/1710.05941'\n    Mathematically identical to SiLU. See note in nn.SiLU for references.\n    \"\"\"\n    \n    \nclass CausalConv1D(nn.Conv1d):\n    \"\"\"\n    A causal version of nn.Conv1d where each step would have limited access to locations on its right or left\n    All arguments are the same as nn.Conv1d except padding.\n\n    If padding is set None, then paddings are set automatically to make it a causal convolution where each location would not see any steps on its right.\n\n    If padding is set as a list (size of 2), then padding[0] would be used as left padding and padding[1] as right padding.\n    It would make it possible to control the number of steps to be accessible on the right and left.\n    This mode is not supported when stride > 1. padding[0]+padding[1] should be equal to (kernel_size - 1).\n    \"\"\"\n\n    def __init__(\n        self,\n        in_channels: int,\n        out_channels: int,\n        kernel_size: int,\n        stride: int = 1,\n        padding: Union[str, int] = 0,\n        dilation: int = 1,\n        groups: int = 1,\n        bias: bool = True,\n        padding_mode: str = 'zeros',\n        device=None,\n        dtype=None,\n    ) -> None:\n        self.cache_drop_size = None\n        if padding is None:\n            self._left_padding = kernel_size - 1\n            self._right_padding = stride - 1\n        else:\n            if stride != 1 and padding != kernel_size - 1:\n                raise ValueError(\"No striding allowed for non-symmetric convolutions!\")\n            if isinstance(padding, int):\n                self._left_padding = padding\n                self._right_padding = padding\n            elif isinstance(padding, list) and len(padding) == 2 and padding[0] + padding[1] == kernel_size - 1:\n                self._left_padding = padding[0]\n                self._right_padding = padding[1]\n            else:\n                raise ValueError(f\"Invalid padding param: {padding}!\")\n\n        self._max_cache_len = self._left_padding\n\n        super(CausalConv1D, self).__init__(\n            in_channels=in_channels,\n            out_channels=out_channels,\n            kernel_size=kernel_size,\n            stride=stride,\n            padding=0,\n            dilation=dilation,\n            groups=groups,\n            bias=bias,\n            padding_mode=padding_mode,\n            device=device,\n            dtype=dtype,\n        )\n\n    def update_cache(self, x, cache=None):\n        if cache is None:\n            new_x = F.pad(x, pad=(self._left_padding, self._right_padding))\n            next_cache = cache\n        else:\n            new_x = F.pad(x, pad=(0, self._right_padding))\n            new_x = torch.cat([cache, new_x], dim=-1)\n            if self.cache_drop_size > 0:\n                next_cache = new_x[:, :, : -self.cache_drop_size]\n            else:\n                next_cache = new_x\n            next_cache = next_cache[:, :, -cache.size(-1) :]\n        return new_x, next_cache\n\n    def forward(self, x, cache=None):\n        x, cache = self.update_cache(x, cache=cache)\n        x = super().forward(x)\n        if cache is None:\n            return x\n        else:\n            return x, cache\n\nclass ConformerFeedForward(nn.Module):\n    \"\"\"\n    feed-forward module of Conformer model.\n    \"\"\"\n\n    def __init__(self, d_model, d_ff, dropout, activation=Swish()):\n        super(ConformerFeedForward, self).__init__()\n        self.d_model = d_model\n        self.d_ff = d_ff\n        self.linear1 = nn.Linear(d_model, d_ff)\n        self.activation = activation\n        self.dropout = nn.Dropout(p=dropout)\n        self.linear2 = nn.Linear(d_ff, d_model)\n\n    def forward(self, x):\n        x = self.linear1(x)\n        x = self.activation(x)\n        x = self.dropout(x)\n        x = self.linear2(x)\n        return x\n\n    def reset_parameters_ff(self):\n        ffn1_max = self.d_model ** -0.5\n        ffn2_max = self.d_ff ** -0.5\n        with torch.no_grad():\n            nn.init.uniform_(self.linear1.weight, -ffn1_max, ffn1_max)\n            nn.init.uniform_(self.linear1.bias, -ffn1_max, ffn1_max)\n            nn.init.uniform_(self.linear2.weight, -ffn2_max, ffn2_max)\n            nn.init.uniform_(self.linear2.bias, -ffn2_max, ffn2_max)\n\nactivation_registry = {\n    \"identity\": nn.Identity,\n    \"hardtanh\": nn.Hardtanh,\n    \"relu\": nn.ReLU,\n    \"selu\": nn.SELU,\n    \"swish\": nn.SiLU,\n    \"silu\": nn.SiLU,\n    \"gelu\": nn.GELU,\n}\n\nclass ConformerConvolution(nn.Module):\n    \"\"\"The convolution module for the Conformer model.\n    Args:\n        d_model (int): hidden dimension\n        kernel_size (int): kernel size for depthwise convolution\n        pointwise_activation (str): name of the activation function to be used for the pointwise conv.\n            Note that Conformer uses a special key `glu_` which is treated as the original default from\n            the paper.\n    \"\"\"\n\n    def __init__(\n        self, d_model, kernel_size, norm_type='batch_norm', conv_context_size=None, pointwise_activation='glu_'\n    ):\n        super(ConformerConvolution, self).__init__()\n        assert (kernel_size - 1) % 2 == 0\n        self.d_model = d_model\n        self.kernel_size = kernel_size\n        self.norm_type = norm_type\n\n        if conv_context_size is None:\n            conv_context_size = (kernel_size - 1) // 2\n\n        if pointwise_activation in activation_registry:\n            self.pointwise_activation = activation_registry[pointwise_activation]()\n            dw_conv_input_dim = d_model * 2\n\n            if hasattr(self.pointwise_activation, 'inplace'):\n                self.pointwise_activation.inplace = True\n        else:\n            self.pointwise_activation = pointwise_activation\n            dw_conv_input_dim = d_model\n\n        self.pointwise_conv1 = nn.Conv1d(\n            in_channels=d_model, out_channels=d_model * 2, kernel_size=1, stride=1, padding=0, bias=True\n        )\n\n        self.depthwise_conv = CausalConv1D(\n            in_channels=dw_conv_input_dim,\n            out_channels=dw_conv_input_dim,\n            kernel_size=kernel_size,\n            stride=1,\n            padding=conv_context_size,\n            groups=dw_conv_input_dim,\n            bias=True,\n        )\n\n        if norm_type == 'batch_norm':\n            self.batch_norm = nn.BatchNorm1d(dw_conv_input_dim)\n        elif norm_type == 'instance_norm':\n            self.batch_norm = nn.InstanceNorm1d(dw_conv_input_dim)\n        elif norm_type == 'layer_norm':\n            self.batch_norm = nn.LayerNorm(dw_conv_input_dim)\n        elif norm_type.startswith('group_norm'):\n            num_groups = int(norm_type.replace(\"group_norm\", \"\"))\n            self.batch_norm = nn.GroupNorm(num_groups=num_groups, num_channels=d_model)\n        else:\n            raise ValueError(f\"conv_norm_type={norm_type} is not valid!\")\n\n        self.activation = Swish()\n        self.pointwise_conv2 = nn.Conv1d(\n            in_channels=dw_conv_input_dim, out_channels=d_model, kernel_size=1, stride=1, padding=0, bias=True\n        )\n\n    def forward(self, x, pad_mask=None, cache=None):\n        x = x.transpose(1, 2)\n        x = self.pointwise_conv1(x)\n\n        # Compute the activation function or use GLU for original Conformer\n        if self.pointwise_activation == 'glu_':\n            x = nn.functional.glu(x, dim=1)\n        else:\n            x = self.pointwise_activation(x)\n\n        if pad_mask is not None:\n            x = x.float().masked_fill(pad_mask.unsqueeze(1), 0.0)\n\n        x = self.depthwise_conv(x, cache=cache)\n        if cache is not None:\n            x, cache = x\n\n        if self.norm_type == \"layer_norm\":\n            x = x.transpose(1, 2)\n            x = self.batch_norm(x)\n            x = x.transpose(1, 2)\n        else:\n            x = self.batch_norm(x)\n\n        x = self.activation(x)\n        x = self.pointwise_conv2(x)\n        x = x.transpose(1, 2)\n        if cache is None:\n            return x\n        else:\n            return x, cache\n\n    def reset_parameters_conv(self):\n        pw1_max = pw2_max = self.d_model ** -0.5\n        dw_max = self.kernel_size ** -0.5\n\n        with torch.no_grad():\n            nn.init.uniform_(self.pointwise_conv1.weight, -pw1_max, pw1_max)\n            nn.init.uniform_(self.pointwise_conv1.bias, -pw1_max, pw1_max)\n            nn.init.uniform_(self.pointwise_conv2.weight, -pw2_max, pw2_max)\n            nn.init.uniform_(self.pointwise_conv2.bias, -pw2_max, pw2_max)\n            nn.init.uniform_(self.depthwise_conv.weight, -dw_max, dw_max)\n            nn.init.uniform_(self.depthwise_conv.bias, -dw_max, dw_max)\n            \nclass ScaleBiasLayer(nn.Module):\n    \"\"\"\n    Computes an affine transformation y = x * scale + bias, either learned via adaptive weights, or fixed.\n    Efficient alternative to LayerNorm where we can avoid computing the mean and variance of the input, and\n    just rescale the output of the previous layer.\n\n    Args:\n        d_model (int): input dimension of layer.\n        adaptive_scale (bool): whether to learn the affine transformation parameters or not. If set to False,\n            the scale is fixed to 1 and bias to 0, effectively performing a No-Op on the input.\n            This is done for export compatibility.\n    \"\"\"\n\n    def __init__(self, d_model: int, adaptive_scale: bool):\n        super().__init__()\n        self.adaptive_scale = adaptive_scale\n        if adaptive_scale:\n            self.scale = nn.Parameter(torch.ones(d_model))\n            self.bias = nn.Parameter(torch.zeros(d_model))\n        else:\n            self.register_buffer('scale', torch.ones(d_model), persistent=True)\n            self.register_buffer('bias', torch.zeros(d_model), persistent=True)\n\n    def forward(self, x):\n        scale = self.scale.view(1, 1, -1)\n        bias = self.bias.view(1, 1, -1)\n        return x * scale + bias\n\nclass DepthWiseConv1d(nn.Module):\n\n  def __init__(self, chan_in, chan_out, kernel_size, padding):\n    super().__init__()\n    self.padding = padding\n    self.conv = nn.Conv1d(chan_in, chan_out, kernel_size, groups=chan_in)\n\n  def forward(self, x):\n    x = F.pad(x, self.padding)\n    return self.conv(x)\n  \n# attention, feedforward, and conv module\n\n\nclass Scale(nn.Module):\n\n  def __init__(self, scale, fn):\n    super().__init__()\n    self.fn = fn\n    self.scale = scale\n\n  def forward(self, x, **kwargs):\n    return self.fn(x, **kwargs) * self.scale\n\n\nclass PreNorm(nn.Module):\n\n  def __init__(self, dim, fn):\n    super().__init__()\n    self.fn = fn\n    self.norm = nn.LayerNorm(dim)\n\n  def forward(self, x, **kwargs):\n    x = self.norm(x)\n    return self.fn(x, **kwargs)","metadata":{"execution":{"iopub.status.busy":"2023-08-27T08:05:08.215622Z","iopub.execute_input":"2023-08-27T08:05:08.216021Z","iopub.status.idle":"2023-08-27T08:05:08.268254Z","shell.execute_reply.started":"2023-08-27T08:05:08.215988Z","shell.execute_reply":"2023-08-27T08:05:08.26722Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"ROPE code from huggingface llma","metadata":{}},{"cell_type":"code","source":"# rotary positional embedding\n# https://arxiv.org/abs/2104.09864\n\n# https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py\nclass LlamaRotaryEmbedding(torch.nn.Module):\n    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):\n        super().__init__()\n\n        self.dim = dim\n        self.max_position_embeddings = max_position_embeddings\n        self.base = base\n        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))\n        self.register_buffer(\"inv_freq\", inv_freq, persistent=False)\n\n        # Build here to make `torch.jit.trace` work.\n        self._set_cos_sin_cache(\n            seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()\n        )\n\n    def _set_cos_sin_cache(self, seq_len, device, dtype):\n        self.max_seq_len_cached = seq_len\n        t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)\n\n        # freqs = torch.einsum(\"i,j->ij\", t, self.inv_freq)\n        freqs = torch.outer(t, self.inv_freq)\n        \n        # Different from paper, but it uses a different permutation in order to obtain the same calculation\n        emb = torch.cat((freqs, freqs), dim=-1)\n        self.register_buffer(\"cos_cached\", emb.cos()[None, None, :, :].to(dtype), persistent=False)\n        self.register_buffer(\"sin_cached\", emb.sin()[None, None, :, :].to(dtype), persistent=False)\n\n    def forward(self, x, seq_len=None):\n        # x: [bs, num_attention_heads, seq_len, head_size]\n        if seq_len > self.max_seq_len_cached:\n            self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)\n\n        return (\n            self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype),\n            self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype),\n        )\n\n\nclass LlamaLinearScalingRotaryEmbedding(LlamaRotaryEmbedding):\n    \"\"\"LlamaRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev\"\"\"\n\n    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):\n        self.scaling_factor = scaling_factor\n        super().__init__(dim, max_position_embeddings, base, device)\n\n    def _set_cos_sin_cache(self, seq_len, device, dtype):\n        self.max_seq_len_cached = seq_len\n        t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)\n        t = t / self.scaling_factor\n\n        # freqs = torch.einsum(\"i,j->ij\", t, self.inv_freq)\n        freqs = torch.outer(t, self.inv_freq)\n        # Different from paper, but it uses a different permutation in order to obtain the same calculation\n        emb = torch.cat((freqs, freqs), dim=-1)\n        self.register_buffer(\"cos_cached\", emb.cos()[None, None, :, :].to(dtype), persistent=False)\n        self.register_buffer(\"sin_cached\", emb.sin()[None, None, :, :].to(dtype), persistent=False)\n\n\nclass LlamaDynamicNTKScalingRotaryEmbedding(LlamaRotaryEmbedding):\n    \"\"\"LlamaRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla\"\"\"\n\n    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):\n        self.scaling_factor = scaling_factor\n        super().__init__(dim, max_position_embeddings, base, device)\n\n    def _set_cos_sin_cache(self, seq_len, device, dtype):\n        self.max_seq_len_cached = seq_len\n\n        if seq_len > self.max_position_embeddings:\n            base = self.base * (\n                (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)\n            ) ** (self.dim / (self.dim - 2))\n            inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))\n            self.register_buffer(\"inv_freq\", inv_freq, persistent=False)\n\n        t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)\n\n        # freqs = torch.einsum(\"i,j->ij\", t, self.inv_freq)\n        freqs = torch.outer(t, self.inv_freq)\n        # Different from paper, but it uses a different permutation in order to obtain the same calculation\n        emb = torch.cat((freqs, freqs), dim=-1)\n        self.register_buffer(\"cos_cached\", emb.cos()[None, None, :, :].to(dtype), persistent=False)\n        self.register_buffer(\"sin_cached\", emb.sin()[None, None, :, :].to(dtype), persistent=False)\n\n\ndef rotate_half(x):\n    \"\"\"Rotates half the hidden dims of the input.\"\"\"\n    x1 = x[..., : x.shape[-1] // 2]\n    x2 = x[..., x.shape[-1] // 2 :]\n    # return torch.cat((-x2, x1), dim=-1)\n    # return torch.cat((torch.neg(x2), x1), dim=-1)\n    return torch.cat((x2 * (-1.), x1), dim=-1)\n\ndef apply_rotary_pos_emb(q, k, cos, sin, position_ids):\n    # The first two dimensions of cos and sin are always 1, so we can `squeeze` them.\n    cos = cos.squeeze(1).squeeze(0)  # [seq_len, dim]\n    sin = sin.squeeze(1).squeeze(0)  # [seq_len, dim]\n    cos = cos[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]\n    sin = sin[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]\n    q_embed = (q * cos) + (rotate_half(q) * sin)\n    k_embed = (k * cos) + (rotate_half(k) * sin)\n    return q_embed, k_embed","metadata":{"execution":{"iopub.status.busy":"2023-08-27T08:05:09.248678Z","iopub.execute_input":"2023-08-27T08:05:09.249293Z","iopub.status.idle":"2023-08-27T08:05:09.27703Z","shell.execute_reply.started":"2023-08-27T08:05:09.249251Z","shell.execute_reply":"2023-08-27T08:05:09.275976Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Attention is important which use ROPE relative pos encodeing","metadata":{}},{"cell_type":"code","source":"class Attention(nn.Module):\n  def __init__(self, dim, dim_head=64, \n                heads=8, dropout=0., max_pos_emb=512, \n                relpos_att=True, rope=False):\n    super().__init__()\n    self.scale = dim_head ** -0.5\n    self.heads = heads\n    inner_dim = heads * dim_head\n\n    self.dropout = nn.Dropout(dropout)\n    self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)\n    self.to_out = nn.Linear(inner_dim, dim, bias = False)\n    self.max_position_embeddings = max_pos_emb\n    self.dim_head = dim_head\n    self.relpos_att = relpos_att\n    self.rope = rope\n    if relpos_att:\n      if rope:\n        self._init_rope()\n      else:\n        self.rel_pos_emb = nn.Embedding(2 * max_pos_emb + 1, dim_head)\n\n  def _init_rope(self):\n    scaling_type = FLAGS.scaling_type\n    if scaling_type is None:\n        self.rotary_emb = LlamaRotaryEmbedding(self.dim_head, max_position_embeddings=self.max_position_embeddings)\n    else:\n        scaling_factor = FLAGS.scaling_factor\n        if scaling_type == \"linear\":\n            self.rotary_emb = LlamaLinearScalingRotaryEmbedding(\n                self.dim_head, max_position_embeddings=self.max_position_embeddings, scaling_factor=scaling_factor\n            )\n        elif scaling_type == \"dynamic\":\n            self.rotary_emb = LlamaDynamicNTKScalingRotaryEmbedding(\n                self.dim_head, max_position_embeddings=self.max_position_embeddings, scaling_factor=scaling_factor\n            )\n        else:\n            raise ValueError(f\"Unknown RoPE scaling type {scaling_type}\")\n      \n\n  def forward(self, x):\n    n, device, h = x.shape[-2], x.device, self.heads\n\n    q, k, v = self.to_qkv(x).chunk(3, dim = -1)\n    \n    q = q.view(q.shape[0], q.shape[1], h, -1).permute(0, 2, 1, 3)\n    k = k.view(k.shape[0], k.shape[1], h, -1).permute(0, 2, 1, 3)\n    v = v.view(v.shape[0], v.shape[1], h, -1).permute(0, 2, 1, 3)\n\n    if self.relpos_att:\n      cos, sin = self.rotary_emb(v, seq_len=n)\n      position_ids = torch.arange(0, n, dtype=torch.long, device=device)\n      position_ids = position_ids.unsqueeze(0).view(-1, n)\n      q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids)\n      sim = torch.matmul(q, k.transpose(2, 3)) * self.scale\n    else:\n      sim = torch.matmul(q, k.permute(0, 1, 3, 2)) * self.scale\n    \n    attn = F.softmax(sim, dim=-1, dtype=torch.float32).to(q.dtype)\n    attn = self.dropout(attn)\n    \n    out = torch.matmul(attn, v)\n    out = out.permute(0, 2, 1, 3)\n    out = out.reshape(out.shape[0], out.shape[1], -1)\n    out = self.to_out(out)\n    out = self.dropout(out)\n    return out","metadata":{"execution":{"iopub.status.busy":"2023-08-27T08:05:10.396932Z","iopub.execute_input":"2023-08-27T08:05:10.397439Z","iopub.status.idle":"2023-08-27T08:05:10.418992Z","shell.execute_reply.started":"2023-08-27T08:05:10.3974Z","shell.execute_reply":"2023-08-27T08:05:10.417642Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class SqueezeformerBlock(nn.Module):\n\n  def __init__(self,\n               *,\n               dim,\n               dim_head=64,\n               heads=8,\n               ff_mult=4,\n               conv_expansion_factor=2,\n               conv_kernel_size=31,\n               attn_dropout=0.,\n               ff_dropout=0.,\n               conv_dropout=0.,\n               conv_causal=False,\n               relpos_att=True,\n               rope=False,\n               inst_drop=None,\n               skip_factor=None):\n    super().__init__()\n    # first feed forward module\n    self.norm_feed_forward1 = nn.LayerNorm(dim)\n    self.feed_forward1 = ConformerFeedForward(d_model=dim, d_ff=dim * ff_mult, dropout=ff_dropout)\n    self.feed_forward1_scale = ScaleBiasLayer(d_model=dim, adaptive_scale=True)\n    \n    # convolution module\n    self.norm_conv = nn.LayerNorm(dim)\n    self.conv = ConformerConvolution(\n        d_model=dim, kernel_size=conv_kernel_size, norm_type='batch_norm', pointwise_activation='swish'\n    )\n    self.conv_scale = ScaleBiasLayer(d_model=dim, adaptive_scale=True)\n    \n    # multi-headed self-attention module\n    self.norm_self_att = nn.LayerNorm(dim)\n    self.self_attn = Attention(dim=dim,\n                          dim_head=dim_head,\n                          heads=heads,\n                          dropout=attn_dropout,\n                          max_pos_emb=FLAGS.n_frames,\n                          relpos_att=relpos_att,\n                          rope=rope)\n    self.self_attn_scale = ScaleBiasLayer(d_model=dim, adaptive_scale=True)\n\n    # second feed forward module\n    self.norm_feed_forward2 = nn.LayerNorm(dim)\n    self.feed_forward2 = ConformerFeedForward(d_model=dim, d_ff=dim * ff_mult, dropout=ff_dropout)\n    self.feed_forward2_scale = ScaleBiasLayer(d_model=dim, adaptive_scale=True)\n    \n    self.inst_drop = inst_drop if inst_drop is not None else FLAGS.inst_drop_rate\n    # 0.2\n    self.dropout = InstanceDropout(self.inst_drop)\n    self.skip_factor = skip_factor if skip_factor is not None else FLAGS.skip_factor\n    \n    self.fc_factor = 0.5\n    \n    self.reset_parameters()\n    \n  def reset_parameters(self):\n    # Used for Squeezeformer initialization only\n    self.feed_forward1.reset_parameters_ff()\n    self.feed_forward2.reset_parameters_ff()\n    self.conv.reset_parameters_conv()\n\n  def forward(self, x):\n    residual = x\n    x = self.self_attn_scale(x)\n    x = self.self_attn(x)\n    x = residual + self.dropout(x) * self.skip_factor\n    \n    x = self.norm_self_att(x)\n    residual = x\n    x = self.feed_forward1_scale(x)\n    x = self.feed_forward1(x)\n    x = residual + self.dropout(x) * self.skip_factor * self.fc_factor\n    x = self.norm_feed_forward1(x)\n    residual = x\n\n    x = self.conv_scale(x)\n    x = self.conv(x)\n    x = residual + self.dropout(x) * self.skip_factor\n    x = self.norm_conv(x)\n    residual = x\n    \n    x = self.feed_forward2_scale(x)\n    x = self.feed_forward2(x)\n    x = residual + self.dropout(x) * self.skip_factor * self.fc_factor\n    x = self.norm_feed_forward2(x)\n\n    return x\n\n\nclass Squeezeformer(nn.Module):\n\n  def __init__(self,\n               dim,\n               *,\n               depth,\n               dim_head=64,\n               heads=8,\n               ff_mult=4,\n               conv_expansion_factor=2,\n               conv_kernel_size=31,\n               attn_dropout=0.,\n               ff_dropout=0.,\n               conv_dropout=0.,\n               conv_causal=False):\n    super().__init__()\n    self.dim = dim\n    self.layers = nn.ModuleList([])\n    self.inst_drops = [None] * depth\n    heads_ = heads\n    conv_kernel_size_ = conv_kernel_size\n    for i in range(depth):    \n      relpos_att = True\n      dim_head_ = dim_head\n      rope = True\n      heads_ = heads\n      if i < FLAGS.time_reduce_idx:\n        heads_ = heads // 2\n      \n      self.layers.append(SqueezeformerBlock(dim=dim,\n                          dim_head=dim_head_,\n                          heads=heads_,\n                          ff_mult=ff_mult,\n                          conv_expansion_factor=conv_expansion_factor,\n                          conv_kernel_size=conv_kernel_size_,\n                          conv_causal=conv_causal,\n                          attn_dropout=attn_dropout,\n                          ff_dropout=ff_dropout,\n                          conv_dropout=conv_dropout,\n                          relpos_att=relpos_att,\n                          rope=rope,\n                          inst_drop=self.inst_drops[i]))\n      \n    if FLAGS.time_reduce:\n      if FLAGS.time_reduce_idx < depth - 1:\n        reduction_module = TimeReductionModule(dim, dim, kernel_size=FLAGS.time_kernel_size, stride=FLAGS.time_stride) \n        self.layers.insert(FLAGS.time_reduce_idx, reduction_module)\n      \n  def forward(self, x):\n    for i, layer in enumerate(self.layers):\n      x = layer(x)\n    return x","metadata":{"execution":{"iopub.status.busy":"2023-08-27T08:05:10.890373Z","iopub.execute_input":"2023-08-27T08:05:10.890755Z","iopub.status.idle":"2023-08-27T08:05:10.916085Z","shell.execute_reply.started":"2023-08-27T08:05:10.890727Z","shell.execute_reply":"2023-08-27T08:05:10.91523Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = Model()\nmodel","metadata":{"execution":{"iopub.status.busy":"2023-08-27T08:05:11.36003Z","iopub.execute_input":"2023-08-27T08:05:11.360455Z","iopub.status.idle":"2023-08-27T08:05:11.862391Z","shell.execute_reply.started":"2023-08-27T08:05:11.360424Z","shell.execute_reply":"2023-08-27T08:05:11.861111Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from torchinfo import summary\ninput_shape = (1, FLAGS.n_frames, get_n_cols())\nsummary(model.get_infer_model(), input_shape)","metadata":{"execution":{"iopub.status.busy":"2023-08-27T08:05:12.664959Z","iopub.execute_input":"2023-08-27T08:05:12.665692Z","iopub.status.idle":"2023-08-27T08:05:13.300407Z","shell.execute_reply.started":"2023-08-27T08:05:12.665652Z","shell.execute_reply":"2023-08-27T08:05:13.299269Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# TODO you may load weights of trained torch model here at first","metadata":{}},{"cell_type":"code","source":"# TFLite model for submission\nclass TFLiteModel(tf.keras.Model):\n\n  def __init__(self, model):\n    super(TFLiteModel, self).__init__()\n\n    # Load the feature generation and main models\n    self.preprocess_layer = PreprocessLayer(FLAGS.n_frames)\n    self.model = model\n\n  @tf.function(jit_compile=True)\n  def infer(self, frames):\n    return self.model(frames)\n\n  @tf.function(input_signature=[\n      tf.TensorSpec(shape=[None, N_COLS], dtype=tf.float32, name='inputs')\n  ])\n  def call(self, inputs):\n    # Preprocess Data\n    frames_inp = self.preprocess_layer(inputs)\n    # x = frames_inp\n    # Add Batch Dimension\n    frames_inp = tf.expand_dims(frames_inp, axis=0)\n\n    outputs = self.infer(frames_inp)\n    # y = outputs\n \n    # Squeeze outputs\n    outputs = tf.squeeze(outputs, axis=0)\n    outputs = decode_phrase(outputs)\n    \n    # for 0 is PAD_IDX\n    outputs -= 1\n    # outputs = tf.one_hot(outputs, get_vocab_size())\n    # vocab_size = 61 if not FLAGS.no_eos else 60\n    vocab_size = get_vocab_size()\n    outputs = tf.one_hot(outputs, vocab_size)\n    if FLAGS.decode_phrase_type:\n      ouputs = outputs[1:]\n\n    # Return a dictionary with the output tensor\n    return {'outputs': outputs}\n    # return {\n    #   'outputs': outputs,\n    #   'frames': x,\n    #   'intermediate': y,\n    # }\n\n  # TODO not work... an intermediate Keras symbolic input/output, to a TF API that does not allow registering custom dispatchers, such as `tf.cond`, `tf.function`, gradient tapes, or `tf.map_fn`. Keras Functional model construction only supports TF API calls that *do* support dispatching, such as `tf.math.add` or `tf.reshape`. Other APIs cannot be called directly on symbolic Kerasinputs/outputs. You can work around this limitation by putting the operation in a custom Keras layer `call` and calling that layer on this symbolic input/output.\n  def get_model(self):\n    inputs = tf.keras.layers.Input([N_COLS],\n                                    dtype=tf.float32,\n                                    name='inputs')\n    out = self.call(inputs)\n    model = tf.keras.models.Model(inputs, out)\n    model.summary()\n    return model","metadata":{"execution":{"iopub.status.busy":"2023-08-27T08:06:30.641217Z","iopub.execute_input":"2023-08-27T08:06:30.641624Z","iopub.status.idle":"2023-08-27T08:06:30.654944Z","shell.execute_reply.started":"2023-08-27T08:06:30.641593Z","shell.execute_reply":"2023-08-27T08:06:30.653926Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"FLAGS.model_dir = './'\n\ndef get_tflite_model(model):\n  tflite_keras_model = TFLiteModel(model)\n  return tflite_keras_model\n\ndef to_tflite_model(model):\n  tflite_keras_model = get_tflite_model(model)\n  # tflite_keras_model = tflite_keras_model.get_model()\n  # plot_model(tflite_func_model, to_file=f'{FLAGS.model_dir}/tflite.png', show_shapes=True, show_layer_names=True)\n  ic(tflite_keras_model)\n  # Create Model Converter\n  converter = tf.lite.TFLiteConverter.from_keras_model(tflite_keras_model)\n  converter.optimizations = [tf.lite.Optimize.DEFAULT]\n  converter.target_spec.supported_types = [tf.float16]\n  #  1/4 size but seems hurt acc a lot like 0.8 to 0.7\n  # converter.target_spec.supported_ops = [tf.lite.OpsSet.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8,\n  #                                        tf.lite.OpsSet.TFLITE_BUILTINS]\n  converter._experimental_default_to_single_batch_in_tensor_list_ops = True\n  \n  # converter.post_training_quantize = True\n  ic(converter)\n  ## converter.experimental_new_converter = True\n  # Convert Model\n  tflite_model = converter.convert()\n  # ic(tflite_model)\n  # Write Model\n  with open(f'{FLAGS.model_dir}/model.tflite', 'wb') as f:\n    f.write(tflite_model)\n\n  os.system(f'du -h {FLAGS.model_dir}/model.tflite')","metadata":{"execution":{"iopub.status.busy":"2023-08-27T08:06:31.481958Z","iopub.execute_input":"2023-08-27T08:06:31.482387Z","iopub.status.idle":"2023-08-27T08:06:31.491776Z","shell.execute_reply.started":"2023-08-27T08:06:31.482353Z","shell.execute_reply":"2023-08-27T08:06:31.49027Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install sty\nsys.path.append('../input/aslfr-nobuco/')\nimport nobuco","metadata":{"execution":{"iopub.status.busy":"2023-08-27T08:06:31.694382Z","iopub.execute_input":"2023-08-27T08:06:31.694815Z","iopub.status.idle":"2023-08-27T08:06:45.210914Z","shell.execute_reply.started":"2023-08-27T08:06:31.694781Z","shell.execute_reply":"2023-08-27T08:06:45.209516Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"FLAGS.convert_trace = False","metadata":{"execution":{"iopub.status.busy":"2023-08-27T08:06:45.213531Z","iopub.execute_input":"2023-08-27T08:06:45.213906Z","iopub.status.idle":"2023-08-27T08:06:45.221966Z","shell.execute_reply.started":"2023-08-27T08:06:45.21387Z","shell.execute_reply":"2023-08-27T08:06:45.220572Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def torch2keras(model):\n  import nobuco\n  from nobuco import ChannelOrder, ChannelOrderingStrategy\n  from nobuco.layers.weight import WeightLayer\n  input_shape = [1, FLAGS.n_frames, get_n_cols()]\n  ic(input_shape)\n  dummy_input = torch.randn(input_shape)\n  model = model.get_infer_model()\n  # # https://www.kaggle.com/competitions/asl-signs/discussion/406411\n  # model = model.half().float()\n  model = model.to('cpu')\n  model = model.eval()\n  \n  ## TODO help merge to nobuco codebase\n  # for this you could use x = F.avg_pool1d(x, FLAGS.encode_pool_size)\n  @nobuco.converter(F.avg_pool1d, channel_ordering_strategy=ChannelOrderingStrategy.FORCE_TENSORFLOW_ORDER)\n  def converter_avg_pool1d(input: torch.Tensor, input2: int, inplace: bool = False):\n    return lambda input, input2, inplace=False: tf.keras.layers.AveragePooling1D(input2)(input)\n  \n  @nobuco.converter(nn.AvgPool1d, channel_ordering_strategy=ChannelOrderingStrategy.FORCE_TENSORFLOW_ORDER)\n  def converter_AvgPool1d(self, input: torch.Tensor):\n    return tf.keras.layers.AveragePooling1D(self.kernel_size)\n    \n  @nobuco.converter(torch.ones_like, channel_ordering_strategy=ChannelOrderingStrategy.MINIMUM_TRANSPOSITIONS)\n  def converter_ones_like(input, *args, **kwargs):\n      def func(input, *args, **kwargs):\n          return tf.ones_like(input)\n      return func\n  \n  @nobuco.converter(torch.Tensor.bmm, channel_ordering_strategy=ChannelOrderingStrategy.MINIMUM_TRANSPOSITIONS_OR_PYTORCH, autocast=True)\n  def converter_bmm(self, value, *args, **kwargs):\n    def func(self, value, *args, **kwargs):\n        return tf.matmul(self, value)\n    return func\n  \n  @nobuco.converter(torch.flip, channel_ordering_strategy=ChannelOrderingStrategy.MINIMUM_TRANSPOSITIONS)\n  def converter_flip(input, input2, *args, **kwargs):\n      def func(input, input2, *args, **kwargs):\n          return tf.reverse(input, input2)\n      return func\n  \n  \n  input_shape[0] = None\n  keras_model = nobuco.pytorch_to_keras(\n      model,\n      args=[dummy_input], \n      input_shapes={dummy_input: input_shape},\n      inputs_channel_order=ChannelOrder.PYTORCH,\n      outputs_channel_order=ChannelOrder.PYTORCH,\n      trace_shape=True, \n      debug_traces=FLAGS.convert_trace,\n    )\n  keras_model.summary()\n  return keras_model\n  ","metadata":{"execution":{"iopub.status.busy":"2023-08-27T08:06:45.225562Z","iopub.execute_input":"2023-08-27T08:06:45.226001Z","iopub.status.idle":"2023-08-27T08:06:45.245537Z","shell.execute_reply.started":"2023-08-27T08:06:45.225968Z","shell.execute_reply":"2023-08-27T08:06:45.244153Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Convert torch to keras using nobuco  \nYou may notice there are errors, but does not matter, it could convert finally and work fine","metadata":{}},{"cell_type":"code","source":"%%capture\nkeras_model = torch2keras(model)","metadata":{"_kg_hide-output":true,"execution":{"iopub.status.busy":"2023-08-27T08:06:45.248777Z","iopub.execute_input":"2023-08-27T08:06:45.249266Z","iopub.status.idle":"2023-08-27T08:07:39.250373Z","shell.execute_reply.started":"2023-08-27T08:06:45.249218Z","shell.execute_reply":"2023-08-27T08:07:39.249031Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"FLAGS.pad_thre = 0.3\nFLAGS.decode_phrase_type = False\nto_tflite_model(keras_model)","metadata":{"execution":{"iopub.status.busy":"2023-08-27T08:11:16.557551Z","iopub.execute_input":"2023-08-27T08:11:16.558342Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}