text: import copy
  import os
  from pathlib import Path
  
  import torch
  import wandb
  from tqdm import tqdm
  import json
  
  from config import read_arguments_evaluation
  from data_loader import get_data_loader
  from intermediate_representation import semQL
  from intermediate_representation.sem2sql.sem2SQL import transform_semQL_to_sql
  from manual_inference.helper import _execute_query_postgresql
  from model.model import IRNet
  from spider import spider_utils
  from spider.example_builder import build_example
  from spider.test_suite_eval.exec_eval import result_eq
  from utils import setup_device, set_seed_everywhere
  import spider.test_suite_eval.evaluation as spider_evaluation
  
  
  def evaluate(model, dev_loader, schema, beam_size):
      model.eval()
  
      sketch_correct, rule_label_correct, found_in_beams, not_all_values_found, total = 0, 0, 0, 0, 0
      predictions = []
      for batch in tqdm(dev_loader, desc="Evaluating"):
  
          for data_row in batch:
              original_row = copy.deepcopy(data_row)
  
              try:
                  example = build_example(data_row, schema)
              except Exception as e:
                  print("Exception while building example (evaluation): {}".format(e))
                  continue
  
              with torch.no_grad():
                  results_all = model.parse(example, beam_size=beam_size)
  
              results = results_all[0]
              all_predictions = []
              try:
                  # here we set assemble the predicted actions (including leaf-nodes) as string
                  full_prediction = " ".join([str(x) for x in results[0].actions])
                  for beam in results:
                      all_predictions.append(" ".join([str(x) for x in beam.actions]))
              except Exception as e:
                  # print(e)
                  full_prediction = ""
  
              prediction = original_row
  
              # here we set assemble the predicted sketch actions as string
              prediction['sketch_result'] = " ".join(str(x) for x in results_all[1])
              prediction['model_result'] = full_prediction
  
              truth_sketch = " ".join([str(x) for x in example.sketch])
              truth_rule_label = " ".join([str(x) for x in example.semql_actions])
  
              if prediction['all_values_found']:
                  if truth_sketch == prediction['sketch_result']:
                      sketch_correct += 1
                  if truth_rule_label == prediction['model_result']:
                      rule_label_correct += 1
                  elif truth_rule_label in all_predictions:
                      found_in_beams += 1
              else:
                  question = prediction['question']
                  tqdm.write(f'Not all values found during pre-processing for question "{question}". Replace values with dummy to make query fail')
                  prediction['values'] = [1] * len(prediction['values'])
                  not_all_values_found += 1
  
              total += 1
  
              predictions.append(prediction)
  
      print(f"in {found_in_beams} times we found the correct results in another beam (failing queries: {total - rule_label_correct})")
  
      return float(sketch_correct) / float(total), float(rule_label_correct) / float(total), float(not_all_values_found) / float(total), predictions
  
  
  def transform_to_sql_and_evaluate_with_spider(predictions, table_data, experiment_dir, data_dir, training_step):
      total_count, failure_count = transform_semQL_to_sql(table_data, predictions, experiment_dir)
  
      spider_eval_results = spider_evaluation.evaluate(os.path.join(experiment_dir, 'ground_truth.txt'),
                                                       os.path.join(experiment_dir, 'output.txt'),
                                                       os.path.join(data_dir, "testsuite_databases"),
                                                       "exec",
                                                       None,
                                                       False,
                                                       False,
                                                       False,
                                                       training_step,
                                                       quickmode=False)
  
      return total_count, failure_count, spider_eval_results
  
  
  def _remove_unnecessary_distinct(p, g):
      p_tokens = p.split(' ')
      g_tokens = g.split(' ')
  
      if 'distinct' not in g_tokens and 'DISTINCT' not in g_tokens:
          return ' '.join([p_token for p_token in p_tokens if p_token != 'DISTINCT'])
      else:
          return p
  
  
  def evaluate_cordis(groundtruth_path: Path, prediction_path: Path, database: str, connection_config: dict, do_not_verify_distinct=True):
  
      with open(groundtruth_path, 'r', encoding='utf-8') as f:
          groundtruth_full_lines = f.readlines()
          groundtruth = [g.strip().split('\t')[0] for g in groundtruth_full_lines]
          questions = [g.strip().split('\t')[2] for g in groundtruth_full_lines]
  
      with open(prediction_path, 'r', encoding='utf-8') as f:
          prediction = f.readlines()
  
      n_success = 0
      for p, g, q in zip(prediction, groundtruth, questions):
  
          if do_not_verify_distinct:
              p = _remove_unnecessary_distinct(p, g)
  
          try:
              results_prediction = _execute_query_postgresql(p, database, connection_config)
          except Exception as ex:
              print("Prediction is not a valid SQL query:")
              print(f"Q: {q.strip()}")
              print(f"P: {p.strip()}")
              print(f"G: {g.strip()}")
              print()
  
              continue
  
          results_groundtruth = _execute_query_postgresql(g, database, connection_config)
  
          if result_eq(results_groundtruth, results_prediction, order_matters=True):
              n_success += 1
          else:
              print("Results not equal for:")
              print(f"Q: {q.strip()}")
              print(f"P: {p.strip()}")
              print(f"G: {g.strip()}")
              print()
  
      print(f"================== Total score: {n_success} of {len(groundtruth)} ({100 / len(groundtruth) * n_success}%) ==================")
  
  
  def main():
      args = read_arguments_evaluation()
  
      device, n_gpu = setup_device()
      set_seed_everywhere(args.seed, n_gpu)
  
      sql_data, table_data, val_sql_data, val_table_data = spider_utils.load_dataset(args.data_dir, use_small=False)
      _, dev_loader = get_data_loader(sql_data, val_sql_data, args.batch_size, True, False)
  
      grammar = semQL.Grammar()
      model = IRNet(args, device, grammar)
      model.to(device)
  
      # load the pre-trained parameters
      model.load_state_dict(torch.load(args.model_to_load))
      print("Load pre-trained model from '{}'".format(args.model_to_load))
  
      sketch_acc, acc, not_all_values_found, predictions = evaluate(model,
                                                                    dev_loader,
                                                                    table_data,
                                                                    args.beam_size)
  
      print("Predicted {} examples. Start now converting them to SQL. Sketch-Accuracy: {}, Accuracy: {}, Not all values found: {}".format(
              len(dev_loader), sketch_acc, acc, not_all_values_found))
  
      with open(os.path.join(args.prediction_dir, 'predictions_sem_ql.json'), 'w', encoding='utf-8') as f:
          json.dump(predictions, f, indent=2)
  
      # with open(os.path.join(args.prediction_dir, 'predictions_sem_ql.json'), 'r', encoding='utf-8') as json_file:
      #     predictions = json.load(json_file)
  
      count_success, count_failed = transform_semQL_to_sql(val_table_data, predictions, args.prediction_dir)
  
      print("Transformed {} samples successful to SQL. {} samples failed. Generated the files a 'ground_truth.txt' "
            "and a 'output.txt' file.".format(
          count_success, count_failed))
  
      if args.evaluation_type == 'spider':
  
          print('We now use the official Spider evaluation script to evaluate the generated/ground truth files.')
          wandb.init(project="proton")
  
          spider_evaluation.evaluate(os.path.join(args.prediction_dir, 'ground_truth.txt'),
                                     os.path.join(args.prediction_dir, 'output.txt'),
                                     os.path.join(args.data_dir, "testsuite_databases"),
                                     'exec', None, False, False, False, 1, quickmode=False)
  
      elif args.evaluation_type == 'cordis':
  
          connection_config = {k: v for k, v in vars(args).items() if k.startswith('database')}
  
          evaluate_cordis(Path(args.prediction_dir) / 'ground_truth.txt',
                          Path(args.prediction_dir) / 'output.txt',
                          args.database,
                          connection_config)
      else:
          raise NotImplemented('Only Spider and CORDIS evaluation are implemented so far.')
  
  
  if __name__ == '__main__':
      main()
  