{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":52950,"databundleVersionId":5973250,"sourceType":"competition"},{"sourceId":2799445,"sourceType":"datasetVersion","datasetId":1710071},{"sourceId":141055433,"sourceType":"kernelVersion"},{"sourceId":141055455,"sourceType":"kernelVersion"},{"sourceId":141116660,"sourceType":"kernelVersion"}],"dockerImageVersionId":31154,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"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":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:10:37.440265Z","iopub.execute_input":"2025-10-08T10:10:37.440547Z","iopub.status.idle":"2025-10-08T10:10:42.06109Z","shell.execute_reply.started":"2025-10-08T10:10:37.44052Z","shell.execute_reply":"2025-10-08T10:10:42.059971Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Import libs","metadata":{}},{"cell_type":"code","source":"!pip install --upgrade transformers huggingface-hub tokenizers -q --no-cache-dir","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:10:42.06214Z","iopub.execute_input":"2025-10-08T10:10:42.062457Z","iopub.status.idle":"2025-10-08T10:10:57.895001Z","shell.execute_reply.started":"2025-10-08T10:10:42.06243Z","shell.execute_reply":"2025-10-08T10:10:57.893899Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:10:57.897443Z","iopub.execute_input":"2025-10-08T10:10:57.898363Z","iopub.status.idle":"2025-10-08T10:11:17.604123Z","shell.execute_reply.started":"2025-10-08T10:10:57.898328Z","shell.execute_reply":"2025-10-08T10:11:17.60346Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:17.604846Z","iopub.execute_input":"2025-10-08T10:11:17.605382Z","iopub.status.idle":"2025-10-08T10:11:17.70518Z","shell.execute_reply.started":"2025-10-08T10:11:17.60536Z","shell.execute_reply":"2025-10-08T10:11:17.704264Z"}},"outputs":[],"execution_count":null},{"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 = 256\n  eval_batch_size = 256\n  awp = False\n  adv_start_epoch = None\n  adv_lr = 0.2\n  adv_eps = 0\n  fp16 = True # 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":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:17.706223Z","iopub.execute_input":"2025-10-08T10:11:17.706471Z","iopub.status.idle":"2025-10-08T10:11:17.751073Z","shell.execute_reply.started":"2025-10-08T10:11:17.706451Z","shell.execute_reply":"2025-10-08T10:11:17.750081Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"FLAGS.batch_size = 128\nFLAGS.acc_steps = 4 # gradient acc\nFLAGS.batch_size = FLAGS.batch_size // FLAGS.acc_steps","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:17.751988Z","iopub.execute_input":"2025-10-08T10:11:17.752678Z","iopub.status.idle":"2025-10-08T10:11:17.771628Z","shell.execute_reply.started":"2025-10-08T10:11:17.75265Z","shell.execute_reply":"2025-10-08T10:11:17.770833Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:17.772423Z","iopub.execute_input":"2025-10-08T10:11:17.772749Z","iopub.status.idle":"2025-10-08T10:11:18.273081Z","execution_failed":"2025-10-08T10:57:47.656Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Tfrecord dataset","metadata":{}},{"cell_type":"code","source":"def gen_inputs(files, \n           decode_fn, \n           batch_size=64,\n           post_decode_fn=None,\n           num_epochs = None, \n           num_threads=None, \n           buffer_size = 15000, #change from 1000 to 15000\n           dynamic_pad=True,\n           shuffle=True,\n           shuffle_batch=None,\n           shuffle_files=None,\n           ordered=None,\n           min_after_dequeue=None, #depreciated\n           seed=None, \n           enqueue_many=False,  #depreciated\n           fixed_random=False, \n           drop_remainder=False, \n           num_prefetch_batches=None, \n           bucket_boundaries=None,\n           length_index=None,\n           length_key=None,\n           length_fn=None,\n           bucket_batch_sizes=None,\n           repeat=True,\n           initializable=False,\n           filter_fn=None,\n           balance_pos_neg=False,\n           pos_filter_fn=None,\n           neg_filter_fn=None,\n           count_fn=None,\n           return_iterator=False,\n           Dataset=None,\n           batch_parse=False, #by default will be line parse\n           hvd_shard=True,\n           shard_by_files=False,\n           training=False,\n           simple_parse=False,\n           repeat_then_shuffle=False,\n           cache=False,\n           cache_file='',\n           cache_after_map=False,\n           device=None,\n           world_size=1,\n           rank=0,\n           parallel_read_files=False,\n           use_feed_dict=False,\n           feed_name=None,\n           padding_values=None,\n           distribute_strategy=None,\n           torch=False,\n           keras=False,\n           subset=None,\n           return_numpy=False,\n           name='input'):\n  Dataset = Dataset or tf.data.TFRecordDataset\n  AUTO = tf.data.AUTOTUNE\n  use_horovod = False\n\n  def shard(d):\n    return d.shard(hvd.size(), hvd.rank())\n\n  # Choose to use cpu outside input function like in dataset.py\n  #with tf.device('/cpu:0'):\n  if isinstance(files, str):\n    files = gezi.list_files(files)\n  assert len(files) > 0\n\n  if not num_threads:\n    num_threads = 8\n\n  if 'batch_size' in inspect.getfullargspec(decode_fn).args:\n    decode_fn_ = decode_fn\n    def decode_function(example):\n      return decode_fn_(example, batch_size)\n    decode_fn = decode_function\n    \n  if not num_epochs: \n    num_epochs = None\n\n  if shuffle:\n    if shuffle_files is None:\n      shuffle_files = True\n    if shuffle_batch is None:\n      shuffle_batch = True\n  else:\n    if shuffle_files is None:\n      shuffle_files = False\n    if shuffle_batch is None:\n      shuffle_batch = False\n    # TDO 并行读取就会打乱顺序？\n    if not shuffle_files:\n      parallel_read_files = False\n\n  if fixed_random:\n    if seed is None:\n      seed = 1024\n  else:\n    pass\n\n  num_files = len(files)\n  if use_feed_dict and feed_name:\n    files = tf.compat.v1.placeholder(tf.string, [None], feed_name)\n    gezi.set_global(feed_name, files)\n\n  if not num_prefetch_batches:\n    #num_prefetch_batches = num_threads + 3\n    if buffer_size:\n      num_prefetch_batches = int(buffer_size / batch_size)\n    # else:\n    #   num_prefetch_batches = 100\n  \n  if not buffer_size and num_prefetch_batches:\n    buffer_size = num_prefetch_batches * batch_size\n    \n  options = tf.data.Options()\n  try:\n    options.threading.private_threadpool_size = num_threads\n    options.threading.max_intra_op_parallelism = 1\n  except Exception:\n    options.experimental_threading.private_threadpool_size = num_threads\n    options.experimental_threading.max_intra_op_parallelism = 1\n\n  options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.DATA\n  options.experimental_deterministic = True\n\n  if shuffle and not fixed_random:\n    options.experimental_deterministic = False\n\n  if not ordered:\n    options.experimental_deterministic = False\n\n  if not parallel_read_files or num_files == 1:\n    d = Dataset(files)\n    d = d.with_options(options)\n    if use_horovod and hvd_shard:\n      d = shard(d)\n    if not use_horovod and world_size > 1:\n      d = d.shard(world_size, rank)\n  else:\n    try:\n      if shffle_files and (use_horovod or world_size > 1):\n        assert seed\n      d = tf.data.Dataset.list_files(files, shuffle=shuffle_files, seed=seed)\n      d = d.with_options(options)\n    except Exception:\n      d = tf.data.Dataset.from_tensor_slices(files)\n      d = d.with_options(options)\n    # here shard by files, not work good, especially for text line dataset with horovod\n    if use_horovod and shard_by_files:\n      d = shard(d)\n    elif world_size > 1 and shard_by_files:\n      d = d.shard(world_size, rank)\n\n    d = d.interleave(Dataset,\n                  #  cycle_length=min(len(files), 1000),  # in tf 1.14 must set and can not set as AUTOTUNE for tf 2.1 with default as AUTOTUNE\n                  block_length=1,\n                  num_parallel_calls=AUTO)\n\n    if world_size > 1 and not shard_by_files:\n      d = d.shard(world_size, rank)\n\n  if repeat and repeat_then_shuffle:\n    d = d.repeat(num_epochs)\n\n  if cache and (not FLAGS.cache_after_map):\n    d = d.cache(cache_file)\n    \n  # must batch then map if use pyfunc which you might use batch_parse, here batch_parse means batch parse otherwise slower but simple and powerfull...\n  if not batch_parse:\n    d = d.map(decode_fn, num_parallel_calls=AUTO)\n    if cache and cache_after_map:\n      d = d.cache(cache_file)\n  \n  if filter_fn is not None and not batch_parse:\n    d = d.filter(filter_fn)\n\n  if shuffle_batch:\n    d = d.shuffle(buffer_size=buffer_size, seed=seed, reshuffle_each_iteration=True)\n\n  # shuffle then repeat\n  if repeat and not repeat_then_shuffle:\n    d = d.repeat(num_epochs)\n  \n  if dynamic_pad:\n    if not batch_parse: \n      d = d.padded_batch(batch_size, drop_remainder=drop_remainder)\n    else:\n      d = d.batch(batch_size, drop_remainder=drop_remainder)\n  else:\n    d = d.batch(batch_size, drop_remainder=drop_remainder)\n\n  if batch_parse:\n    d = d.map(decode_fn, num_parallel_calls=AUTO)\n    if filter_fn is not None:\n      try:\n        d = d.unbatch()\n        d = d.filter(filter_fn)\n      except Exception:\n        d = d.unbatch()\n\n      d = d.batch(batch_size, drop_remainder=drop_remainder)\n\n  if post_decode_fn is not None:\n    d = d.map(post_decode_fn, num_parallel_calls=AUTO)\n\n  if cache and FLAGS.cache_after_map:\n    logging.debug('Cache datase after map')\n    d = d.cache(cache_file)\n\n  d = d.prefetch(FLAGS.prefetch or AUTO)\n\n  if not return_numpy:    \n    return d\n  else:\n    return d.as_numpy_iterator()\n","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:18.273967Z","iopub.execute_input":"2025-10-08T10:11:18.274604Z","iopub.status.idle":"2025-10-08T10:11:18.292167Z","execution_failed":"2025-10-08T10:57:47.657Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def decode_example(x):\n  if tf.executing_eagerly():\n    x = x.numpy()\n  x = tf.train.Example.FromString(x).features.feature\n  features = {}\n  for key in x.keys():\n    typenames = ['bytes_list', 'float_list', 'int64_list']\n    dtypes = [object, np.float32, np.int64]\n    for typename, dtype in zip(typenames, dtypes):\n      value = getattr(x[key], typename).value\n      if value:\n        features[key] = np.array(value, dtype=dtype)\n  return features\n\ndef first_example(record_file):\n  if isinstance(record_file, (list, tuple)):\n    record_file = record_file[0]\n  if tf.executing_eagerly():\n    for item in tf.data.TFRecordDataset(record_file):\n      x = decode_example(item)\n      return x\n  else:\n    for item in tf.compat.v1.python_io.tf_record_iterator(record_file):\n      x = decode_example(item)\n      return x\n\ndef npdtype2tfdtype(dtype, large=False):\n  if dtype == np.float32:\n    return tf.float32\n  if dtype == np.int32:\n    if not large:\n      return tf.int32\n    else:\n      return tf.int64\n  if dtype == np.int64:\n    return tf.int64\n  if dtype == np.float64:\n    return tf.float32\n  return tf.string\n\ndef sparse_tensor_to_dense(input_tensor, default_value=0):  \n  return tf.sparse.to_dense(input_tensor, default_value=default_value, validate_indices=False)\n\ndef sparse2dense(features, key=None, default_value=0):\n\n  def sparse2dense_(features, key, default_value):\n    val = features[key]\n    if val.values.dtype == tf.string:\n      default_value = None\n    val = sparse_tensor_to_dense(val, default_value)\n    features[key] = val\n\n  modified = False\n  if key:\n    sparse2dense_(features, key)\n    modified = True\n  else:\n    from tensorflow.python.framework.sparse_tensor import SparseTensor\n    for key, val in features.items():\n      if isinstance(val, SparseTensor):\n        sparse2dense_(features, key, default_value)\n        modified = True\n  return modified\n\ndef get_num_records_single(tf_record_file, recount=False):\n  if not recount:\n    filename = os.path.basename(tf_record_file)\n    filename = filename.replace('-', '.').replace('_', '.')\n    l = filename.split('.')\n\n    for item in reversed(l):\n      if item.isdigit():\n        return int(item)\n\n  # try:\n  return sum(\n      1 for _ in tf.compat.v1.python_io.tf_record_iterator(tf_record_file))\n  # except Exception:\n  #   return 0\n\n\ndef get_num_records(files, recount=False):\n  if isinstance(files, str):\n    files = gezi.list_files(files)\n  res = sum([\n      get_num_records_single(file, recount=recount)\n      for file in tqdm(files, ascii=False, desc='get_num_records', leave=False)\n  ])\n  return res\n","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:18.295644Z","iopub.execute_input":"2025-10-08T10:11:18.295892Z","iopub.status.idle":"2025-10-08T10:11:18.317993Z","execution_failed":"2025-10-08T10:57:47.658Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# A wrapper base class for tfrecords related dataset \nclass TfrecordsDataset(object):\n  def __init__(self, \n               subset='valid',\n               batch_size=None,\n               Type=None, \n               files=None,\n               num_instances=None,\n               batch_parse=None,\n               sparse_to_dense=None,\n               hvd_shard=True,\n               use_int32=True,\n               is_info=False,\n               eval_keys=[],\n               incl_keys=[],\n               excl_keys=[],\n               str_keys=[],\n               varlen_keys=[],\n               use_tpu=False,\n               recount=None):\n    self.subset = subset\n    self.filter_fn = None\n    self.pos_filter_fn = None\n    self.neg_filter_fn = None \n    self.count_fn = None\n    self.Type = Type\n    self.batch_parse = batch_parse if batch_parse is not None else FLAGS.batch_parse\n    self.sparse_to_dense = sparse_to_dense if sparse_to_dense is not None else FLAGS.sparse_to_dense\n    self.use_post_decode = None\n    # if self.batch_parse:\n    #   self.sparse_to_dense = False\n    self.batch_size = batch_size or FLAGS.batch_size\n    self.hvd_shard = hvd_shard\n    self.indexes = {'train': -1, 'valid': -1, 'test': -1}\n    self.is_info = is_info\n    self.eval_keys = eval_keys or FLAGS.eval_keys\n    if subset == 'test':\n      self.eval_keys = gezi.get('test_keys') or self.eval_keys\n    self.show_keys = set()  # 如果用户不指定eval_keys 可以用self.show_keys所有非变成以及长度为0,1的key 前提需要使用.adds不能自己外部定义\n    self.excl_keys = excl_keys or FLAGS.excl_keys\n    self.incl_keys = incl_keys or FLAGS.incl_keys\n    self.str_keys = str_keys\n    self.varlen_keys = varlen_keys\n\n    self.parse_fn = tf.io.parse_single_example if not self.batch_parse else tf.io.parse_example\n\n    self.features_dict = {}\n    self.has_varlen_feats = False\n    self.use_tpu = use_tpu\n    try:\n      # TPU detection. No parameters necessary if TPU_NAME environment variable is\n      # set: this is always the case on Kaggle.\n      tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\n      # print('Running on TPU ', tpu.master())\n    except ValueError:\n      tpu = None\n    if tpu is not None:\n      self.use_tpu = True\n    self.use_int32 = use_int32\n    if self.use_tpu:\n      self.use_int32 = True\n\n    self.num_instances_ = num_instances\n    self.files_ = files\n\n    # self.use_post_decode = use_post_decode\n    self.recount = recount or FLAGS.recount_tfrecords\n\n    assert self.subset in ['train', 'valid', 'test'], \\\n          'subset is {} but should in [train, valid, test]'.format(self.subset)\n\n  @staticmethod\n  def get_filenames_(subset=None, shuffle=False):\n    try:\n      if subset in ['train', 'valid', 'test']:\n        if subset == 'train':\n          return FLAGS.train_files\n        elif subset == 'valid':\n          return FLAGS.valid_files\n      else:\n        raise ValueError('Invalid data subset \"%s\"' % subset)\n    except Exception:\n      return None\n\n  def get_filenames(self, subset=None, shuffle=False):\n    subset = subset or self.subset\n    return TfrecordsDataset.get_filenames_(subset, shuffle=False)\n\n  def basic_parse(self, example):\n    self.auto_parse(keys=self.incl_keys, exclude_keys=self.excl_keys)\n    if self.varlen_keys:\n      self.adds_varlens(self.varlen_keys)\n    fe = self.parse_(serialized=example)\n    return fe\n  \n  # override this\n  def parse(self, example):\n    return self.basic_parse(example)\n\n  def decode(self, example):\n    l = self.parse(example)\n    \n    if isinstance(l, (list, tuple)):\n      features = l[0]\n    else:\n      features = l\n    # self.use_tpu = True\n    if isinstance(features, dict):\n      if self.use_tpu:\n        def decode_label(label):\n          label = tf.io.decode_raw(label, tf.uint8)  # tf.string -> [tf.uint8]\n          label = tf.reshape(label, [])  # label is a scalar\n          return tf.cast(label, tf.int32) \n        for key in features.keys():\n          if features[key].dtype in [tf.int64, tf.uint8, tf.uint16, tf.uint32]:\n            features[key] = tf.cast(features[key], tf.int32)\n\n        if not self.is_info:\n          keys = list(features.keys())\n          for key in keys:\n            if features[key].dtype ==tf.string:\n              del features[key]\n\n              if key in self.eval_keys:\n                FLAGS.use_info_dataset = True  # 因为训练model的dataset不再包含eval的某个信息 需要依赖再遍历一遍info_dataset\n              # features[key] = tf.ones_like(features[key], tf.int32)\n              # features[key] = decode_label(features[key]) ## not work TODO\n\n      else:\n        def _cast_dict(features):\n          for key in features:\n            if isinstance(features[key], dict):\n              _cast_dict(features[key])\n            else:\n              # tf.print(key, features[key])\n              if features[key].dtype == tf.int64 and self.use_int32:\n                features[key] = tf.cast(features[key], tf.int32)\n        _cast_dict(features)\n \n      # is_info 只在tf2 keras模式下生效, 都创建 但是有可能不用 只有 FLAGS.use_info_dataset = True 才使用\n      if self.is_info:\n        keys = list(features.keys())\n        if not FLAGS.predict_on_batch:\n          if not self.eval_keys:\n            for key in keys:\n              dim = 1 if self.batch_parse else 0\n              if not (len(features[key].shape) == dim or features[key].shape[dim] == 1):\n                del features[key]\n              else:\n                self.show_keys.add(key)\n          else:\n            for key in keys:\n              if key not in self.eval_keys:\n                del features[key]\n      else:\n        keys = list(features.keys())\n        for key in keys:\n          if key in self.excl_keys:\n            del features[key]\n\n    return l\n\n  def adjust(self, result):\n    return result\n\n  def parse_(self, serialized, features=None):\n    features = features or self.features_dict\n    # ic(features)\n    features = self.parse_fn(serialized=serialized, features=features)\n    if FLAGS.exclude_varlen_keys:\n      from tensorflow.python.framework.sparse_tensor import SparseTensor\n      sparse_keys = [key for key in features if isinstance(key, SparseTensor)]\n      for key in sparse_keys:\n        del features[key]\n    else:\n      if self.sparse_to_dense:\n        modified = sparse2dense(features, default_value=FLAGS.padding_idx)\n        self.has_varlen_feats = modified\n    self.features = features\n    return features\n  \n  def gen_example(self, files=None):\n    if not files:\n      files = self.get_filenames()\n    if not isinstance(files, (list, tuple)):\n      files = [files]\n    example = {}\n    if files:\n      for file in files:\n        try:\n          example = first_example(file)\n        except Exception:\n          ic(traceback.format_exc())\n          ic('bad tfrecord:', file)\n        if example:\n          self.example = example\n          break\n    self.example = example\n    return example\n\n  def gen_input(self, files=None):\n    example = self.gen_example().copy()\n    for key in example:\n      example[key] = np.asarray([example[key]])\n    return example\n\n  def first_input(self, files=None):\n    return self.gen_input(files)\n\n  def add(self, key, dtype=None, length=None, features_dict=None):\n    features_dict = features_dict or self.features_dict\n    dtype_ = dtype\n    if key in self.example:\n      dtype = dtype_ or self.example[key].dtype \n      if length is None:\n        features_dict[key] = tf.io.VarLenFeature(dtype)\n      elif length > 0:\n        features_dict[key] = tf.io.FixedLenFeature([length], dtype)\n      else:\n        features_dict[key] = tf.io.FixedLenFeature([], dtype)\n    \n  def adds(self, keys, dtype=None, length=None, features_dict=None):\n    features_dict = features_dict or self.features_dict\n    dtype_ = dtype\n    for key in keys:\n      if key in self.example:\n        dtype = dtype_ or self.example[key].dtype \n        if length is None:\n          features_dict[key] = tf.io.VarLenFeature(dtype)\n        elif length > 0:\n          features_dict[key] = tf.io.FixedLenFeature([length], dtype)\n        else:\n          features_dict[key] = tf.io.FixedLenFeature([], dtype)\n\n  def auto_parse(self, keys=[], exclude_keys=[], features_dict=None):\n    keys = keys or FLAGS.dataset_keys or self.example.keys()\n    exclude_keys = exclude_keys or FLAGS.dataset_excl_keys\n    keys = [key for key in keys if key not in exclude_keys]\n\n    for key in keys:\n      if key not in self.example:\n        continue\n      length = self.example[key].shape[0]\n      \n      if length == 1:\n        # just to (bs,), tf keras will auto change to (bs,1), also for string 0 is ok\n        length = 0 \n\n      dtype = npdtype2tfdtype(self.example[key].dtype)\n      # print(key, dtype, length, self.example[key])\n      self.adds([key], dtype, length, features_dict)\n\n  def adds_varlens(self, keys=[], exclude_keys=[], features_dict=None):\n    keys = keys or self.example.keys()\n    keys = [key for key in keys if key not in exclude_keys]\n\n    for key in keys:\n      if not key in self.example:\n        continue\n      length = self.example[key].shape[0]\n      dtype = npdtype2tfdtype(self.example[key].dtype)\n      length = None\n      if dtype == tf.string:\n        length = 1\n      self.adds([key], dtype, length, features_dict)  \n  \n  def make_batch(self, \n                 batch_size=None, \n                 filenames=None,\n                 subset=None,\n                 initializable=False,\n                 repeat=None,\n                 shuffle=None,\n                 return_iterator=True,\n                 hvd_shard=None,\n                 simple_parse=False,\n                 num_epochs=None,\n                 cache=False,\n                 cache_file='',\n                 buffer_size=None,\n                 batch_sizes=None,\n                 buckets=None,\n                 drop_remainder=None,\n                 world_size=1,\n                 rank=0,\n                 shard_by_files=None,\n                 distribute_strategy=None,\n                 return_numpy=False):\n    # with tf.device('/cpu:0'):\n    subset = subset or self.subset\n    hvd_shard = hvd_shard if hvd_shard is not None else self.hvd_shard\n    if batch_size is None:\n      is_test = True\n    else:\n      is_test = False\n    batch_size = batch_size or self.batch_size\n    self.batch_size = batch_size\n    batch_sizes = batch_sizes if batch_sizes is not None else FLAGS.batch_sizes\n    buffer_size = buffer_size if buffer_size is not None else FLAGS.buffer_size\n    buckets = buckets if buckets is not None else FLAGS.buckets\n    drop_remainder = drop_remainder if drop_remainder is not None else FLAGS.drop_remainder\n    shard_by_files = shard_by_files if shard_by_files is not None else FLAGS.shard_by_files\n\n    self.return_numpy = return_numpy\n\n    filenames = filenames or self.files_ or self.get_filenames(subset)\n    \n    self.gen_example(filenames)\n\n    is_eager = tf.executing_eagerly()\n\n    self.files_ = filenames\n\n    self.indexes[self.subset] += 1\n    \n    if repeat is None:\n      num_gpus = 1\n      # if subset == 'train' or num_gpus > 1:\n      if subset == 'train':\n        repeat = True\n      else:\n        repeat = False\n      if is_eager and num_gpus == 1 and tf.__version__ < '2':\n        # let tf eager similary to pytorch\n        repeat = False\n\n    if shuffle is None:\n      if subset == 'train':\n        shuffle = FLAGS.shuffle \n      else:\n        shuffle = FLAGS.shuffle_valid \n\n    if drop_remainder is None:\n      if subset == 'train':\n        drop_remainder = True\n      else:\n        drop_remainder = False\n\n    balance_pos_neg=False\n    ic(self.subset, repeat, drop_remainder)\n\n    seed = FLAGS.seed \n    if seed is not None:\n      FLAGS.seed += 1\n\n    ## put on cpu or dummy\n    with tf.device('/cpu'):\n      result = gen_inputs(\n        filenames, \n        decode_fn=self.decode,\n        batch_size=batch_size,\n        post_decode_fn=self.post_decode if hasattr(self, 'post_decode') and self.use_post_decode != False else None,\n        shuffle=shuffle,\n        shuffle_batch=FLAGS.shuffle_batch,\n        shuffle_files=FLAGS.shuffle_files,\n        ordered=FLAGS.dataset_ordered if subset == 'train' else True,\n        num_threads=FLAGS.num_dataset_threads,\n        buffer_size=buffer_size,\n        num_prefetch_batches=FLAGS.num_prefetch_batches,\n        initializable=initializable,\n        repeat=repeat,\n        repeat_then_shuffle=FLAGS.repeat_then_shuffle,\n        drop_remainder=drop_remainder,\n        bucket_boundaries=buckets,\n        bucket_batch_sizes=batch_sizes,\n        length_index=FLAGS.length_index,\n        length_key=FLAGS.length_key,\n        seed=seed,\n        return_iterator=return_iterator,\n        filter_fn=self.filter_fn,  # inside filter_fn judge subset train or valid or test\n        balance_pos_neg=balance_pos_neg,\n        pos_filter_fn=self.pos_filter_fn if subset == 'train' else None,\n        neg_filter_fn=self.neg_filter_fn if subset == 'train' else None,\n        count_fn=self.count_fn if subset == 'train' else None,\n        name=subset,\n        Dataset=self.Type,\n        batch_parse=self.batch_parse,\n        hvd_shard=hvd_shard,\n        shard_by_files=shard_by_files,\n        training=subset == 'train',\n        simple_parse=simple_parse,\n        num_epochs=num_epochs,\n        dynamic_pad=FLAGS.dynamic_pad, #如果有varlen feats才需要 padded_batch 同时batch_parse模式其实也不需要因为sparse2dense就可以自动padd\n        cache=cache,\n        cache_file=cache_file,\n        cache_after_map=FLAGS.cache_after_map,\n        device='/gpu:0',\n        world_size=world_size,\n        rank=rank,\n        fixed_random=FLAGS.fixed_random,\n        parallel_read_files=FLAGS.parallel_read_files,\n        #use_feed_dict=FLAGS.train_loop and FLAGS.rounds > 1 and not is_eager and FLAGS.feed_dataset and tf.__version__ < '2',\n        feed_name=f'{self.subset}_{self.indexes[self.subset]}' if not is_test else None,\n        padding_values=FLAGS.padding_idx, \n        distribute_strategy=distribute_strategy,\n        torch=FLAGS.torch,\n        keras=FLAGS.keras,\n        subset=self.subset,\n        return_numpy=return_numpy,\n        ) \n      \n    result = self.adjust(result)\n    return result\n    \n  @staticmethod\n  def num_examples_per_epoch(subset, dir=None):\n    if subset == 'train':\n      num_examples = get_num_records(FLAGS.train_files)\n    elif subset == 'valid':\n      num_examples = get_num_records(FLAGS.valid_files)\n    else:\n      raise ValueError('Invalid data subset \"%s\"' % subset)\n    \n    assert num_examples\n    return num_examples\n\n  @staticmethod\n  def num_examples(subset, dir=None):\n    return Dataset.num_examples_per_epoch(subset, dir)\n\n  @property\n  def num_instances(self):\n    if self.num_instances_:\n      return self.num_instances_\n    assert self.files_\n    self.num_instances_ = get_num_records(self.files_, recount=self.recount)\n    return self.num_instances_\n\n  @property\n  def files(self):\n    return self.files_\n\n  @property\n  def records(self):\n    return self.files_\n\n  def __len__(self):\n    return self.num_instances or Dataset.num_examples_per_epoch(self.subset)\n\n  @property\n  def num_steps(self):\n    return -(-len(self) // self.batch_size)\n","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:18.319266Z","iopub.execute_input":"2025-10-08T10:11:18.3196Z","iopub.status.idle":"2025-10-08T10:11:18.361761Z","execution_failed":"2025-10-08T10:57:47.664Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Aug and preprocess","metadata":{}},{"cell_type":"code","source":"def apply(func, x, p=0.5):\n  if p <= 0:\n    return x\n  \"\"\"Randomly apply function func to x with probability p.\"\"\"\n  return tf.cond(\n      tf.less(tf.random.uniform([], minval=0, maxval=1, dtype=tf.float32),\n              tf.cast(p, tf.float32)),\n      lambda: func(x),\n      lambda: x)\n\nclass Apply(object):\n  def __init__(self, func, p):\n    self.func = func\n    self.p = p\n  \n  def __call__(self, x):\n    return apply(self.func, x, self.p)\n\ndef oneof(funcs, x, p=1.):\n  if p <= 0:\n    return x\n  num_choices = len(funcs)\n  prob = tf.random.uniform([], minval=0, maxval=1, dtype=tf.float32)\n  for i in range(num_choices):\n    if prob <= float(i + 1) * p / float(num_choices):\n      x = funcs[i](x)\n  return x\n\nclass OneOf(object):\n  def __init__(self, funcs, p=1.):\n    self.funcs = funcs\n    self.p = p\n\n  def __call__(self, x):\n    return oneof(self.funcs, x, self.p)\n\ndef 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 try_reshape(data):\n  if tf.shape(data)[-1] == N_COLS:\n    data = reshape(data)\n  return data\n\ndef splits(data):\n  l = [\n      N_HAND_POINTS, N_HAND_POINTS, \n      N_POSE_POINTS, N_POSE_POINTS, \n      N_LIP_POINTS, N_LIP_POINTS,\n      N_EYE_POINTS, N_EYE_POINTS,\n      N_NOSE_POINTS, N_NOSE_POINTS,\n      N_MID_POINTS, \n  ]\n  return tf.split(data, l, axis=-2)\n\ndef split(data):\n  x, y, z = tf.split(data, 3, axis=-1)\n  return x, y, z\n\ndef concat(x, y, z):\n  data = tf.concat([x, y, z], axis=-1)\n  return data\n\ndef flip_lr(data):\n  x, y, z = split(data)\n  x = 1. - x\n  data = concat(x, y, z)\n  # lhand, rhand, lpose, rpose, llip, rlip = splits(data)\n  # data = tf.concat([rhand, lhand, rpose, lpose, rlip, llip], axis=-2)\n  lhand, rhand, lpose, rpose, llip, rlip, leye, reye, lnose, rnose, mid = splits(data)\n  data = tf.concat([rhand, lhand, rpose, lpose, rlip, llip, reye, leye, rnose, lnose, mid], axis=-2)\n  return data\n\ndef zero_mean(x, axis=0, keepdims=False):\n  upper = tf.reduce_sum(x, axis=axis, keepdims=keepdims)\n  bottom = tf.reduce_sum(tf.where(tf.math.equal(x, 0.), tf.zeros_like(x), tf.ones_like(x)), axis=axis, keepdims=keepdims)\n  return upper / bottom\n\ndef zero_std(x, center=None, axis=0, keepdims=False):\n  if center is None:\n    center = zero_mean(x, axis=axis, keepdims=True)\n  d = x - center\n  return tf.math.sqrt(zero_mean(d * d, axis=axis, keepdims=keepdims))\n\ndef interp1d(x, target_len, method='random'):\n  target_len = tf.maximum(1, target_len)\n  if method == 'random':\n    if tf.random.uniform(()) < 0.33:\n      x = tf.image.resize(x, (1, target_len), 'bilinear')\n    else:\n      if tf.random.uniform(()) < 0.5:\n        x = tf.image.resize(x, (1, target_len), 'bicubic')\n      else:\n        x = tf.image.resize(x, (1, target_len), 'nearest')\n  else:\n    x = tf.image.resize(x, (1, target_len), method)\n  x = tf.reshape(x, (1, target_len, -1))\n  return x\n\n# x [None, n_frames, n_cols]\ndef resample(x, rate=(0.5, 1.5)):\n  rate = tf.random.uniform((), rate[0], rate[1])\n  length = tf.shape(x)[1]\n  new_size = tf.cast(rate * tf.cast(length, tf.float32), tf.int32)\n  new_x = interp1d(x, new_size)\n  return new_x\n\n# x [None, n_frames, n_cols]\ndef temporal_mask(x):\n  mask_prob = FLAGS.temporal_mask_prob if not FLAGS.temporal_mask_range else tf.random.uniform((), *FLAGS.temporal_mask_range)\n  mask = tf.random.uniform((tf.shape(x)[0], tf.shape(x)[1], 1))  > mask_prob\n  mask = tf.cast(mask, dtype=x.dtype)\n  new_x = x * mask\n  return new_x\n\n# x [1, n_frames, n_cols//3, 3]\ndef temporal_seq_mask(x):\n  x = x[0]\n  mask_value = float('nan')\n  l = tf.shape(x)[0]\n  mask_size = tf.random.uniform((), *FLAGS.temporal_seq_mask_range)\n  mask_size = tf.cast(tf.cast(l, tf.float32) * mask_size, tf.int32)\n  mask_offset = tf.random.uniform((), 0, tf.clip_by_value(l-mask_size,1,l), dtype=tf.int32)\n  new_x = tf.tensor_scatter_nd_update(x, \n                                      tf.range(mask_offset, mask_offset+mask_size)[...,None],\n                                      tf.fill([mask_size, tf.shape(x)[-2], tf.shape(x)[-1]],\n                                      mask_value))\n  new_x = new_x[None]\n  return new_x\n\n\n# x [None, n_frames, n_cols]\ndef spatio_mask(x):\n  mask = tf.random.uniform((tf.shape(x)[0], 1, tf.shape(x)[2]))  > FLAGS.spatio_mask_prob\n  mask = tf.cast(mask, dtype=x.dtype)\n  new_x = x * mask\n  return new_x\n\ndef shift(data):\n  range_ = FLAGS.shift_range\n  if FLAGS.shift_method == 0:\n    shift_ = tf.random.uniform((),*range_)\n    data += shift_\n  elif FLAGS.shift_method == 1:\n    shift_x = tf.random.uniform((),*range_)\n    shift_y = tf.random.uniform((),*range_)\n    shift_z = tf.random.uniform((),*range_)\n    x, y, z = split(data)\n    x += shift_x\n    y += shift_y\n    z += shift_z\n    data = concat(x, y, z)\n  else:\n    raise ValueError('not supported shift method')\n  return data\n\ndef scale(data):\n  scale_ = tf.random.uniform((), *FLAGS.scale_range)\n  if FLAGS.scale_method == 0:\n    data *= scale_\n  elif FLAGS.scale_method == 1:\n    data -= 0.5\n    data *= scale_\n    data += 0.5\n  return data\n\ndef rotate(data):\n  data -= 0.5\n  degree = tf.random.uniform((), *FLAGS.rotate_range)\n  radian = degree / 180 * np.pi\n  c = tf.math.cos(radian)\n  s = tf.math.sin(radian)\n  rotate_mat = tf.identity([\n            [c,s],\n            [-s, c],\n        ])\n  \n  x, y, z = split(data)\n  xy = tf.concat([x, y], axis=-1)\n  # [None, 88, 2]\n  # tf.print(xy.shape)\n  xy = xy @ rotate_mat\n  \n  data = tf.concat([xy, z], axis=-1)  \n  data += 0.5\n  return data\n  \ndef norm_frame(data):  \n  x, y, z = split(data)\n  x_mean = zero_mean(x, axis=-1, keepdims=True)\n  y_mean = zero_mean(y, axis=-1, keepdims=True)\n  z_mean = zero_mean(z, axis=-1, keepdims=True)\n  x_std = zero_std(x, center=x_mean, axis=-1, keepdims=True)\n  y_std = zero_std(y, center=y_mean, axis=-1, keepdims=True)\n  z_std = zero_std(z, center=z_mean, axis=-1, keepdims=True)\n\n  x = x - x_mean / x_std\n  y = y - y_mean / y_std\n  z = z - z_mean / z_std\n  data = concat(x, y, z)\n  return data\n\n# [1, n_frames, n_cols//3, 3]\ndef norm_hands(data):  \n  # lhand, rhand, lpose, rpose, llip, rlip = splits(data)\n  lhand, rhand, lpose, rpose, llip, rlip, leye, reye, lnose, rnose, mid = splits(data)\n  lhand = lhand[:,:,1:] - lhand[:,:,0:1]\n  rhand = rhand[:,:,1:] - rhand[:,:,0:1]\n  data = tf.concat([lhand, rhand], axis=2)\n  if FLAGS.norm_hands_size:\n    max_val = tf.math.maximum(tf.reduce_max(tf.abs(data)), 0.0001)\n    data /= max_val\n  return data\n\ndef add_motion(data):\n  lhand, rhand, lpose, rpose, llip, rlip, leye, reye, lnose, rnose, mid = splits(data)\n  hand = tf.concat([lhand, rhand], axis=2)\n  hand_motion = tf.concat([tf.zeros_like(hand[:, :1]), hand[:, 1:] - hand[:, :-1]], axis=1)\n  hand_motion2 = tf.concat([tf.zeros_like(hand[:, :2]), hand[:, 2:] - hand[:, :-2]], axis=1)\n  return tf.concat([hand_motion, hand_motion2], axis=2)\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","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:18.362803Z","iopub.execute_input":"2025-10-08T10:11:18.363113Z","iopub.status.idle":"2025-10-08T10:11:18.394336Z","execution_failed":"2025-10-08T10:57:47.664Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:18.395422Z","iopub.execute_input":"2025-10-08T10:11:18.396213Z","iopub.status.idle":"2025-10-08T10:11:18.490322Z","execution_failed":"2025-10-08T10:57:47.665Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class TfDataset(TfrecordsDataset):\n  def __init__(self, subset='valid', **kwargs):\n    super().__init__(subset, **kwargs)\n    assert not FLAGS.batch_parse, 'PreprocessLayer will be used as FLAGS.batch_parse = False'\n    self.prepocess = PreProcssor(subset)\n\n  def parse(self, example):\n    dynamic_keys = ['frames']\n    self.auto_parse(exclude_keys=dynamic_keys)\n    self.adds(dynamic_keys)\n    fe = self.parse_(serialized=example)\n    return self.prepocess(fe)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:18.491121Z","iopub.execute_input":"2025-10-08T10:11:18.491337Z","iopub.status.idle":"2025-10-08T10:11:18.49613Z","execution_failed":"2025-10-08T10:57:47.665Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:18.497121Z","iopub.execute_input":"2025-10-08T10:11:18.497347Z","iopub.status.idle":"2025-10-08T10:11:18.517449Z","execution_failed":"2025-10-08T10:57:47.666Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Lets display the first example of dataset","metadata":{}},{"cell_type":"markdown","source":"Wrap tfrecord dataset with a torch iterable dataset","metadata":{}},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:18.518316Z","iopub.execute_input":"2025-10-08T10:11:18.518592Z","iopub.status.idle":"2025-10-08T10:11:18.538543Z","shell.execute_reply.started":"2025-10-08T10:11:18.518574Z","shell.execute_reply":"2025-10-08T10:11:18.537415Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:18.539608Z","iopub.execute_input":"2025-10-08T10:11:18.539869Z","iopub.status.idle":"2025-10-08T10:11:18.559601Z","shell.execute_reply.started":"2025-10-08T10:11:18.539842Z","shell.execute_reply":"2025-10-08T10:11:18.558647Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:18.56082Z","iopub.execute_input":"2025-10-08T10:11:18.561221Z","iopub.status.idle":"2025-10-08T10:11:18.577291Z","shell.execute_reply.started":"2025-10-08T10:11:18.561194Z","shell.execute_reply":"2025-10-08T10:11:18.576344Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:18.578211Z","iopub.execute_input":"2025-10-08T10:11:18.578634Z","iopub.status.idle":"2025-10-08T10:11:18.60146Z","shell.execute_reply.started":"2025-10-08T10:11:18.578609Z","shell.execute_reply":"2025-10-08T10:11:18.600156Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:18.60253Z","iopub.execute_input":"2025-10-08T10:11:18.60314Z","iopub.status.idle":"2025-10-08T10:11:18.61811Z","shell.execute_reply.started":"2025-10-08T10:11:18.603114Z","shell.execute_reply":"2025-10-08T10:11:18.617072Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:18.618946Z","iopub.execute_input":"2025-10-08T10:11:18.61923Z","iopub.status.idle":"2025-10-08T10:11:18.632233Z","shell.execute_reply.started":"2025-10-08T10:11:18.61921Z","shell.execute_reply":"2025-10-08T10:11:18.631553Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:18.632914Z","iopub.execute_input":"2025-10-08T10:11:18.633197Z","iopub.status.idle":"2025-10-08T10:11:18.651841Z","shell.execute_reply.started":"2025-10-08T10:11:18.63317Z","shell.execute_reply":"2025-10-08T10:11:18.651081Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:18.652662Z","iopub.execute_input":"2025-10-08T10:11:18.652946Z","iopub.status.idle":"2025-10-08T10:11:18.669639Z","shell.execute_reply.started":"2025-10-08T10:11:18.652922Z","shell.execute_reply":"2025-10-08T10:11:18.668784Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:18.670571Z","iopub.execute_input":"2025-10-08T10:11:18.670892Z","iopub.status.idle":"2025-10-08T10:11:18.688455Z","shell.execute_reply.started":"2025-10-08T10:11:18.670867Z","shell.execute_reply":"2025-10-08T10:11:18.687443Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:18.689471Z","iopub.execute_input":"2025-10-08T10:11:18.689889Z","iopub.status.idle":"2025-10-08T10:11:18.714576Z","shell.execute_reply.started":"2025-10-08T10:11:18.689822Z","shell.execute_reply":"2025-10-08T10:11:18.713555Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:18.715922Z","iopub.execute_input":"2025-10-08T10:11:18.716448Z","iopub.status.idle":"2025-10-08T10:11:18.759398Z","shell.execute_reply.started":"2025-10-08T10:11:18.716396Z","shell.execute_reply":"2025-10-08T10:11:18.758402Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:18.76055Z","iopub.execute_input":"2025-10-08T10:11:18.761728Z","iopub.status.idle":"2025-10-08T10:11:18.781601Z","shell.execute_reply.started":"2025-10-08T10:11:18.761698Z","shell.execute_reply":"2025-10-08T10:11:18.780814Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:18.786194Z","iopub.execute_input":"2025-10-08T10:11:18.786468Z","iopub.status.idle":"2025-10-08T10:11:18.801647Z","shell.execute_reply.started":"2025-10-08T10:11:18.786449Z","shell.execute_reply":"2025-10-08T10:11:18.800867Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:18.802513Z","iopub.execute_input":"2025-10-08T10:11:18.802819Z","iopub.status.idle":"2025-10-08T10:11:18.824402Z","shell.execute_reply.started":"2025-10-08T10:11:18.802789Z","shell.execute_reply":"2025-10-08T10:11:18.823465Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model = Model()\nmodel","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:18.825308Z","iopub.execute_input":"2025-10-08T10:11:18.825621Z","iopub.status.idle":"2025-10-08T10:11:19.318637Z","shell.execute_reply.started":"2025-10-08T10:11:18.825595Z","shell.execute_reply":"2025-10-08T10:11:19.317827Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:19.319377Z","iopub.execute_input":"2025-10-08T10:11:19.319667Z","iopub.status.idle":"2025-10-08T10:11:20.600667Z","shell.execute_reply.started":"2025-10-08T10:11:19.31965Z","shell.execute_reply":"2025-10-08T10:11:20.599936Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class AWP():\n\n  def __init__(self,\n               model,\n               start_epoch=1,\n               param_name=\"weight\",\n               lr=1.,\n               eps=0.001):\n    self.model = model\n    self.param_name = param_name\n    self.lr = 1 if lr is None else lr\n    self.eps = 0.001 if eps is None else eps\n    self.start_epoch = start_epoch\n    ic(self.start_epoch, self.lr, self.eps)\n    self.backup = {}\n    self.backup_eps = {}\n\n  def will_skip(self, epoch):\n    return (self.lr == 0) or (epoch < self.start_epoch)\n\n  def on_retrain_begin(self, epoch=0):\n    if self.will_skip(epoch):\n      return False\n\n    self.save()\n    self.attack()\n    return True\n\n  def on_retrain_end(self, epoch=0):\n    if self.will_skip(epoch):\n      return False\n\n    self.restore()\n    return True\n\n  def attack(self):\n    e = 1e-6\n    for name, param in self.model.named_parameters():\n      if param.requires_grad and param.grad is not None and self.param_name in name:\n        norm1 = torch.norm(param.grad)\n        norm2 = torch.norm(param.data.detach())\n        if norm1 != 0 and not torch.isnan(norm1):\n          r_at = self.lr * param.grad / (norm1 + e) * (norm2 + e)\n          param.data.add_(r_at)\n          param.data = torch.min(\n              torch.max(param.data, self.backup_eps[name][0]),\n              self.backup_eps[name][1])\n        # param.data.clamp_(*self.backup_eps[name])\n\n  def save(self):\n    for name, param in self.model.named_parameters():\n      if param.requires_grad and param.grad is not None and self.param_name in name:\n        if name not in self.backup:\n          self.backup[name] = param.data.clone()\n          grad_eps = self.eps * param.abs().detach()\n          self.backup_eps[name] = (\n              self.backup[name] - grad_eps,\n              self.backup[name] + grad_eps,\n          )\n\n  def restore(self,):\n    for name, param in self.model.named_parameters():\n      if name in self.backup:\n        param.data = self.backup[name]\n    self.backup = {}\n    self.backup_eps = {}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:20.601366Z","iopub.execute_input":"2025-10-08T10:11:20.601612Z","iopub.status.idle":"2025-10-08T10:11:20.611444Z","shell.execute_reply.started":"2025-10-08T10:11:20.601596Z","shell.execute_reply":"2025-10-08T10:11:20.610753Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"FLAGS.optimizer = 'Adam'\nFLAGS.lr = 2e-3\nFLAGS.opt_eps = 1e-6\nFLAGS.warmup_proportion = 0.1\nFLAGS.epochs = 400\nFLAGS.pad_thre = -.3","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:20.612218Z","iopub.execute_input":"2025-10-08T10:11:20.612626Z","iopub.status.idle":"2025-10-08T10:11:20.633667Z","shell.execute_reply.started":"2025-10-08T10:11:20.612608Z","shell.execute_reply":"2025-10-08T10:11:20.632804Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"Optimizer = getattr(torch.optim, FLAGS.optimizer)\nkwargs = {\n    'lr': FLAGS.lr,\n    'eps': FLAGS.opt_eps,\n}\noptimizer = Optimizer(model.parameters(), **kwargs)\noptimizer","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:20.634692Z","iopub.execute_input":"2025-10-08T10:11:20.634979Z","iopub.status.idle":"2025-10-08T10:11:23.524052Z","shell.execute_reply.started":"2025-10-08T10:11:20.634954Z","shell.execute_reply":"2025-10-08T10:11:23.523251Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"records_pattern = f'../input/3rd-place-step1-gen-tfrecords-for-train/tfrecords/train/*.tfrec'\nfiles = glob.glob(records_pattern) \n\nif FLAGS.online:\n  FLAGS.train_files = files\nelse:\n  FLAGS.train_files = [x for x in files if int(os.path.basename(x).split('.')[0]) % FLAGS.folds != FLAGS.fold]\n\nFLAGS.valid_files = [x for x in files if int(os.path.basename(x).split('.')[0]) % FLAGS.folds == FLAGS.fold]\n\nif FLAGS.mix_sup:\n  FLAGS.train_files += glob.glob(f'../input/3rd-place-step2-gen-tfrecords-for-supplement/tfrecords/sup/*.tfrec')\n  np.random.shuffle(FLAGS.train_files)\nic(FLAGS.train_files[:3], FLAGS.valid_files[:2])\n\ntrain_dl, eval_dl = get_dataloaders()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:23.524907Z","iopub.execute_input":"2025-10-08T10:11:23.525587Z","iopub.status.idle":"2025-10-08T10:11:29.003267Z","shell.execute_reply.started":"2025-10-08T10:11:23.525566Z","shell.execute_reply":"2025-10-08T10:11:29.002626Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"num_steps_per_epoch = train_dl.dataset.num_steps\ntotal_steps = num_steps_per_epoch * FLAGS.epochs","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:29.004009Z","iopub.execute_input":"2025-10-08T10:11:29.004206Z","iopub.status.idle":"2025-10-08T10:11:29.008287Z","shell.execute_reply.started":"2025-10-08T10:11:29.004191Z","shell.execute_reply":"2025-10-08T10:11:29.007373Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def get_scheduler_infos():\n  num_train_steps = int(num_steps_per_epoch * FLAGS.epochs)\n  num_warmup_steps = int(num_train_steps * FLAGS.warmup_proportion)\n  ic(num_warmup_steps)\n  return num_train_steps, num_warmup_steps","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:29.00927Z","iopub.execute_input":"2025-10-08T10:11:29.009848Z","iopub.status.idle":"2025-10-08T10:11:29.028705Z","shell.execute_reply.started":"2025-10-08T10:11:29.00983Z","shell.execute_reply":"2025-10-08T10:11:29.027962Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\ndevice","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:29.029479Z","iopub.execute_input":"2025-10-08T10:11:29.029764Z","iopub.status.idle":"2025-10-08T10:11:29.046463Z","shell.execute_reply.started":"2025-10-08T10:11:29.02974Z","shell.execute_reply":"2025-10-08T10:11:29.045776Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def to_device(x, y):\n  for key in x:\n    x[key] = x[key].to(device, non_blocking=True)\n  y = y.to(device, non_blocking=True)\n  return x, y","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:29.047237Z","iopub.execute_input":"2025-10-08T10:11:29.047487Z","iopub.status.idle":"2025-10-08T10:11:29.061065Z","shell.execute_reply.started":"2025-10-08T10:11:29.047457Z","shell.execute_reply":"2025-10-08T10:11:29.060266Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# for eval","metadata":{}},{"cell_type":"code","source":"def 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":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:29.061859Z","iopub.execute_input":"2025-10-08T10:11:29.062081Z","iopub.status.idle":"2025-10-08T10:11:29.073209Z","shell.execute_reply.started":"2025-10-08T10:11:29.062066Z","shell.execute_reply":"2025-10-08T10:11:29.072485Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install python-Levenshtein","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:29.07395Z","iopub.execute_input":"2025-10-08T10:11:29.07421Z","iopub.status.idle":"2025-10-08T10:11:34.90804Z","shell.execute_reply.started":"2025-10-08T10:11:29.074193Z","shell.execute_reply":"2025-10-08T10:11:34.907098Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from Levenshtein import distance\nfrom sklearn.metrics import log_loss\n\ndef is_normal_char(idx):\n  return not IDX2CHAR[idx].startswith('<')\n\ndef to_original_phrase(phrase_pred):\n  phrase = simple_decode(phrase_pred).numpy()\n  return phrase\n\ndef to_str(phrase):\n  phrase = ''.join([IDX2CHAR[idx] for idx in phrase if is_normal_char(idx)])\n  return phrase\n\ndef max_char_idx(phrase):\n  idxs = [i + 1 for i, idx in enumerate(phrase) if is_normal_char(idx)]\n  return max(idxs) if idxs else 0\n\ndef calc_score(phrase_true, phrase_pred):\n  return distance(phrase_true, phrase_pred)\n\n\ndef evaluate(dataset, model, eval_step, steps, step, is_last, num_examples, loss_fn, outdir):\n  return eval_seq(dataset, model, steps)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:34.909409Z","iopub.execute_input":"2025-10-08T10:11:34.910141Z","iopub.status.idle":"2025-10-08T10:11:35.031389Z","shell.execute_reply.started":"2025-10-08T10:11:34.910104Z","shell.execute_reply":"2025-10-08T10:11:35.030768Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def try_numpy(x):\n  if hasattr(x, 'numpy'):\n    return x.numpy()\n  return x\n\ndef eval_seq(dataset, model, steps):\n  metrics = {}\n  l = []\n  outputs = []\n  feats = []\n      \n  t = tqdm(enumerate(dataset), total=steps, ascii=True, desc= 'eval_loop')\n  for step_, (x, y) in t:\n    if step_ == steps:\n      break\n    \n    x['frames'] = x['frames'].cuda()\n    frames = x['frames']\n\n    with torch.no_grad():\n      y_ = model.infer(frames)\n    \n    for key in x:\n      if hasattr(x[key], 'cpu'):\n        x[key] = x[key].cpu().detach()\n        \n    y_ = y_.cpu().detach().numpy()        \n    outputs.append(y_)\n\n    y = y.numpy()\n    sequence_ids = try_numpy(x['sequence_id'])\n    phrase_types = try_numpy(x['phrase_type'])\n    phrase_dups = try_numpy(x['phrase_dup'])\n    n_frames = try_numpy(x['n_frames'])\n    frame_means = try_numpy(x['frame_mean'])\n    idxes = try_numpy(x['idx'])\n    cls_labels = try_numpy(x['cls_label'])\n    for i, (sequence_id, phrase_true, phrase_pred, phrase_type, phrase_dup, n_frame, frame_mean, cls_label, idx) \\\n      in enumerate(zip(sequence_ids, y, y_, phrase_types, phrase_dups, n_frames, frame_means, cls_labels, idxes)):  \n      phrase_ori = to_original_phrase(phrase_pred)\n      char_max_idx = max_char_idx(phrase_ori)\n      char_ori_rate = char_max_idx  / len(phrase_ori)\n      phrase_ori = ''.join([IDX2CHAR[idx] for idx in phrase_ori])\n      phrase_pred = decode_phrase(phrase_pred).numpy()\n      phrase_true = to_str(phrase_true)\n      phrase_pred = to_str(phrase_pred)\n      # try:\n      char_true_rate = 0. if not len(phrase_true) else char_max_idx / len(phrase_true)\n      char_pred_rate = 0. if not len(phrase_pred) else char_max_idx / len(phrase_pred)\n\n      cls_pred = np.zeros_like(cls_label)\n      for ch in phrase_pred:\n        cidx = CHAR2IDX[ch] - 1\n        cls_pred[cidx] = 1\n        \n      phrase_type = phrase_type.decode('utf-8')\n      phrase_type_pred = get_phrase_type(phrase_pred)\n      distance = calc_score(phrase_true, phrase_pred)\n      m = {\n        'sequence_id': sequence_id,\n        'phrase_type': phrase_type,\n        'phrase_dup': phrase_dup,\n        'phrase_true': phrase_true,\n        'phrase_pred': phrase_pred,\n        'phrase_ori': phrase_ori,\n        'char/max_idx': char_max_idx,\n        'char/ori_rate': char_ori_rate,\n        'char/true_rate': char_true_rate,\n        'char/pred_rate': char_pred_rate,\n        'phrase_len_true': len(phrase_true),\n        'phrase_len_pred': len(phrase_pred),\n        'phrase_len_rate': len(phrase_pred) / len(phrase_true),\n        'idx': idx,\n        'distance': distance,\n        'acc/char': (cls_label == cls_pred).mean(),\n        'acc/type': phrase_type == phrase_type_pred,\n        'acc/first': phrase_true[0] == phrase_pred[0] if phrase_pred else False,\n        'acc/last': phrase_true[-1] == phrase_pred[-1] if phrase_pred else False,\n        'n_frame': n_frame,\n        'frame_mean': frame_mean,\n        'score': max(len(phrase_true) - distance, 0.) / len(phrase_true),\n      }\n      l.append(m)\n      \n  outputs = np.concatenate(outputs, axis=0)\n    \n  df = pd.DataFrame(l)\n  assert len(df) == len(df.idx.unique())\n  def get_metrics(df):\n    metrics = {\n      'phrase_len_true': df['phrase_len_true'].mean(),\n      'phrase_len_pred': df['phrase_len_pred'].mean(),\n      'phrase_len_rate': df['phrase_len_rate'].mean(),\n      'len/l1': (df['phrase_len_true'] - df['phrase_len_pred']).abs().mean(),\n      'char/max_idx_max': df['char/max_idx'].max(),\n      'char/max_idx': df['char/max_idx'].mean(),\n      'char/ori_rate': df['char/ori_rate'].mean(),\n      'char/true_rate': df['char/true_rate'].mean(),\n      'char/pred_rate': df['char/pred_rate'].mean(),\n      'distance': df['distance'].mean(),\n      'acc/char': df['acc/char'].mean(),\n      'acc/type': df['acc/type'].mean(),\n      'acc/first': df['acc/first'].mean(),\n      'acc/last': df['acc/last'].mean(),\n      'score': max(df['phrase_len_true'].sum() - df['distance'].sum(), 0.) / df['phrase_len_true'].sum(),\n    }\n    return metrics\n  \n  metrics = get_metrics(df)\n  \n  df.to_csv(f'{FLAGS.working}/eval.csv', index=False)\n  df_ = df\n  df = df[df.idx < 20]\n  df = df.sort_values(by=['idx'], ascending=True)\n  df2 = df[['sequence_id', 'phrase_type', 'phrase_true', 'phrase_pred', 'phrase_ori', 'phrase_len_true', 'phrase_len_pred', 'distance', 'score']].head(20)\n  # display(df2)\n  # ic(metrics['score'])\n  return metrics\n\nTRAIN_LOSSES = []      # Training loss per step\nEVAL_LOSSES = []       # Validation loss proxy per epoch\nEVAL_SCORES = []       # Validation score per epoch\nLRS = []               # Learning rate per step\nGRAD_NORMS = []        # Gradient norm per accumulation step","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:35.032274Z","iopub.execute_input":"2025-10-08T10:11:35.032583Z","iopub.status.idle":"2025-10-08T10:11:35.048915Z","shell.execute_reply.started":"2025-10-08T10:11:35.032559Z","shell.execute_reply":"2025-10-08T10:11:35.047909Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Training \nAs P100 slow here I change to run less epochs, where on my local train using 1 4090, it will train for 400 epochs, and will finetune with train data only for addtional 10 epochs    \nFor simplify I commented out awp part here  \nyou may change FLAGS.epochs to 400 and also you may set FLAGS.online = True which means using all folds for train  ","metadata":{}},{"cell_type":"code","source":"# def train_loop():\n#   global_step = 0\n#   pbar = tqdm(range(FLAGS.epochs), total=FLAGS.epochs, desc='train_epochs')\n#   ds = iter(train_dl)\n#   for epoch in pbar:\n#     pbar2 = tqdm(range(num_steps_per_epoch), total=num_steps_per_epoch, desc=f'train_epoch:{epoch}')\n#     for _ in pbar2:\n#       x, y = next(ds)\n#       x, y = to_device(x, y)\n      \n#       def loss_fn_(x, y):\n#         y_ = model(x)\n#         loss = loss_fn(y_, y, x)\n#         return loss\n    \n#       loss = loss_fn_(x, y)\n#       loss = loss / FLAGS.acc_steps  \n#       loss.backward()\n\n#   #     awp.on_retrain_begin(epoch)\n#   #     loss = loss_fn_(x, y)\n#   #     loss = loss / FLAGS.acc_steps\n#   #     loss.backward()\n#   #     awp.on_retrain_end(epoch)\n\n#       is_acc_step = ((global_step + 1) % FLAGS.acc_steps == 0) or (global_step + 1 == total_steps)\n#       if is_acc_step:\n#         optimizer.step()\n#         model.zero_grad()\n              \n#       scheduler.step()\n#       global_step += 1\n#       pbar2.set_postfix({\n#           'loss': loss.item()\n#       })\n        \n#     metrics = eval_seq(eval_dl, model, eval_dl.dataset.num_steps)\n#     pbar2.set_postfix({\n#           'loss': loss.item(),\n#           'score': metrics['score'],\n#     })\n    \n#     pbar.set_postfix({\n#         'loss': loss.item(),\n#         'score': metrics['score'],\n#     })\ndef train_loop(model, optimizer, scheduler, train_dl, eval_dl, num_steps_per_epoch, total_steps, num_epochs, loss_fn):\n    global_step = 0\n    ds = iter(train_dl)\n    pbar = tqdm(range(num_epochs), total=num_epochs, desc='train_epochs')\n\n    for epoch in pbar:\n        pbar2 = tqdm(range(num_steps_per_epoch), total=num_steps_per_epoch, desc=f'train_epoch:{epoch}')\n\n        for _ in pbar2:\n            try:\n                x, y = next(ds)\n            except StopIteration:\n                ds = iter(train_dl)\n                x, y = next(ds)\n\n            x, y = to_device(x, y)\n\n            def loss_fn_(x, y):\n                y_ = model(x)\n                loss = loss_fn(y_, y, x)\n                return loss\n\n            loss = loss_fn_(x, y)\n            loss = loss / FLAGS.acc_steps  \n            loss.backward()\n\n            is_acc_step = ((global_step + 1) % FLAGS.acc_steps == 0) or (global_step + 1 == total_steps)\n            if is_acc_step:\n                # gradient norm\n                total_norm = 0\n                for p in model.parameters():\n                    if p.grad is not None:\n                        param_norm = p.grad.data.norm(2)\n                        total_norm += param_norm.item() ** 2\n                total_norm = total_norm ** 0.5\n                GRAD_NORMS.append(total_norm)\n\n                optimizer.step()\n                scheduler.step()   # must be after optimizer.step()\n                model.zero_grad()\n\n            LRS.append(scheduler.get_last_lr()[0])\n            global_step += 1\n            TRAIN_LOSSES.append(loss.item() * FLAGS.acc_steps)\n\n            pbar2.set_postfix({'loss': loss.item() * FLAGS.acc_steps})\n\n        metrics = eval_seq(eval_dl, model, eval_dl.dataset.num_steps)\n\n        eval_loss_val = metrics.get('distance', 0)  # distance as proxy for \"loss\"\n        eval_score_val = metrics['score']\n\n        EVAL_LOSSES.append(eval_loss_val)\n        EVAL_SCORES.append(eval_score_val)\n\n        # save per-epoch CSV for later analysis\n        pd.DataFrame([metrics]).to_csv(f\"eval_epoch_{epoch+1}.csv\", index=False)\n\n        print(f\"\\n[Epoch {epoch+1}/{num_epochs}] Score: {eval_score_val:.4f}, Distance: {eval_loss_val:.4f}\")\n\n        pbar2.set_postfix({'loss': loss.item(), 'score': eval_score_val})\n        pbar.set_postfix({'loss': loss.item(), 'score': eval_score_val})\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-08T10:11:35.049982Z","iopub.execute_input":"2025-10-08T10:11:35.050245Z","iopub.status.idle":"2025-10-08T10:11:35.07112Z","shell.execute_reply.started":"2025-10-08T10:11:35.050228Z","shell.execute_reply":"2025-10-08T10:11:35.070173Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# # FLAGS.epochs = 2\n# FLAGS.epochs = 15\n# # FLAGS.epochs = 400  \n# FLAGS.lr = 2e-3\n# Optimizer = getattr(torch.optim, FLAGS.optimizer)\n# kwargs = {\n#     'lr': FLAGS.lr,\n#     'eps': FLAGS.opt_eps,\n# }\n# optimizer = Optimizer(model.parameters(), **kwargs)\n# num_steps_per_epoch = train_dl.dataset.num_steps\n# total_steps = num_steps_per_epoch * FLAGS.epochs\n# num_train_steps, num_warmup_steps = get_scheduler_infos()\n# scheduler = transformers.get_polynomial_decay_schedule_with_warmup(optimizer, \n#                                                                    num_warmup_steps=num_warmup_steps, \n#                                                                    num_training_steps=num_train_steps + 1)\n\nFLAGS.epochs = 30   # change to 50, 100, 150, 200 for different runs\nFLAGS.lr = 2e-3\n\nOptimizer = getattr(torch.optim, FLAGS.optimizer)\noptimizer = Optimizer(model.parameters(), lr=FLAGS.lr, eps=FLAGS.opt_eps)\n\nnum_steps_per_epoch = train_dl.dataset.num_steps\ntotal_steps = num_steps_per_epoch * FLAGS.epochs\nnum_train_steps, num_warmup_steps = get_scheduler_infos()\n\nscheduler = transformers.get_polynomial_decay_schedule_with_warmup(\n    optimizer, \n    num_warmup_steps=num_warmup_steps, \n    num_training_steps=num_train_steps + 1\n)\n\nloss_fn = model.get_loss_fn()\nmodel = model.to(device)\n\nprint(f\"--- STARTING MAIN TRAINING (Epochs={FLAGS.epochs}, LR={FLAGS.lr}) ---\")\ntrain_loop(model, optimizer, scheduler, train_dl, eval_dl, num_steps_per_epoch, total_steps, FLAGS.epochs, loss_fn)\nprint(\"--- MAIN TRAINING COMPLETE ---\")\n","metadata":{"trusted":true,"execution":{"execution_failed":"2025-10-08T10:57:47.671Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# model = model.to(device)\n# ## p100 not support torch.compile\n# # model = torch.compile(model)\n# loss_fn = model.get_loss_fn()\n# awp = AWP(model, start_epoch=int(FLAGS.epochs * 0.15), lr=0.2, eps=0)","metadata":{"trusted":true,"execution":{"execution_failed":"2025-10-08T10:57:47.667Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# train_loop()","metadata":{"trusted":true,"execution":{"execution_failed":"2025-10-08T10:57:47.667Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# finetune with train data only(excluding supplement data)","metadata":{}},{"cell_type":"code","source":"files = glob.glob(records_pattern) \n\nif FLAGS.online:\n  FLAGS.train_files = files\nelse:\n  FLAGS.train_files = [x for x in files if int(os.path.basename(x).split('.')[0]) % FLAGS.folds != FLAGS.fold]\n\nFLAGS.valid_files = [x for x in files if int(os.path.basename(x).split('.')[0]) % FLAGS.folds == FLAGS.fold]\n\nic(FLAGS.train_files[:3], FLAGS.valid_files[:2])\ntrain_dl, eval_dl = get_dataloaders()","metadata":{"trusted":true,"execution":{"execution_failed":"2025-10-08T10:57:47.668Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# # FLAGS.epochs = 2\n# FLAGS.epochs = 4\n# # FLAGS.epochs = 10\n# FLAGS.lr = 1e-4\n# Optimizer = getattr(torch.optim, FLAGS.optimizer)\n# kwargs = {\n#     'lr': FLAGS.lr,\n#     'eps': FLAGS.opt_eps,\n# }\n# optimizer = Optimizer(model.parameters(), **kwargs)\n# num_steps_per_epoch = train_dl.dataset.num_steps\n# total_steps = num_steps_per_epoch * FLAGS.epochs\n# num_train_steps, num_warmup_steps = get_scheduler_infos()\n# scheduler = transformers.get_polynomial_decay_schedule_with_warmup(optimizer, \n#                                                                    num_warmup_steps=num_warmup_steps, \n#                                                                    num_training_steps=num_train_steps + 1)\n\nFLAGS.epochs = 5\nFLAGS.lr = 1e-4\n\nOptimizer = getattr(torch.optim, FLAGS.optimizer)\noptimizer = Optimizer(model.parameters(), lr=FLAGS.lr, eps=FLAGS.opt_eps)\n\nnum_steps_per_epoch = train_dl.dataset.num_steps\ntotal_steps = num_steps_per_epoch * FLAGS.epochs\nnum_train_steps, num_warmup_steps = get_scheduler_infos()\n\nscheduler = transformers.get_polynomial_decay_schedule_with_warmup(\n    optimizer, \n    num_warmup_steps=num_warmup_steps, \n    num_training_steps=num_train_steps + 1\n)\n\nprint(f\"--- STARTING FINETUNING (Epochs={FLAGS.epochs}, LR={FLAGS.lr}) ---\")\ntrain_loop(model, optimizer, scheduler, train_dl, eval_dl, num_steps_per_epoch, total_steps, FLAGS.epochs, loss_fn)\nprint(\"--- FINETUNING COMPLETE ---\")\n\ntorch.save(model.state_dict(), './model.pt')\n","metadata":{"trusted":true,"execution":{"execution_failed":"2025-10-08T10:57:47.668Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# torch.save(model.state_dict(), './model.pt')","metadata":{"trusted":true,"execution":{"execution_failed":"2025-10-08T10:57:47.669Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Visualization","metadata":{}},{"cell_type":"code","source":"import matplotlib.pyplot as plt\n\ndef plot_training_metrics(train_losses, eval_losses, eval_scores, lrs, grad_norms):\n    num_eval_epochs = len(eval_scores)\n\n    # 1. Training vs Validation Loss\n    plt.figure(figsize=(15, 5))\n    plt.subplot(1, 2, 1)\n    plt.plot(train_losses, label='Training Loss (per step)', alpha=0.7)\n    plt.title('Training Loss vs Steps')\n    plt.xlabel('Training Steps'); plt.ylabel('Loss'); plt.grid(True)\n\n    plt.subplot(1, 2, 2)\n    plt.plot(range(1, num_eval_epochs+1), eval_losses, marker='o', color='red')\n    plt.title('Validation Distance vs Epochs')\n    plt.xlabel('Epochs'); plt.ylabel('Distance'); plt.grid(True)\n    plt.tight_layout(); plt.show()\n\n    # 2. Score Curve\n    plt.figure(figsize=(7,5))\n    plt.plot(range(1, num_eval_epochs+1), eval_scores, marker='o', color='green')\n    plt.title('Validation Score vs Epochs')\n    plt.xlabel('Epochs'); plt.ylabel('Score'); plt.grid(True)\n    plt.show()\n\n    # 3. LR and Gradient Norms\n    plt.figure(figsize=(15,5))\n    plt.plot(lrs, label='Learning Rate', color='purple', alpha=0.7)\n    plt.plot(np.linspace(0, len(lrs), len(grad_norms)), grad_norms, '--', color='orange', label='Grad Norm')\n    plt.title('Learning Rate & Gradient Norms')\n    plt.xlabel('Training Steps'); plt.legend(); plt.grid(True)\n    plt.show()\n\nplot_training_metrics(TRAIN_LOSSES, EVAL_LOSSES, EVAL_SCORES, LRS, GRAD_NORMS)","metadata":{"trusted":true,"execution":{"execution_failed":"2025-10-08T10:57:47.67Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def visualize_eval_csv(log_dir, show_worst=True, top_k=10):\n    \"\"\"\n    Visualize metrics from eval.csv (per run) and eval_epoch_X.csv (if available).\n    Args:\n        log_dir (str): Directory where eval.csv / eval_epoch_X.csv are stored.\n        show_worst (bool): Whether to print worst-case samples.\n        top_k (int): Number of worst-case samples to display.\n    \"\"\"\n    \n    # --- Case 1: Latest eval.csv (per-run detailed results) ---\n    try:\n        df = pd.read_csv(f\"{log_dir}/eval.csv\")\n        print(f\"[INFO] Loaded eval.csv with {len(df)} samples.\")\n\n        # (a) Score Distribution\n        plt.figure(figsize=(10,5))\n        df['score'].hist(bins=50, alpha=0.7)\n        plt.title(\"Distribution of Scores (Eval Set)\")\n        plt.xlabel(\"Score\")\n        plt.ylabel(\"Frequency\")\n        plt.grid(True, alpha=0.3)\n        plt.show()\n\n        # (b) Distance vs Phrase Length\n        plt.figure(figsize=(8,6))\n        plt.scatter(df['phrase_len_true'], df['distance'], alpha=0.5)\n        plt.xlabel(\"True Phrase Length\")\n        plt.ylabel(\"Edit Distance\")\n        plt.title(\"Edit Distance vs Phrase Length\")\n        plt.grid(True, alpha=0.3)\n        plt.show()\n\n        # (c) Score by Phrase Type\n        if \"phrase_type\" in df.columns:\n            acc_by_type = df.groupby(\"phrase_type\")[\"score\"].mean().sort_values()\n            plt.figure(figsize=(12,6))\n            acc_by_type.plot(kind=\"barh\")\n            plt.title(\"Average Score by Phrase Type\")\n            plt.xlabel(\"Score\")\n            plt.ylabel(\"Phrase Type\")\n            plt.show()\n\n        # (d) Worst-case Predictions\n        if show_worst:\n            print(f\"\\n--- {top_k} WORST CASES (Lowest Scores) ---\")\n            worst_cases = df.sort_values(\"score\").head(top_k)\n            display(worst_cases[[\"sequence_id\", \"phrase_true\", \"phrase_pred\", \"score\", \"distance\"]])\n\n    except Exception as e:\n        print(f\"[WARN] Could not load {log_dir}/eval.csv. Error: {e}\")\n\n    # --- Case 2: Per-epoch eval logs ---\n    eval_files = sorted(glob.glob(f\"{log_dir}/eval_epoch_*.csv\"))\n    if eval_files:\n        print(f\"[INFO] Found {len(eval_files)} epoch-level eval logs.\")\n        scores, epochs = [], []\n        for f in eval_files:\n            df_epoch = pd.read_csv(f)\n            scores.append(df_epoch['score'].mean())\n            epoch = int(f.split(\"_\")[-1].split(\".\")[0].replace(\".csv\", \"\"))\n            epochs.append(epoch)\n\n        plt.figure(figsize=(10,6))\n        plt.plot(epochs, scores, marker='o')\n        plt.xlabel(\"Epoch\")\n        plt.ylabel(\"Mean Validation Score\")\n        plt.title(\"Validation Score per Epoch\")\n        plt.grid(True, alpha=0.3)\n        plt.show()\n    else:\n        print(\"[INFO] No per-epoch eval logs found.\")\nvisualize_eval_csv(FLAGS.working, show_worst=True, top_k=10)","metadata":{"trusted":true,"execution":{"execution_failed":"2025-10-08T10:57:47.67Z"}},"outputs":[],"execution_count":null}]}