text: import argparse
  import json
  from typing import List
  
  import torch
  from nltk.translate.bleu_score import corpus_bleu, SmoothingFunction, sentence_bleu
  import spacy
  import tqdm
  import numpy as np
  import rouge
  import edlib
  import os
  import pandas as pd
  import re
  import glob
  
  from pytorch_pretrained_bert import BertTokenizer, BertModel
  from wmd import WMD
  from torch.nn.modules.distance import CosineSimilarity
  
  torch_emb_sim = CosineSimilarity()
  
  from bert_score import score as bert_score
  
  
  nlp = spacy.load("en_core_web_md")
  nlp.add_pipe(WMD.SpacySimilarityHook(nlp), last=True)
  
  
  def _clean_text(txt):
      return txt.lower()
  
  
  class CFRInstance(object):
      def __init__(self,
                   original_context: str,
                   cf_context: str,
                   original_ending: str,
                   predicted_ending: str,
                   gold_cf_endings: List[str],
                   ):
          self.original_context = original_context
          self.cf_context = cf_context
  
          self.predicted_ending = predicted_ending
          self.original_ending = original_ending
          self.gold_cf_endings = gold_cf_endings
  
          self.spacy_docs = {
              'original_context': nlp(_clean_text(self.original_context)),
              'original_ending': nlp(_clean_text(self.original_ending)),
              'cf_context': nlp(_clean_text(self.cf_context)),
              'predicted_ending': nlp(_clean_text(self.predicted_ending)),
              'gold_cf_endings': [nlp(_clean_text(g)) for g in self.gold_cf_endings]
          }
  
          self.original_context_tokens = [t.text for t in self.spacy_docs['original_context']]
          self.original_ending_tokens = [t.text for t in self.spacy_docs['original_ending']]
          self.cf_context_tokens = [t.text for t in self.spacy_docs['cf_context']]
          self.predicted_ending_tokens = [t.text for t in self.spacy_docs['predicted_ending']]
          self.gold_cf_endings_tokens = [[t.text for t in _spacy_doc] for _spacy_doc in
                                         self.spacy_docs['gold_cf_endings']]
  
  
  def read_lines(filename):
      lines = []
      with open(filename, "r") as f:
          for line in tqdm.tqdm(f):
              l = line.strip()
              if len(re.sub(r'[^\w\s]', '', l)) == 0:
                  lines.append("")
              else:
                  lines.append(l)
      return lines
  
  
  def read_jsonl_lines(filename):
      with open(filename) as f:
          for line in tqdm.tqdm(f):
              yield json.loads(line.strip())
  
  
  def _read_gold_cf_endings(gold_cf_endings_dir):
      gold_cf_endings = []
      for f in os.listdir(gold_cf_endings_dir):
          df = pd.read_csv(os.path.join(gold_cf_endings_dir, f))
          new_3 = df["new_3"]
          new_4 = df["new_4"]
          new_5 = df["new_5"]
  
          gold_cf_ending_lst = []
  
          for s3, s4, s5 in zip(new_3, new_4, new_5):
              gold_cf_ending_lst.append(' '.join([s3, s4, s5]))
          gold_cf_endings.append(gold_cf_ending_lst)
  
      return [list(Z) for Z in zip(*gold_cf_endings)]
  
  
  def eval_bleu(instances: List[CFRInstance]):
      references = []
      hypotheses = []
      for instance in tqdm.tqdm(instances):
          references.append(instance.gold_cf_endings_tokens)
          hypotheses.append(instance.predicted_ending_tokens)
  
      corpus_bleu_scores = corpus_bleu(
          references, hypotheses, smoothing_function=SmoothingFunction().method4
      )
  
      sentence_bleu_scores = []
      total_skipped = 0
      for r, h in tqdm.tqdm(zip(references, hypotheses)):
          if len(h) == 0:
              sentence_bleu_scores.append(0)
              continue
          else:
              try:
                  sentence_bleu_scores.append(
                      sentence_bleu(r, h, smoothing_function=SmoothingFunction().method4))
              except:
                  sentence_bleu_scores.append(0.0)
                  total_skipped+=1
  
      print("Total skipped = {}".format(total_skipped))
  
      metrics = {
          'corpus_bleu': corpus_bleu_scores,
          'mean_sentence_bleu': np.mean(sentence_bleu_scores),
          'sentence_bleu_by_instance': sentence_bleu_scores
      }
      return metrics
  
  
  def eval_rouge(instances: List[CFRInstance]):
      references = []
      hypotheses = []
  
      evaluator = rouge.Rouge(metrics=['rouge-n', 'rouge-l', 'rouge-w'],
                              max_n=4,
                              limit_length=True,
                              length_limit=100,
                              length_limit_type='words',
                              apply_avg=True,
                              apply_best=False,
                              alpha=0.5,  # Default F1_score
                              weight_factor=1.2,
                              stemming=True)
  
      by_instance = []
      for instance in instances:
          _r = [_clean_text(g) for g in instance.gold_cf_endings]
          _h = _clean_text(instance.predicted_ending)
          references.append(_r)
          hypotheses.append(_h)
          try:
              by_instance.append(evaluator.get_scores(_h, _r))
          except:
              by_instance.append({})
  
      scores = evaluator.get_scores(hypotheses, references)
      return {'rouge_all' : scores,
              'rouge_by_instance': by_instance
              }
  
  
  def eval_bert_score(instances: List[CFRInstance], bert_model="bert-base-uncased"):
      references = []
      hypotheses = []
      for instance in instances:
          clean_reference = _clean_text(instance.original_context + ' ' + instance.original_ending)
          clean_hypothesis = _clean_text(instance.cf_context + ' ' + instance.predicted_ending)
          if len(clean_hypothesis) == 0:
              continue
          references.append(clean_reference)
          hypotheses.append(clean_hypothesis)
  
      P, R, F1 = bert_score(hypotheses, references, bert=bert_model, verbose=True)
      return {
          "bert_score_P": P.mean().item(),
          "bert_score_R": R.mean().item(),
          "bert_score_F1": F1.mean().item(),
          "bert_score_P_by_instance": [float(f) for f in list(P.numpy())],
          "bert_score_R_by_instance": [float(f) for f in list(R.numpy())],
          "bert_score_F1_by_instance": [float(f) for f in list(F1.numpy())],
      }
  
  
  def cigar_to_word_sets(cigar_path: str, modification_lst, reference_lst):
      diff_counts = [int(c) for c in re.split("[D=IX]", cigar_path) if c != '']
      diff_types = [c for c in cigar_path if c in {'D', '=', 'I', 'X'}]
  
      reference_ctr = 0
      modification_ctr = 0
  
      deleted_items = []
      inserted_items = []
      replaced_items = []
      equal_items = []
  
      deleted_counts = 0
      inserted_counts = 0
      replaced_counts = 0
      equal_counts = 0
  
      for count, type in zip(diff_counts, diff_types):
          if type == "D":
              # Deleted token
              for i in range(count):
                  deleted_items.append(reference_lst[reference_ctr])
                  reference_ctr += 1
              deleted_counts += count
          elif type == "I":
              # Inserted token
              for i in range(count):
                  inserted_items.append(modification_lst[modification_ctr])
                  modification_ctr += 1
              inserted_counts += count
          elif type == "=":
              # Same token
              for i in range(count):
                  equal_items.append(reference_lst[reference_ctr])
                  reference_ctr += 1
                  modification_ctr += 1
              equal_counts += count
          elif type == "X":
              # Exchanged items
              for i in range(count):
                  replaced_items.append(
                      (reference_lst[reference_ctr], modification_lst[modification_ctr]))
                  reference_ctr += 1
                  modification_ctr += 1
              replaced_counts += count
  
      return {
          "deleted": deleted_items,
          "inserted": inserted_items,
          "replaced": replaced_items,
          "equal": equal_items
      }
  
  
  def compare_edit_sets_unigram(predicted_edits, reference_edits):
      keys = ["deleted", "inserted", "replaced", "equal"]
  
      total_edits = np.sum([len(set(reference_edits[k])) for k in keys])
  
      equivalent_edits_count = 0
      for k in keys:
          equivalent_edits_count += len(
              set(reference_edits[k]).intersection(set(predicted_edits[k])))
  
      return equivalent_edits_count / total_edits
  
  
  def compare_edit_sets(predicted_edits, reference_edits, method="unigram"):
      if method == "unigram":
          return compare_edit_sets_unigram(predicted_edits, reference_edits)
  
  
  def eval_rewrite(instances: List[CFRInstance]):
      instance_score = []
  
      for instance in tqdm.tqdm(instances):
          if len(instance.predicted_ending_tokens) == 0:
              instance_score.append(0)
              continue
          predicted_edits = edlib.align(
              instance.predicted_ending_tokens,
              instance.original_ending_tokens,
              mode="NW",
              task="path"
          )
          predicted_edits_set = \
              cigar_to_word_sets(predicted_edits['cigar'],
                                 instance.predicted_ending_tokens,
                                 instance.original_ending_tokens
                                 )
  
          scores = []
          gold_edit_sets = []
          for gold_cf in instance.gold_cf_endings_tokens:
              gold_edits = edlib.align(
                  gold_cf,
                  instance.original_ending_tokens,
                  mode="NW",
                  task="path"
              )
              gold_edits_set = \
                  cigar_to_word_sets(gold_edits['cigar'],
                                     gold_cf,
                                     instance.original_ending_tokens
                                     )
              gold_edit_sets.append(gold_edits_set)
  
              scores.append(compare_edit_sets(predicted_edits_set, gold_edits_set))
  
          instance_score.append(max(scores))
  
      return {
          'CFR_METRIC': np.mean(instance_score),
          'CFR_METRIC_by_instance': instance_score
      }
  
  
  def eval_wmd(instances: List[CFRInstance]):
      wmd_scores = []
      for instance in instances:
          pred_ending_spacy_doc = instance.spacy_docs['predicted_ending']
          wmd_scores.append(
              np.min(
                  [pred_ending_spacy_doc.similarity(gold_spacy_doc)
                   for gold_spacy_doc in instance.spacy_docs['gold_cf_endings']]
              )
          )
      return {
          'mean_wmd': np.mean(wmd_scores),
          'wmd_by_instance': wmd_scores
      }
  
  
  def _bert_embed_sentence(sentence, bert_model: BertModel, bert_tokenizer: BertTokenizer):
      text = "[CLS] {} [SEP]".format(sentence)
      tokenized_text = bert_tokenizer.tokenize(text)
      indexed_tokens = bert_tokenizer.convert_tokens_to_ids(tokenized_text)
      tokens_tensor = torch.tensor([indexed_tokens])
      segments_ids = [0] * len(indexed_tokens)
      segments_tensors = torch.tensor([segments_ids])
  
      device = 'cuda' if torch.cuda.is_available() else 'cpu'
      tokens_tensor = tokens_tensor.to(device)
      segments_tensors = segments_tensors.to(device)
  
      with torch.no_grad():
          encoded_layers, _ = bert_model(tokens_tensor, segments_tensors, output_all_encoded_layers=False)
  
      # Embedding of the [CLS] token
      return encoded_layers[0][0]
  
  
  def drift_similarity(original_story_emb, predicted_ending_emb, gold_cf_emb):
      drift_1 = predicted_ending_emb - original_story_emb
      drift_2 = gold_cf_emb - original_story_emb
      return torch_emb_sim(drift_1.unsqueeze(0), drift_2.unsqueeze(0)).item()
  
  
  def eval_semantic_sim_score(instances: List[CFRInstance], bert_model_type="bert-base-uncased"):
  
      tokenizer = BertTokenizer.from_pretrained(bert_model_type)
      model = BertModel.from_pretrained(bert_model_type)
      model.eval()
      device = 'cuda' if torch.cuda.is_available() else 'cpu'
      model.to(device)
  
      drift_similarities = []
  
      for instance in instances:
          clean_original_story = _clean_text(instance.original_context + ' ' + instance.original_ending)
          predicted_ending = _clean_text(instance.cf_context + ' ' + instance.predicted_ending)
  
          original_story_emb = _bert_embed_sentence(clean_original_story, model, tokenizer)
          predicted_ending_emb = _bert_embed_sentence(predicted_ending, model, tokenizer)
  
          all_sims = []
          for gold_cf in instance.gold_cf_endings:
              clean_gold_cf = _clean_text(instance.cf_context + ' ' + gold_cf)
              gold_cf_emb = _bert_embed_sentence(clean_gold_cf, model, tokenizer)
  
              all_sims.append(drift_similarity(original_story_emb, predicted_ending_emb, gold_cf_emb))
  
          drift_similarities.append(np.max(all_sims))
  
      return {
          "drift_similarity": np.mean(drift_similarities),
          "drift_similarity_by_instance": [float(f) for f in  drift_similarities]
      }
  
  
  def main(pred_endings_file, gold_file, bert_model):
      pred_endings = read_lines(filename=pred_endings_file)
      gold_records = read_jsonl_lines(gold_file)
  
      instances = []
      for pe, record in tqdm.tqdm(
              zip(pred_endings, gold_records)
      ):
          instance = CFRInstance(
              original_context=record['ori_context'],
              cf_context=record['cf_context'],
              predicted_ending=pe,
              original_ending=' '.join(record['ori_endinng']),
              gold_cf_endings=[' '.join(_ge) for _ge in record['gold_end']]
          )
          instances.append(instance)
      metrics = {}
      print("Eval BLEU ... ")
      metrics.update(eval_bleu(instances))
      print("Eval ROUGE ... ")
      metrics.update(eval_rouge(instances))
      print("Eval BertScore ... ")
      metrics.update(eval_bert_score(instances, bert_model=bert_model))
      # print("Eval CFRScore ... ")
      # metrics.update(eval_rewrite(instances))
      # print("Eval WMD ... ")
      # metrics.update(eval_wmd(instances))
      print("Eval Drift Similarity ... ")
      metrics.update(eval_semantic_sim_score(instances, bert_model_type=bert_model))
      print(metrics)
      print(json.dumps(metrics, indent=2))
      return metrics
  
  
  if __name__ == '__main__':
      parser = argparse.ArgumentParser(
          prog='evaluate.py',
          usage='%(prog)s gold_annotations predictions',
          description='Evaluate story rewrite'
      )
  
      parser.add_argument('--pred-endings-file', type=str,
                          dest="pred_endings_file",
                          help='Location of prediction file. Usually named *_pred.txt',
                          default=None)
  
      parser.add_argument('--all-preds-dir', type=str,
                          dest="all_preds_dir",
                          help='Location of prediction file. Usually named *_pred.txt',
                          default=None)
  
      parser.add_argument('--gold-file', type=str,
                          dest="gold_file",
                          help='Location of human annotated cf endings and rest of the data. Usually named [dev/test].jsonl')
  
      parser.add_argument('--bert_model', type=str,
                          dest="bert_model",
                          help='Location of human annotated cf endings and rest of the data. Usually named [dev/test].jsonl',
                          default=None)
  
      parser.add_argument('--output_file', type=str,
                          dest="output_file",
                          help='')
  
      args = parser.parse_args()
  
      # Run seed selection if args valid
      print('====Input Arguments====')
      print(json.dumps(vars(args), indent=2, sort_keys=True))
      print("=======================")
  
      assert args.all_preds_dir is not None or args.pred_endings_file is not None
  
      all_metrics = {}
      if args.all_preds_dir is not None:
          for f in glob.iglob(args.all_preds_dir + "/*/*.txt"):
              print("Processing file {}".format(f))
              metrics = main(f, args.gold_file, args.bert_model)
              model_name = os.path.basename(f).split(".")[0]
              all_metrics[model_name] = metrics
  
      else:
          all_metrics = main(args.pred_endings_file, args.gold_file, args.bert_model)
  
      with open(args.output_file, "w") as f:
          f.write(json.dumps(all_metrics))
          f.close()
  