Given a binary file of tokens in Megatron format, calculate its number of tokens. Expect each token occupies exactly 4 bytes.

calc_tokens_by_filesize.py

import re

def calc_num_token_of_bin(filesize: str) -> int:
  m = re.findall(r"(\d+)([G|T])", filesize)
  assert len(m) == 1, f"Expect string like '10G', '1T', but got {filesize}"
  m = m[0]
  digit, dim = int(m[0]), m[1]
  pow = 2 if dim == 'G' else 3
  return digit / 4 * (1024 / 1000) ** pow