def evaluate_sequence_accuracy(p_pred_step,
                               p_init_cache,
                               state,
                               ds,
                               config,
                               split,
                               workdir,
                               num_eval_steps=-1):
  """Evaluate classification on the given dataset."""
  prediction_dir = os.path.join(workdir, 'predictions')
  tf.io.gfile.makedirs(prediction_dir)
  logging.info('Starting evaluating sequence accuracy on %s split.', split)
  outputs = []

  test_metrics = collections.defaultdict(list)

  data_dir = config.dataset.data_dir
  input_vocab_file = os.path.join(data_dir, 'training_input_vocab.txt')
  target_vocab_file = os.path.join(data_dir, 'training_target_vocab.txt')
  dataset_file = os.path.join(data_dir, 'dataset.txt')
  eos_idx = config.dataset.eos_idx

  with tf.io.gfile.GFile(input_vocab_file, 'r') as f:
    input_vocab = json.load(f)
  with tf.io.gfile.GFile(target_vocab_file, 'r') as f:
    target_vocab = json.load(f)
  with tf.io.gfile.GFile(dataset_file, 'r') as f:
    annotations = json.load(f)

  for step, batch in enumerate(ds):  # pytype: disable=wrong-arg-types
    batch = jax.tree.map(np.asarray, batch)
    cache = p_init_cache(batch)
    batch['predictions'] = p_pred_step(batch, state, cache, eos_idx)
    batch = remove_pad(tohost(batch))
    target_token = batch['target_token']
    predictions = batch['predictions']
    for i, (prediction, target) in enumerate(zip(predictions, target_token)):
      prediction = remove_special_tokens(prediction.tolist(), eos_idx)
      target = remove_special_tokens(target.tolist(), eos_idx)
      acc = evaluation.sequence_accuracy(prediction, target)
      test_metrics['test_accuracy'].append(acc)
      exact_match = 100 if acc == 100 else 0
      test_metrics['test_exact_match'].append(exact_match)

      input_command = remove_special_tokens(batch['token'][i].tolist(), eos_idx)
      index = int(batch['index'][i][0])
      example = annotations['examples'][split][index]
      outputs.append({
          'split': split,
          'index': index,
          'input': array_to_sentence(input_command, input_vocab),
          'prediction': array_to_sentence(prediction, target_vocab),
          'target': array_to_sentence(target, target_vocab),
          'derivation': [example['derivation']],
          'situation': [example['situation']],
          'accuracy': acc,
          'exact_match': True if acc == 100 else False,
          'attention_weights_input': [],
          'attention_weights_situation': [],
      })
    if num_eval_steps > 0 and step + 1 == num_eval_steps:
      break
  test_metrics = {k: sum(v) / len(v) for k, v in test_metrics.items()}
  step = flax_utils.unreplicate(state).step
  out_path = os.path.join(prediction_dir, f'{split}_predict_{step}.json')
  with tf.io.gfile.GFile(out_path, 'w') as f:
    json.dump(outputs, f, indent=2)
  return test_metrics
Evaluate classification on the given dataset.
language/gscan/xattn_model/predict.py