text: # Copyright (c) Facebook, Inc. and its affiliates.
  #
  # This source code is licensed under the MIT license found in the
  # LICENSE file in the root directory of this source tree.
  
  import numpy as np
  import nltk
  from misc.bleu_utils import sentence_bleu
  import warnings
  
  
  def get_target_sequences(manifest, ground_truth, to_take=1000):
      import json
      import pathlib
  
      with open(ground_truth, 'r') as fin:
          original_continuations = json.loads(fin.read())
  
      sequence2length = [(k, v[0]) for k, v in original_continuations.items()]
      assert all(float(v) >= 6.0 for (_, v) in sequence2length)  # 6 seconds
  
      sequence2length.sort(key=lambda x: x[1])
      to_take_sequences = set(v[0] for v in sequence2length[:to_take])
      to_take_ids = []
  
      with open(manifest, 'r') as f:
          f.readline()
  
          for i, line in enumerate(f.readlines()):
              seq_id = line.split()[0]
              seq_id = pathlib.Path(seq_id).name.split('__')[0]
  
              if seq_id in to_take_sequences:
                  to_take_ids.append(i)
  
      print(f'Took {len(to_take_ids)} ids')
      return set(to_take_ids)
  
  
  def get_args():
      import argparse
  
      parser = argparse.ArgumentParser()
      parser.add_argument('--asr-transcript', type=str,
                          help='Path to the transcript file.')
  
      parser.add_argument('--manifest', required=True)
      parser.add_argument('--prompts-description', required=True)
  
      parser.add_argument('--cut-id', action='store_true',
                          help='Whether cut the first token (typically a seq id)')
      parser.add_argument('--cut-tail', action='store_true',
                          help='Whether cut the last token (typically a speaker id)')
      parser.add_argument('--debug', action='store_true')
  
      args = parser.parse_args()
  
      return args
  
  
  def get_self_bleu(utterances, averaging_mode, weights):
      self_bleu = []
  
      for i in range(len(utterances)):
          hypo = utterances[i]
          rest = utterances[:i] + utterances[i+1:]
  
          self_bleu.append(sentence_bleu(rest, hypo, weights,
                           no_length_penalty=True, averaging_mode=averaging_mode))
  
      return self_bleu
  
  
  def get_self_bleu2_arithmetic(utterances):
      weights = (0.5, 0.5)  # equal weight for unigrams and bigrams
      return get_self_bleu(utterances, averaging_mode='arithmetic', weights=weights)
  
  
  def get_self_bleu2_geometric(utterances):
      weights = (0.5, 0.5)
      return get_self_bleu(utterances, averaging_mode='geometric', weights=weights)
  
  
  def get_auto_bleu2_arithmetic(utterances):
      weights = (0.5, 0.5)
      return [auto_bleu(u, mean_mode='arithmetic', weights=weights) for u in utterances]
  
  
  def get_auto_bleu2_geometric(utterances):
      weights = (0.5, 0.5)
      return [auto_bleu(u, mean_mode='geometric', weights=weights) for u in utterances]
  
  
  def get_auto_bleu3_geometric(utterances):
      weights = (1./3, 1./3, 1./3)
      return [auto_bleu(u, mean_mode='geometric', weights=weights) for u in utterances]
  
  
  def get_auto_bleu3_arithmetic(utterances):
      weights = (1./3, 1./3, 1./3)
      return [auto_bleu(u, mean_mode='arithmetic', weights=weights) for u in utterances]
  
  
  def get_self_bleu3_arithmetic(utterances):
      weights = (1./3, 1./3, 1./3)
      return get_self_bleu(utterances, averaging_mode='arithmetic', weights=weights)
  
  
  def get_self_bleu3_geometric(utterances):
      weights = (1./3, 1./3, 1./3)
      return get_self_bleu(utterances, averaging_mode='geometric', weights=weights)
  
  
  def auto_bleu(sentence, weights, mean_mode='arithmetic'):
      if len(sentence) <= 1:
          return 0
  
      N = len(weights)
  
      bleu_n = np.zeros([N])
      for n in range(N):
          targ_ngrams = list(nltk.ngrams(sentence, n+1))
          for p in range(len(targ_ngrams)):
              left = sentence[:p]
              right = sentence[(p+n+1):]
              rest_ngrams = list(nltk.ngrams(left, n+1)) + \
                  list(nltk.ngrams(right, n+1))
              # compute the nb of matching ngrams
              bleu_n[n] += targ_ngrams[p] in rest_ngrams
          bleu_n[n] /= len(targ_ngrams)  # average them to get a proportion
  
      weights = np.array(weights)
      if mean_mode == 'arithmetic':
          return (bleu_n * weights).sum()
      elif mean_mode == 'geometric':
          return (bleu_n ** weights).prod()
      else:
          raise ValueError(f'Unknown agggregation mode {mean_mode}')
  
  
  def main():
      from multiprocessing import Pool
  
      args = get_args()
      target_ids = get_target_sequences(args.manifest, args.prompts_description)
  
      with open(args.asr_transcript, 'r') as fin:
          lines = fin.readlines()
  
      terms = [x.strip().split() for x in lines]
      filtered = []
      for term in terms:
          line_id = int(term[-1].split('-')[1][:-1])
          if line_id in target_ids:
              filtered.append(term)
      terms = filtered
  
      if args.cut_id:
          terms = [x[1:] for x in terms]
      if args.cut_tail:
          terms = [x[:-1] for x in terms]
  
      if args.debug:
          terms = terms[:10]
  
      tasks = [
          ('Self-BLEU2-arithmetic', get_self_bleu2_arithmetic),
          ('Self-BLEU2-geometric', get_self_bleu2_geometric),
          ('Auto-BLEU2-arithmetic', get_auto_bleu2_arithmetic),
          ('Auto-BLEU2-geometric', get_auto_bleu2_geometric),
  
          ('Self-BLEU3-arithmetic', get_self_bleu3_arithmetic),
          ('Self-BLEU3-geometric', get_self_bleu3_geometric),
          ('Auto-BLEU3-arithmetic', get_auto_bleu3_arithmetic),
          ('Auto-BLEU3-geometric', get_auto_bleu3_geometric),
      ]
  
      n_processes = min(16, len(tasks))
      with Pool(n_processes) as pool:
          metrics = pool.map(run_f, [(t[1], terms) for t in tasks])
  
      for (metric_name, _), metric in zip(tasks, metrics):
          metric, sem = np.mean(metric), np.std(metric) / np.sqrt(len(metric))
  
          metric, sem = [
              round(100 * x, 2) for x in [metric, sem]
          ]
  
          print(f'{metric_name} {metric} +- {sem}')
  
  
  def run_f(task_params):
      f, terms = task_params
      return f(terms)
  
  
  if __name__ == '__main__':
      # NLTK produces warnings
      warnings.filterwarnings("ignore")
  
      main()
  