text: import os
  import sys
  import numpy as np
  import tensorflow as tf
  import logging
  tf.get_logger().setLevel(logging.ERROR)
  from model import TrRefiner
  
  n_chunk_default = 5
  n_filter_default = 128
  def parse_input_args():
      import argparse
      opt = argparse.ArgumentParser\
              (description='''AngPred: Backbone torsion angle prediction using ResNet''')
      #
      # Training config
      opt.add_argument('-train_param', dest='train_param', metavar='TRAIN_PARAM', default='param_train.json',\
                           help="json file containing all the training related parameters (e.g. learning rate, etc)")
      #
      # Model config
      opt.add_argument('-num_chunks', '--num_chunks', dest='num_chunks', metavar='num_chunks', type=int, default=n_chunk_default, \
                           help='The number of ResNet chunks [%d]'%n_chunk_default)
      opt.add_argument('-n2d_filters', '--n2d_filters', dest='n2d_filters', metavar='n2d_filters', type=int, default=n_filter_default, \
                           help='The number of 2D filters [%d]'%n_filter_default)
      #
      # Flag for training or not
      opt.add_argument('-train', '--train', dest='train', default=False, action='store_true',\
                           help='for training [False]')
      opt.add_argument('-att', '--att', dest='att', default=False, action='store_true',\
                           help='use pixel-attention')
      #
      # Inputs for prediction mode
      opt.add_argument('-a3m_fn', dest='a3m_fn', default=None, \
                          help="MSA file in a3m format")
      opt.add_argument('-pdb_fn_s', dest='pdb_fn_s', default=[None], nargs='+',\
                          help="Input structures in pdb format")
      opt.add_argument('-npz_fn_s', dest='npz_fn_s', default=[None], nargs='+',\
                          help="Previous 6D coords distribution prediction results in npz format. It must contain dist, omega, theta, phi")
      opt.add_argument('-SS_fn_s', dest='SS_fn_s', default=[None], nargs='+',\
                          help="predicted SS in npz file from Emin")
      opt.add_argument('-acc_fn_s', dest='acc_fn_s', default=[None], nargs='+',\
                          help="predicted acc in npz file for all input pdbs")
      opt.add_argument('-out_fn', dest='out_fn', default=None,\
                          help="Output file name")
      opt.add_argument('-model_dir', dest='model_dir', default="%s/model"%(os.path.dirname(__file__)),\
                          help="Directory for pre-trained ML model checkpoint files")
      
      if len(sys.argv) == 1:
          opt.print_help()
          sys.exit()
      #
      FLAGS = opt.parse_args()
  
      return FLAGS
  
  def main():
      args = parse_input_args()
      #
      if not args.train: # prediction mode
          error = False
          if args.a3m_fn == None:
              sys.stderr.write("ERROR: Please provide MSA in a3m format\n")
              error = True
          if args.out_fn == None:
              sys.stderr.write("ERROR: Please provide output file name\n")
              error = True
          if args.pdb_fn_s[0] == None:
              sys.stderr.write("ERROR: Please provide structure model generated by trRosetta\n")
              error = True
          if args.npz_fn_s[0] == None:
              sys.stderr.write("ERROR: please provide 6d coords prediction results in npz format\n")
              error = True
          if not os.path.exists("%s_rep1"%args.model_dir):
              sys.stderr.write("ERROR: please check directory path for pre-trained checkpoint files, %s\n"%args.model_dir)
              error = True
          if error:
              return
  
      run_config = tf.ConfigProto(
          gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.92)
      )
      with tf.Session(config=run_config) as sess:
  
          if args.train:
              model = TrRefiner(sess, num_chunks=args.num_chunks, n2d_filters=args.n2d_filters, pdb_in=False, train=True, config_file=args.train_param, attention=args.att)
              model.train(args.train_param)
          else:
              model = TrRefiner(sess, num_chunks=args.num_chunks, n2d_filters=args.n2d_filters, pdb_in=True, attention=args.att)
              model.predict(args.model_dir, args.a3m_fn, args.npz_fn_s, args.pdb_fn_s, args.acc_fn_s, args.SS_fn_s, args.out_fn)
  
  main()
  