text: import logging
  import os
  from argparse import ArgumentParser
  from ast import literal_eval
  from types import SimpleNamespace
  from typing import List
  
  from robustnessgym import Dataset, Spacy, CachedOperation
  from robustnessgym.core.constants import CACHEDOPS
  from robustnessgym.core.tools import strings_as_json
  from robustnessgym.logging.utils import set_logging_level
  from spacy import load
  from spacy.attrs import DEP, IS_ALPHA, IS_PUNCT, IS_STOP, LEMMA, LOWER, TAG, SENT_END, \
      SENT_START, ORTH, POS, ENT_IOB
  from spacy.tokens import Doc
  
  from align import BertscoreAligner, NGramAligner, StaticEmbeddingAligner
  from utils import clean_text
  
  set_logging_level('critical')
  logger = logging.getLogger(__name__)
  logger.setLevel(logging.CRITICAL)
  
  
  def _spacy_encode(self, x):
      arr = x.to_array(
          [DEP, IS_ALPHA, IS_PUNCT, IS_STOP, LEMMA, LOWER, TAG, SENT_END, SENT_START,
           ORTH, POS, ENT_IOB])
      return {
          'arr': arr.flatten(),
          'shape': list(arr.shape),
          'words': [t.text for t in x]
      }
  
  
  def _spacy_decode(self, x):
      doc = Doc(self.nlp.vocab, words=x['words'])
      return doc.from_array(
          [DEP, IS_ALPHA, IS_PUNCT, IS_STOP, LEMMA, LOWER,
           TAG, SENT_END, SENT_START, ORTH, POS, ENT_IOB],
          x['arr'].reshape(x['shape'])
      )
  
  
  Spacy.encode = _spacy_encode
  Spacy.decode = _spacy_decode
  
  
  class AlignerCap(CachedOperation):
      def __init__(
              self,
              aligner,
              spacy,
              **kwargs,
      ):
          super(AlignerCap, self).__init__(**kwargs)
          self.spacy = spacy
          self.aligner = aligner
  
      @classmethod
      def encode(cls, x):
          # Convert to built-in types from np.int / np.float
          return super(AlignerCap, cls).encode([
              {str(k): [(int(t[0]), float(t[1])) for t in v] for k, v in d.items()}
              for d in x
          ])
  
      @classmethod
      def decode(cls, x):
          x = super(AlignerCap, cls).decode(x)
          x = [{literal_eval(k): v for k, v in d.items()} for d in x]
          return x
  
      def apply(self, batch, columns, *args, **kwargs):
          # Run the aligner on the first example of the batch
          return [
              self.aligner.align(
                  self.spacy.retrieve(batch, columns[0])[0],
                  [self.spacy.retrieve(batch, col)[0] for col in columns[1:]]
                  if len(columns) > 2 else
                  [self.spacy.retrieve(batch, columns[1])[0]],
              )
          ]
  
  
  class BertscoreAlignerCap(AlignerCap):
      def __init__(
              self,
              threshold: float,
              top_k: int,
              spacy,
      ):
          super(BertscoreAlignerCap, self).__init__(
              aligner=BertscoreAligner(threshold=threshold, top_k=top_k),
              spacy=spacy,
              threshold=threshold,
              top_k=top_k,
          )
  
  
  class NGramAlignerCap(AlignerCap):
      def __init__(
              self,
              spacy,
      ):
          super(NGramAlignerCap, self).__init__(
              aligner=NGramAligner(),
              spacy=spacy
          )
  
  
  class StaticEmbeddingAlignerCap(AlignerCap):
      def __init__(
              self,
              threshold: float,
              top_k: int,
              spacy,
      ):
          super(StaticEmbeddingAlignerCap, self).__init__(
              aligner=StaticEmbeddingAligner(threshold=threshold, top_k=top_k),
              spacy=spacy,
              threshold=threshold,
              top_k=top_k,
          )
  
  
  def _run_aligners(
          dataset: Dataset,
          aligners: List[CachedOperation],
          doc_column: str,
          reference_column: str,
          summary_columns: List[str] = None,
  ):
      if not summary_columns:
          summary_columns = []
  
      to_columns = []
      if reference_column is not None:
          to_columns.append(reference_column)
      to_columns.extend(summary_columns)
  
      for aligner in aligners:
  
          # Run the aligner on (document, summary) pairs
  
          dataset = aligner(
              dataset,
              [doc_column] + to_columns,
              # Must use `batch_size = 1`
              batch_size=1,
          )
  
          if reference_column is not None and len(summary_columns):
              # Run the aligner on (reference, summary) pairs
              dataset = aligner(
                  dataset,
                  [reference_column] + summary_columns,
                  # Must use `batch_size = 1`
                  batch_size=1,
              )
  
          if len(to_columns) > 1:
              # Instead of having one column for (document, summary) comparisons, split
              # off into (1 + |summary_columns|) total columns, one for each comparison
  
              # Retrieve the (document, summary) column
              doc_summary_column = aligner.retrieve(
                  dataset[:],
                  [doc_column] + to_columns,
              )[tuple([doc_column] + to_columns)]
  
              for i, col in enumerate(to_columns):
                  # Add as a new column after encoding with the aligner's `encode` method
                  dataset.add_column(
                      column=str(aligner.identifier(columns=[doc_column, col])),
                      values=[aligner.encode([row[i]]) for row in doc_summary_column],
                  )
  
              # Remove the (document, summary) column
              dataset.remove_column(
                  str(
                      aligner.identifier(
                          columns=[doc_column] + to_columns
                      )
                  )
              )
              del dataset.interactions[CACHEDOPS].history[
                  (
                      aligner.identifier,
                      strings_as_json(
                          strings=[doc_column] + to_columns
                      )
                  )
              ]
  
          if reference_column is not None and len(summary_columns) > 1:
              # Instead of having one column for (reference, summary) comparisons, split
              # off into (|summary_columns|) total columns, one for each comparison
  
              # Retrieve the (reference, summary) column
              reference_summary_column = aligner.retrieve(
                  dataset[:],
                  [reference_column] + summary_columns,
              )[tuple([reference_column] + summary_columns)]
  
              for i, col in enumerate(summary_columns):
                  # Add as a new column
                  dataset.add_column(
                      column=str(aligner.identifier(columns=[reference_column, col])),
                      values=[
                          aligner.encode([row[i]]) for row in reference_summary_column
                      ]
                  )
  
              # Remove the (reference, summary) column
              dataset.remove_column(
                  str(
                      aligner.identifier(
                          columns=[reference_column] + summary_columns
                      )
                  )
              )
              del dataset.interactions[CACHEDOPS].history[
                  (
                      aligner.identifier,
                      strings_as_json(
                          strings=[reference_column] + summary_columns
                      )
                  )
              ]
  
      return dataset
  
  
  def deanonymize_dataset(
          rg_path: str,
          standardized_dataset: Dataset,
          processed_dataset_path: str = None,
          n_samples: int = None,
  
  ):
      """Take an anonymized dataset and add back the original dataset columns."""
      assert processed_dataset_path is not None, \
          "Please specify a path to save the dataset."
  
      # Load the dataset
      dataset = Dataset.load_from_disk(rg_path)
  
      if n_samples:
          dataset.set_visible_rows(list(range(n_samples)))
          standardized_dataset.set_visible_rows(list(range(n_samples)))
  
      text_columns = []
  
      # Add columns from the standardized dataset
      dataset.add_column('document', standardized_dataset['document'])
      text_columns.append('document')
  
      if 'summary:reference' in standardized_dataset.column_names:
          dataset.add_column('summary:reference', standardized_dataset['summary:reference'])
          text_columns.append('summary:reference')
  
      # Preprocessing all the text columns
      dataset = dataset.update(
          lambda x:
              {
                  f'preprocessed_{k}': x[k] if args.no_clean else clean_text(x[k])
                  for k in text_columns
              }
      )
  
      # Run the Spacy pipeline on all preprocessed text columns
      try:
          nlp = load('en_core_web_lg')
      except OSError:
          nlp = load('en_core_web_sm')
  
      nlp.add_pipe('sentencizer', before="parser")
      spacy = Spacy(nlp=nlp)
      dataset = spacy(
          dataset,
          [f'preprocessed_{col}' for col in text_columns],
          batch_size=100,
      )
  
      # Directly save to disk
      dataset.save_to_disk(processed_dataset_path)
  
      return dataset
  
  
  def run_workflow(
          jsonl_path: str = None,
          dataset: Dataset = None,
          doc_column: str = None,
          reference_column: str = None,
          summary_columns: List[str] = None,
          bert_aligner_threshold: float = 0.5,
          bert_aligner_top_k: int = 3,
          embedding_aligner_threshold: float = 0.5,
          embedding_aligner_top_k: int = 3,
          processed_dataset_path: str = None,
          n_samples: int = None,
          anonymize: bool = False,
  ):
      assert (jsonl_path is None) != (dataset is None), \
          "One of `jsonl_path` and `dataset` must be specified."
      assert processed_dataset_path is not None, \
          "Please specify a path to save the dataset."
  
      # Load the dataset
      if jsonl_path is not None:
          dataset = Dataset.from_jsonl(jsonl_path)
  
      if doc_column is None:
          # Assume `doc_column` is called "document"
          doc_column = 'document'
          assert doc_column in dataset.column_names, \
              f"`doc_column={doc_column}` is not a column in dataset."
          print("Assuming `doc_column` is called 'document'.")
  
      if reference_column is None:
          # Assume `reference_column` is called "summary:reference"
          reference_column = 'summary:reference'
          print("Assuming `reference_column` is called 'summary:reference'.")
          if reference_column not in dataset.column_names:
              print("No reference summary loaded")
              reference_column = None
  
      if summary_columns is None or len(summary_columns) == 0:
          # Assume `summary_columns` are prefixed by "summary:"
          summary_columns = []
          for col in dataset.column_names:
              if col.startswith("summary:") and col != "summary:reference":
                  summary_columns.append(col)
          print(f"Reading summary columns from dataset. Found {summary_columns}.")
  
      if len(summary_columns) == 0 and reference_column is None:
          raise ValueError("At least one summary is required")
  
      # Set visible rows to restrict to the first `n_samples`
      if n_samples:
          dataset.set_visible_rows(list(range(n_samples)))
  
      # Combine the text columns into one list
      text_columns = [doc_column] + ([reference_column] if reference_column else []) + summary_columns
  
      # Preprocessing all the text columns
      dataset = dataset.update(
          lambda x: {
              f'preprocessed_{k}': x[k] if args.no_clean else clean_text(x[k])
              for k in text_columns
          }
      )
  
      # Run the Spacy pipeline on all preprocessed text columns
      try:
          nlp = load('en_core_web_lg')
      except OSError:
          nlp = None
  
      if nlp is None:
          raise OSError('Missing spaCy model "en_core_web_lg". Please run "python -m spacy download en_core_web_lg"')
  
      nlp.add_pipe('sentencizer', before="parser")
      spacy = Spacy(nlp=nlp)
      dataset = spacy(
          dataset,
          [f'preprocessed_{col}' for col in text_columns],
          batch_size=100,
      )
  
      # Run the 3 align pipelines
      bert_aligner = BertscoreAlignerCap(
          threshold=bert_aligner_threshold,
          top_k=bert_aligner_top_k,
          spacy=spacy,
      )
  
      embedding_aligner = StaticEmbeddingAlignerCap(
          threshold=embedding_aligner_threshold,
          top_k=embedding_aligner_top_k,
          spacy=spacy,
      )
  
      ngram_aligner = NGramAlignerCap(
          spacy=spacy,
      )
  
      dataset = _run_aligners(
          dataset=dataset,
          aligners=[bert_aligner, embedding_aligner, ngram_aligner],
          doc_column=f'preprocessed_{doc_column}',
          reference_column=f'preprocessed_{reference_column}' if reference_column else None,
          summary_columns=[f'preprocessed_{col}' for col in summary_columns],
      )
  
      # Save the dataset
      if anonymize:
          # Remove certain columns to anonymize and save to disk
          for col in [doc_column, reference_column]:
              if col is not None:
                  dataset.remove_column(col)
                  dataset.remove_column(f'preprocessed_{col}')
                  dataset.remove_column(
                      str(spacy.identifier(columns=[f'preprocessed_{col}']))
                  )
                  del dataset.interactions[CACHEDOPS].history[
                      (spacy.identifier, f'preprocessed_{col}')
                  ]
          dataset.save_to_disk(f'{processed_dataset_path}.anonymized')
      else:
          # Directly save to disk
          dataset.save_to_disk(processed_dataset_path)
  
      return dataset
  
  
  def parse_prediction_jsonl_name(prediction_jsonl: str):
      """Parse the name of the prediction_jsonl to extract useful information."""
      # Analyze the name of the prediction_jsonl
      filename = prediction_jsonl.split("/")[-1]
  
      # Check that the filename ends with `.results.anonymized`
      if filename.endswith(".results.anonymized"):
          # Fmt: <model>-<training dataset>.<eval dataset>.<eval split>.results.anonymized
  
          # Split using a period
          model_train_dataset, eval_dataset, eval_split = filename.split(".")[:-2]
          model, train_dataset = model_train_dataset.split("-")
  
          return SimpleNamespace(
              model_train_dataset=model_train_dataset,
              model=model,
              train_dataset=train_dataset,
              eval_dataset=eval_dataset,
              eval_split=eval_split,
          )
  
      raise NotImplementedError(
          "Prediction files must be named "
          "<model>-<training dataset>.<eval dataset>.<eval split>.results.anonymized. "
          f"Please rename the prediction file {filename} and run again."
      )
  
  
  def join_predictions(
          dataset_jsonl: str = None,
          prediction_jsonls: str = None,
          save_jsonl_path: str = None,
  ):
      """Join predictions with a dataset."""
      assert prediction_jsonls is not None, "Must have prediction jsonl files."
  
      print(
          "> Warning: please inspect the prediction .jsonl file to make sure that "
          "predictions are aligned with the examples in the dataset. "
          "Use `get_dataset` to inspect the dataset."
      )
  
      # Load the dataset
      dataset = get_dataset(dataset_jsonl=dataset_jsonl)
  
      # Parse names of all prediction files to get metadata
      metadata = [
          parse_prediction_jsonl_name(prediction_jsonl)
          for prediction_jsonl in prediction_jsonls
      ]
  
      # Load the predictions
      predictions = [
          Dataset.from_jsonl(json_path=prediction_jsonl)
          for prediction_jsonl in prediction_jsonls
      ]
  
      # Predictions for a model
      for i, prediction_data in enumerate(predictions):
          # Get metadata for i_th prediction file
          metadata_i = metadata[i]
  
          # Construct a prefix for columns added to the dataset for this prediction file
          prefix = metadata_i.model_train_dataset
  
          # Add the predictions column to the dataset
          for col in prediction_data.column_names:
              # Don't add the indexing information since the dataset has it already
              if col not in {'index', 'ix', 'id'}:
                  # `add_column` will automatically ensure that column lengths match
                  if col == 'decoded':  # rename decoded to summary
                      dataset.add_column(f'summary:{prefix}', prediction_data[col])
                  else:
                      dataset.add_column(f'{prefix}:{col}', prediction_data[col])
  
      # Save the dataset back to disk
      if save_jsonl_path:
          dataset.to_jsonl(save_jsonl_path)
      else:
          print("Dataset with predictions was not saved since `save_jsonl_path` "
                "was not specified.")
  
      return dataset
  
  
  def standardize_dataset(
          dataset_name: str = None,
          dataset_version: str = None,
          dataset_split: str = 'test',
          dataset_jsonl: str = None,
          doc_column: str = None,
          reference_column: str = None,
          save_jsonl_path: str = None,
          no_save: bool = False,
  ):
      """Load a dataset from Huggingface and dump it to disk."""
      # Load the dataset from Huggingface
      dataset = get_dataset(
          dataset_name=dataset_name,
          dataset_version=dataset_version,
          dataset_split=dataset_split,
          dataset_jsonl=dataset_jsonl,
      )
  
      if doc_column is None:
          if reference_column is not None:
              raise ValueError("You must specify `doc_column` if you specify `reference_column`")
          try:
              doc_column, reference_column = {
                  'cnn_dailymail': ('article', 'highlights'),
                  'xsum': ('document', 'summary')
              }[dataset_name]
          except:
              raise NotImplementedError(
                  "Please specify `doc_column`."
              )
  
      # Rename the columns
      if doc_column != 'document':
          dataset.add_column('document', dataset[doc_column])
          dataset.remove_column(doc_column)
      dataset.add_column('summary:reference', dataset[reference_column])
      dataset.remove_column(reference_column)
  
      # Save the dataset back to disk
      if save_jsonl_path:
          dataset.to_jsonl(save_jsonl_path)
  
      elif (save_jsonl_path is None) and not no_save:
          # Auto-create a path to save the standardized dataset
          os.makedirs('preprocessing', exist_ok=True)
          if not dataset_jsonl:
              dataset.to_jsonl(
                  f'preprocessing/'
                  f'standardized_{dataset_name}_{dataset_version}_{dataset_split}.jsonl'
              )
          else:
              dataset.to_jsonl(
                  f'preprocessing/'
                  f'standardized_{dataset_jsonl.split("/")[-1]}'
              )
  
      return dataset
  
  
  def get_dataset(
          dataset_name: str = None,
          dataset_version: str = None,
          dataset_split: str = 'test',
          dataset_jsonl: str = None,
  ):
      """Load a dataset."""
      assert (dataset_name is not None) != (dataset_jsonl is not None), \
          "Specify one of `dataset_name` or `dataset_jsonl`."
  
      # Load the dataset
      if dataset_name is not None:
          return get_hf_dataset(dataset_name, dataset_version, dataset_split)
  
      return Dataset.from_jsonl(json_path=dataset_jsonl)
  
  
  def get_hf_dataset(name: str, version: str = None, split: str = 'test'):
      """Get dataset from Huggingface."""
      if version:
          return Dataset.load_dataset(name, version, split=split)
      return Dataset.load_dataset(name, split=split)
  
  
  if __name__ == '__main__':
      parser = ArgumentParser()
      parser.add_argument('--dataset', type=str, choices=['cnn_dailymail', 'xsum'],
                          help="Huggingface dataset name.")
      parser.add_argument('--version', type=str,
                          help="Huggingface dataset version.")
      parser.add_argument('--split', type=str, default='test',
                          help="Huggingface dataset split.")
      parser.add_argument('--dataset_jsonl', type=str,
                          help="Path to a jsonl file for the dataset.")
      parser.add_argument('--dataset_rg', type=str,
                          help="Path to a dataset stored in the Robustness Gym format. "
                               "All processed datasets are stored in this format.")
      parser.add_argument('--prediction_jsonls', nargs='+', default=[],
                          help="Path to one or more jsonl files for the predictions.")
      parser.add_argument('--save_jsonl_path', type=str,
                          help="Path to save the processed jsonl dataset.")
  
      parser.add_argument('--doc_column', type=str,
                          help="Name of the document column in the dataset.")
      parser.add_argument('--reference_column', type=str,
                          help="Name of the reference summary column in the dataset.")
      parser.add_argument('--summary_columns', nargs='+', default=[],
                          help="Name of other summary columns in/added to the dataset.")
  
      parser.add_argument('--bert_aligner_threshold', type=float, default=0.1,
                          help="Minimum threshold for BERT alignment.")
      parser.add_argument('--bert_aligner_top_k', type=int, default=10,
                          help="Top-k for BERT alignment.")
      parser.add_argument('--embedding_aligner_threshold', type=float, default=0.1,
                          help="Minimum threshold for embedding alignment.")
      parser.add_argument('--embedding_aligner_top_k', type=int, default=10,
                          help="Top-k for embedding alignment.")
      parser.add_argument('--processed_dataset_path', type=str,
                          help="Path to store the final processed dataset.")
      parser.add_argument('--n_samples', type=int,
                          help="Number of dataset samples to process.")
  
      parser.add_argument('--workflow', action='store_true', default=False,
                          help="Whether to run the preprocessing workflow.")
      parser.add_argument('--standardize', action='store_true', default=False,
                          help="Whether to standardize the dataset and save to jsonl.")
      parser.add_argument('--join_predictions', action='store_true', default=False,
                          help="Whether to add predictions to the dataset and save to "
                               "jsonl.")
      parser.add_argument('--try_it', action='store_true', default=False,
                          help="`Try it` mode is faster and runs processing on 10 "
                               "examples.")
      parser.add_argument('--deanonymize', action='store_true', default=False,
                          help="Deanonymize the dataset provided by summvis.")
      parser.add_argument('--anonymize', action='store_true', default=False,
                          help="Anonymize by removing document and reference summary "
                               "columns of the original dataset.")
      parser.add_argument('--no_clean', action='store_true', default=False,
                          help="Do not clean text (remove extraneous spaces, newlines).")
      args = parser.parse_args()
  
      if args.standardize:
          # Dump a dataset to jsonl on disk after standardizing it
          standardize_dataset(
              dataset_name=args.dataset,
              dataset_version=args.version,
              dataset_split=args.split,
              dataset_jsonl=args.dataset_jsonl,
              doc_column=args.doc_column,
              reference_column=args.reference_column,
              save_jsonl_path=args.save_jsonl_path,
          )
  
      if args.join_predictions:
          # Join the predictions with the dataset
          dataset = join_predictions(
              dataset_jsonl=args.dataset_jsonl,
              prediction_jsonls=args.prediction_jsonls,
              save_jsonl_path=args.save_jsonl_path,
          )
  
      if args.workflow:
          # Run the processing workflow
          dataset = None
          # Check if `args.dataset_rg` was passed in
          if args.dataset_rg:
              # Load the dataset directly
              dataset = Dataset.load_from_disk(args.dataset_rg)
  
          run_workflow(
              jsonl_path=args.dataset_jsonl,
              dataset=dataset,
              doc_column=args.doc_column,
              reference_column=args.reference_column,
              summary_columns=args.summary_columns,
              bert_aligner_threshold=args.bert_aligner_threshold,
              bert_aligner_top_k=args.bert_aligner_top_k,
              embedding_aligner_threshold=args.embedding_aligner_threshold,
              embedding_aligner_top_k=args.embedding_aligner_top_k,
              processed_dataset_path=args.processed_dataset_path,
              n_samples=args.n_samples if not args.try_it else 10,
              anonymize=args.anonymize,
          )
  
      if args.deanonymize:
          # Deanonymize an anonymized dataset
          # Check if `args.dataset_rg` was passed in
          assert args.dataset_rg is not None, \
              "Must specify `dataset_rg` path to be deanonymized."
          assert args.dataset_rg.endswith('anonymized'), \
              "`dataset_rg` must end in 'anonymized'."
          assert (args.dataset is None) != (args.dataset_jsonl is None), \
              "`dataset_rg` points to an anonymized dataset that will be " \
              "deanonymized. Please pass in relevant arguments: either " \
              "`dataset`, `version` and `split` OR `dataset_jsonl`."
  
          # Load the standardized dataset
          standardized_dataset = standardize_dataset(
              dataset_name=args.dataset,
              dataset_version=args.version,
              dataset_split=args.split,
              dataset_jsonl=args.dataset_jsonl,
              doc_column=args.doc_column,
              reference_column=args.reference_column,
              no_save=True,
          )
          # Use it to deanonymize
          dataset = deanonymize_dataset(
              rg_path=args.dataset_rg,
              standardized_dataset=standardized_dataset,
              processed_dataset_path=args.processed_dataset_path,
              n_samples=args.n_samples if not args.try_it else 10,
          )
  