text: import json
  from .lib.dbengine import DBEngine
  import re
  import numpy as np
  #from nltk.tokenize import StanfordTokenizer
  
  def load_data(sql_paths, table_paths, use_small=False):
      if not isinstance(sql_paths, list):
          sql_paths = (sql_paths, )
      if not isinstance(table_paths, list):
          table_paths = (table_paths, )
      sql_data = []
      table_data = {}
  
      max_col_num = 0
      for SQL_PATH in sql_paths:
          print("Loading data from %s"%SQL_PATH)
          with open(SQL_PATH) as inf:
              for idx, line in enumerate(inf):
                  if use_small and idx >= 1000:
                      break
                  sql = json.loads(line.strip())
                  sql_data.append(sql)
  
  
      for TABLE_PATH in table_paths:
          print("Loading data from %s"%TABLE_PATH)
          with open(TABLE_PATH) as inf:
              for line in inf:
                  tab = json.loads(line.strip())
                  table_data[tab['id']] = tab
  
      for sql in sql_data:
          assert sql['table_id'] in table_data
      
      return sql_data, table_data
  
  def load_dataset(dataset_id, use_small=False):
      if dataset_id == 0:
          print("Loading from original dataset")
          sql_data, table_data = load_data('data/train_tok.jsonl',
                  'data/train_tok.tables.jsonl', use_small=use_small)
          val_sql_data, val_table_data = load_data('data/dev_tok.jsonl',
                  'data/dev_tok.tables.jsonl', use_small=use_small)
          test_sql_data, test_table_data = load_data('data/test_tok.jsonl',
                  'data/test_tok.tables.jsonl', use_small=use_small)
          TRAIN_DB = 'data/train.db'
          DEV_DB = 'data/dev.db'
          TEST_DB = 'data/test.db'
      else:
          print("Loading from re-split dataset")
          sql_data, table_data = load_data('data_resplit/train.jsonl',
                  'data_resplit/tables.jsonl', use_small=use_small)
          val_sql_data, val_table_data = load_data('data_resplit/dev.jsonl',
                  'data_resplit/tables.jsonl', use_small=use_small)
          test_sql_data, test_table_data = load_data('data_resplit/test.jsonl',
                  'data_resplit/tables.jsonl', use_small=use_small)
          TRAIN_DB = 'data_resplit/table.db'
          DEV_DB = 'data_resplit/table.db'
          TEST_DB = 'data_resplit/table.db'
  
      return sql_data, table_data, val_sql_data, val_table_data,\
              test_sql_data, test_table_data, TRAIN_DB, DEV_DB, TEST_DB
  
  def best_model_name(args, for_load=False):
      new_data = 'new' if args.dataset > 0 else 'old'
      mode = 'seq2sql' if args.baseline else 'sqlnet'
      if for_load:
          use_emb = use_rl = ''
      else:
          use_emb = '_train_emb' if args.train_emb else ''
          use_rl = 'rl_' if args.rl else ''
      use_ca = '_ca' if args.ca else ''
  
      agg_model_name = 'saved_model/%s_%s%s%s.agg_model'%(new_data,
              mode, use_emb, use_ca)
      sel_model_name = 'saved_model/%s_%s%s%s.sel_model'%(new_data,
              mode, use_emb, use_ca)
      cond_model_name = 'saved_model/%s_%s%s%s.cond_%smodel'%(new_data,
              mode, use_emb, use_ca, use_rl)
  
      if not for_load and args.train_emb:
          agg_embed_name = 'saved_model/%s_%s%s%s.agg_embed'%(new_data,
                  mode, use_emb, use_ca)
          sel_embed_name = 'saved_model/%s_%s%s%s.sel_embed'%(new_data,
                  mode, use_emb, use_ca)
          cond_embed_name = 'saved_model/%s_%s%s%s.cond_embed'%(new_data,
                  mode, use_emb, use_ca)
  
          return agg_model_name, sel_model_name, cond_model_name,\
                  agg_embed_name, sel_embed_name, cond_embed_name
      else:
          return agg_model_name, sel_model_name, cond_model_name
  
  
  def to_batch_seq(sql_data, table_data, idxes, st, ed, ret_vis_data=False):
      q_seq = []
      col_seq = []
      col_num = []
      ans_seq = []
      query_seq = []
      gt_cond_seq = []
      vis_seq = []
      for i in range(st, ed):
          sql = sql_data[idxes[i]]
          q_seq.append(sql['question_tok'])
          col_seq.append(table_data[sql['table_id']]['header_tok'])
          col_num.append(len(table_data[sql['table_id']]['header']))
          ans_seq.append((sql['sql']['agg'],
              sql['sql']['sel'], 
              len(sql['sql']['conds']),
              tuple(x[0] for x in sql['sql']['conds']),
              tuple(x[1] for x in sql['sql']['conds'])))
          query_seq.append(sql['query_tok'])
          gt_cond_seq.append(sql['sql']['conds'])
          vis_seq.append((sql['question'],
              table_data[sql['table_id']]['header'], sql['query']))
      if ret_vis_data:
          return q_seq, col_seq, col_num, ans_seq, query_seq, gt_cond_seq, vis_seq
      else:
          return q_seq, col_seq, col_num, ans_seq, query_seq, gt_cond_seq
  
  def to_batch_query(sql_data, idxes, st, ed):
      query_gt = []
      table_ids = []
      for i in range(st, ed):
          query_gt.append(sql_data[idxes[i]]['sql'])
          table_ids.append(sql_data[idxes[i]]['table_id'])
      return query_gt, table_ids
  
  def epoch_exec_acc_from_user(model, batch_size, sql_data, table_data, db_path):
      engine = DBEngine(db_path)
      model.eval()
      perm = list(range(len(sql_data)))
      tot_acc_num = 0.0
      acc_of_log = 0.0
      st = 0
      while st < len(sql_data):
          ed = st+batch_size if st+batch_size < len(perm) else len(perm)
          q_seq, col_seq, col_num, ans_seq, query_seq, gt_cond_seq, raw_data = \
              to_batch_seq(sql_data, table_data, perm, st, ed, ret_vis_data=True)
          raw_q_seq = [x[0] for x in raw_data]
          raw_col_seq = [x[1] for x in raw_data]
          print('q_seq ----- >', q_seq)
          print('col_seq ----- >', col_seq)
          print('query_seq ----- >', query_seq)
          gt_where_seq = model.generate_gt_where_seq(q_seq, col_seq, query_seq)
          query_gt, table_ids = to_batch_query(sql_data, perm, st, ed)
          gt_sel_seq = [x[1] for x in ans_seq]
          score = model.forward(q_seq, col_seq, col_num,
                  (True, True, True), gt_sel=gt_sel_seq)
          pred_queries = model.gen_query(score, q_seq, col_seq,
                  raw_q_seq, raw_col_seq, (True, True, True))
                  
          # print(pred_queries)
          
          for idx, (sql_gt, sql_pred, tid) in enumerate(
                  zip(query_gt, pred_queries, table_ids)):
  #             ret_gt = engine.execute(tid,
  #                     sql_gt['sel'], sql_gt['agg'], sql_gt['conds'])
              try:
                  print("\n\n---->select index", sql_pred['sel'], "\n\n")
                  ret_pred = engine.execute(tid,
                          sql_pred['sel'], sql_pred['agg'], sql_pred['conds'])
                  #return ret_pred
                 #print('select prediction ---> ', sql_pred['sel'])
                  #print('agg prediction ---> ', sql_pred['agg'])
                  #print('conds prediction ---> ', sql_pred['conds'])
              except:
                  ret_pred = None
              #tot_acc_num += (ret_gt == ret_pred)
          
          st = ed
      
      predicted_dict = {}
      predicted_dict['col_index'] = sql_pred['sel'] 
      predicted_dict['value'] = ret_pred
      return predicted_dict
  
  def epoch_train(model, optimizer, batch_size, sql_data, table_data, pred_entry):
      model.train()
      perm=np.random.permutation(len(sql_data))
      cum_loss = 0.0
      st = 0
      while st < len(sql_data):
          ed = st+batch_size if st+batch_size < len(perm) else len(perm)
  
          q_seq, col_seq, col_num, ans_seq, query_seq, gt_cond_seq = \
                  to_batch_seq(sql_data, table_data, perm, st, ed)
          gt_where_seq = model.generate_gt_where_seq(q_seq, col_seq, query_seq)
          gt_sel_seq = [x[1] for x in ans_seq]
          score = model.forward(q_seq, col_seq, col_num, pred_entry,
                  gt_where=gt_where_seq, gt_cond=gt_cond_seq, gt_sel=gt_sel_seq)
          loss = model.loss(score, ans_seq, pred_entry, gt_where_seq)
  #         if (!is)
          print('indices------------>', loss.data.cpu().numpy())
          cum_loss += loss.data.cpu().numpy()[0]*(ed - st)
          optimizer.zero_grad()
          loss.backward()
          optimizer.step()
  
          st = ed
  
      return cum_loss / len(sql_data)
  
  def epoch_exec_acc(model, batch_size, sql_data, table_data, db_path):
      engine = DBEngine(db_path)
  
      model.eval()
      perm = list(range(len(sql_data)))
      tot_acc_num = 0.0
      acc_of_log = 0.0
      st = 0
      while st < len(sql_data):
          ed = st+batch_size if st+batch_size < len(perm) else len(perm)
          q_seq, col_seq, col_num, ans_seq, query_seq, gt_cond_seq, raw_data = \
              to_batch_seq(sql_data, table_data, perm, st, ed, ret_vis_data=True)
          raw_q_seq = [x[0] for x in raw_data]
          raw_col_seq = [x[1] for x in raw_data]
          gt_where_seq = model.generate_gt_where_seq(q_seq, col_seq, query_seq)
          query_gt, table_ids = to_batch_query(sql_data, perm, st, ed)
          gt_sel_seq = [x[1] for x in ans_seq]
          score = model.forward(q_seq, col_seq, col_num,
                  (True, True, True), gt_sel=gt_sel_seq)
          pred_queries = model.gen_query(score, q_seq, col_seq,
                  raw_q_seq, raw_col_seq, (True, True, True))
  
          for idx, (sql_gt, sql_pred, tid) in enumerate(
                  zip(query_gt, pred_queries, table_ids)):
              ret_gt = engine.execute(tid,
                      sql_gt['sel'], sql_gt['agg'], sql_gt['conds'])
              try:
                  ret_pred = engine.execute(tid,
                          sql_pred['sel'], sql_pred['agg'], sql_pred['conds'])
              except:
                  ret_pred = None
              tot_acc_num += (ret_gt == ret_pred)
          
          st = ed
  
      return tot_acc_num / len(sql_data)
  
  def epoch_acc(model, batch_size, sql_data, table_data, pred_entry):
      model.eval()
      perm = list(range(len(sql_data)))
      st = 0
      one_acc_num = 0.0
      tot_acc_num = 0.0
      while st < len(sql_data):
          ed = st+batch_size if st+batch_size < len(perm) else len(perm)
  
          q_seq, col_seq, col_num, ans_seq, query_seq, gt_cond_seq, raw_data = to_batch_seq(sql_data, table_data, perm, st, ed, ret_vis_data=True)
          raw_q_seq = [x[0] for x in raw_data]
          raw_col_seq = [x[1] for x in raw_data]
          query_gt, table_ids = to_batch_query(sql_data, perm, st, ed)
          gt_sel_seq = [x[1] for x in ans_seq]
          score = model.forward(q_seq, col_seq, col_num,
                  pred_entry, gt_sel = gt_sel_seq)
          pred_queries = model.gen_query(score, q_seq, col_seq,
                  raw_q_seq, raw_col_seq, pred_entry)
          one_err, tot_err = model.check_acc(raw_data,
                  pred_queries, query_gt, pred_entry)
  
          one_acc_num += (ed-st-one_err)
          tot_acc_num += (ed-st-tot_err)
  
          st = ed
      return tot_acc_num / len(sql_data), one_acc_num / len(sql_data)
  
  def epoch_reinforce_train(model, optimizer, batch_size, sql_data, table_data, db_path):
      engine = DBEngine(db_path)
  
      model.train()
      perm = np.random.permutation(len(sql_data))
      cum_reward = 0.0
      st = 0
      while st < len(sql_data):
          ed = st+batch_size if st+batch_size < len(perm) else len(perm)
  
          q_seq, col_seq, col_num, ans_seq, query_seq, gt_cond_seq, raw_data =\
              to_batch_seq(sql_data, table_data, perm, st, ed, ret_vis_data=True)
          gt_where_seq = model.generate_gt_where_seq(q_seq, col_seq, query_seq)
          raw_q_seq = [x[0] for x in raw_data]
          raw_col_seq = [x[1] for x in raw_data]
          query_gt, table_ids = to_batch_query(sql_data, perm, st, ed)
          gt_sel_seq = [x[1] for x in ans_seq]
          score = model.forward(q_seq, col_seq, col_num, (True, True, True),
                  reinforce=True, gt_sel=gt_sel_seq)
          pred_queries = model.gen_query(score, q_seq, col_seq, raw_q_seq,
                  raw_col_seq, (True, True, True), reinforce=True)
  
          query_gt, table_ids = to_batch_query(sql_data, perm, st, ed)
          rewards = []
          for idx, (sql_gt, sql_pred, tid) in enumerate(
                  zip(query_gt, pred_queries, table_ids)):
              ret_gt = engine.execute(tid,
                      sql_gt['sel'], sql_gt['agg'], sql_gt['conds'])
              try:
                  ret_pred = engine.execute(tid,
                          sql_pred['sel'], sql_pred['agg'], sql_pred['conds'])
              except:
                  ret_pred = None
  
              if ret_pred is None:
                  rewards.append(-2)
              elif ret_pred != ret_gt:
                  rewards.append(-1)
              else:
                  rewards.append(1)
  
          cum_reward += (sum(rewards))
          optimizer.zero_grad()
          model.reinforce_backward(score, rewards)
          optimizer.step()
  
          st = ed
  
      return cum_reward / len(sql_data)
  
  
  def load_word_emb(file_name, load_used=False, use_small=False):
      if not load_used:
          print(('Loading word embedding from %s'%file_name))
          ret = {}
          with open(file_name, encoding='utf-8') as inf:
              for idx, line in enumerate(inf):
                  if (use_small and idx >= 5000):
                      break
                  info = line.strip().split(' ')
                  if info[0].lower() not in ret:
                      ret[info[0]] = np.array([float(x) for x in info[1:]])
          return ret
      else:
          print ('Load used word embedding')
          with open('glove/word2idx.json') as inf:
              w2i = json.load(inf)
          with open('glove/usedwordemb.npy', mode='rb') as inf:
              word_emb_val = np.load(inf)
          return w2i, word_emb_val
  