class_docstr: Wrapper that always initializes weights from a pre-trained checkpoint.
class_docstr_tok: ["Wrapper","that","always","initializes","weights","from","a","pre-trained","checkpoint","."]
docstr: Returns model dir path with model filename.
docstr_tok: ["Returns","model","dir","path","with","model","filename","."]
code: 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)
  
imports: import os
  import pathlib
  import urllib
  import zipfile
  import gin
  import jax
  import tensorflow
  from trax import fastmath
  from trax import layers
  from trax.fastmath import numpy
label_desc: Use image node_id 2 for calling the PretrainedBERT obj's underlying member method code with example usage: obj.download_model(cls, model_name) and returns: download_dir,  find_ckpt_file_in_dir, dir_path, download_dir,  find_ckpt_file_in_dir, unknown