def download_model(cls, model_name):
    """Returns model dir path with model filename."""
    try:
      model_link = _MODEL_LINKS[model_name]
    except KeyError:
      raise KeyError(
          f'Not known model name, please make sure the model name'
          f' is in the list of available models. If this is a path'
          f' to a model it should contain at least one {os.path.sep}')

    def find_ckpt_file_in_dir(dir_path):
      # TODO(piotrekp1): get last checkpoint instead of any
      for fname in os.listdir(dir_path):
        if '.ckpt' in fname:
          return fname.split('.ckpt')[0] + '.ckpt'
      raise FileNotFoundError('Selected directory doesn\'t contain a ckpt file')

    def cd_nested_directory(dir_path):
      while len(os.listdir(dir_path)) == 1:
        filename = os.listdir(dir_path)[0]
        new_path = os.path.join(dir_path, filename)
        if os.path.isdir(new_path):
          dir_path = new_path
      return dir_path

    download_dir = os.path.join('~', 'trax', model_name)
    download_dir = os.path.expanduser(download_dir)

    if os.path.exists(download_dir):
      if os.path.isdir(download_dir) and len(os.listdir(download_dir)):
        # model already exists, mdoel directory as input
        download_dir = cd_nested_directory(
            download_dir)  # go in if nested single directories inside
        return download_dir, find_ckpt_file_in_dir(download_dir)
    with pathlib.Path(download_dir) as p:
      p.mkdir(parents=True, exist_ok=True)

    if not model_link.endswith('.zip'):
      raise NotImplementedError(
          'Only downloading models packed with zip is implemented')

    # assumes model is packed with zip
    download_path = os.path.join(download_dir, 'model_temp')

    print(f'Downloading model from {model_link} to {download_dir}'
         )  # TODO(piotrekp1) logging to stderr?
    urllib.request.urlretrieve(model_link, download_path)

    if not zipfile.is_zipfile(download_path):
      raise NotImplementedError(
          'Only downloading models packed with zip is implemented')

    # TODO(piotrekp1): some locks to handle crashes during unpacking or download
    with zipfile.ZipFile(download_path, 'r') as zip_ref:
      zip_ref.extractall(download_dir)

    os.remove(download_path)
    download_dir = cd_nested_directory(download_dir)

    return download_dir, find_ckpt_file_in_dir(download_dir)
Returns model dir path with model filename.
trax/models/research/bert.py