from transformers import Seq2SeqTrainingArguments
from dataclasses import dataclass, field
from typing import Optional

@dataclass
class BaselineSeq2SeqTrainingArguments(Seq2SeqTrainingArguments):
    load_states_in_eval_from_model_path: bool = field(
        default=True,
        metadata={
            "help": "In case of only using --do_eval without --do_train, use it to load the states before eval."
            "keep this to true, it causes otherwise an issue with huggingface when doing only --do_eval."
        },
    )
    top_p: Optional[float] = field(default=None, metadata={"help": "top_p value for nucleus (top_p) sampling."})
    temperature: float = field(default=1.0, metadata={"help": "Defines the softmax temperature before doing the sampling."})


baselines/arguments.py
allenai/tess-diffusion
""" Finetuning the library models for sequence classification on GLUE."""

import logging
import os
import random
import sys
from dataclasses import dataclass, field

import datasets
import numpy as np
from datasets import load_dataset

import evaluate
import transformers
from transformers import (
    AutoConfig,
    AutoModelForSequenceClassification,
    AutoTokenizer,
    DataCollatorWithPadding,
    EvalPrediction,
    HfArgumentParser,
    PretrainedConfig,
    Trainer,
    default_data_collator,
    set_seed,
)
from transformers.trainer_utils import get_last_checkpoint
from transformers.utils import check_min_version, send_example_telemetry
from transformers.utils.versions import require_version
from sdlm.arguments import DataTrainingArguments as BaseDataTrainingArguments
from sdlm.arguments import ModelArguments as BaseModelArguments
from sdlm.arguments import TrainingArguments
from sdlm.data.data_utils import split_glue

# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
check_min_version("4.25.0")

require_version("datasets>=1.8.0")

task_to_keys = {
    "cola": ("sentence", None),
    "mnli": ("premise", "hypothesis"),
    "mrpc": ("sentence1", "sentence2"),
    "qnli": ("question", "sentence"),
    "qqp": ("question1", "question2"),
    "rte": ("sentence1", "sentence2"),
    "sst2": ("sentence", None),
    "stsb": ("sentence1", "sentence2"),
    "wnli": ("sentence1", "sentence2"),
}

task_to_metric = {
    "cola": "matthews_correlation",
    "mnli": "accuracy",
    "mrpc": "combined_score",
    "qnli": "accuracy",
    "qqp": "combined_score",
    "rte": "accuracy",
    "sst2": "accuracy",
    "stsb": "combined_score",
    "wnli": "accuracy"
}

logger = logging.getLogger(__name__)


@dataclass
class DataTrainingArguments(BaseDataTrainingArguments):
    """
    Arguments pertaining to what data we are going to input our model for training and eval.

    Using `HfArgumentParser` we can turn this class
    into argparse arguments to be able to specify them on
    the command line.
    """

    def __post_init__(self):
        self.dataset_name = self.dataset_name.lower()
        if self.dataset_name not in task_to_keys.keys():
            raise ValueError("Unknown task, you should pick one in " + ",".join(task_to_keys.keys()))


@dataclass
class ModelArguments(BaseModelArguments):
    """
    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
    """

    ignore_mismatched_sizes: bool = field(
        default=False,
        metadata={"help": "Will enable to load a pretrained model whose head dimensions are different."},
    )


def main():
    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        # If we pass only one argument to the script and it's the path to a json file,
        # let's parse it to get our arguments.
        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    if training_args.checkpoint_best_model:
        training_args.metric_for_best_model = task_to_metric[data_args.dataset_name]

    # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
    # information sent is the one passed as arguments along with your Python/PyTorch versions.
    send_example_telemetry("run_glue", model_args, data_args)

    # Setup logging
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        handlers=[logging.StreamHandler(sys.stdout)],
    )

    log_level = training_args.get_process_log_level()
    logger.setLevel(log_level)
    datasets.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.enable_default_handler()
    transformers.utils.logging.enable_explicit_format()

    # Log on each process the small summary:
    logger.warning(
        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
        + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
    )
    logger.info(f"Training/evaluation parameters {training_args}")

    # Detecting last checkpoint.
    last_checkpoint = None
    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
            raise ValueError(
                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
                "Use --overwrite_output_dir to overcome."
            )
        elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
            logger.info(
                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
            )

    # Set seed before initializing model.
    set_seed(training_args.seed)

    raw_datasets = load_dataset(
        "glue",
        data_args.dataset_name,
        cache_dir=model_args.cache_dir,
        use_auth_token=True if model_args.use_auth_token else None,
    )
    # Split dataset, since test sets of GLUE do not have the labels.
    raw_datasets = split_glue(raw_datasets, data_args.dataset_name, training_args.seed)

    # Labels
    is_regression = data_args.dataset_name == "stsb"
    if not is_regression:
        label_list = raw_datasets["train"].features["label"].names
        num_labels = len(label_list)
    else:
        num_labels = 1

    # Load pretrained model and tokenizer
    config = AutoConfig.from_pretrained(
        model_args.model_name_or_path,
        num_labels=num_labels,
        finetuning_task=data_args.dataset_name,
        cache_dir=model_args.cache_dir,
        revision=model_args.model_revision,
        use_auth_token=True if model_args.use_auth_token else None,
    )
    tokenizer = AutoTokenizer.from_pretrained(
        model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        use_fast=model_args.use_fast_tokenizer,
        revision=model_args.model_revision,
        use_auth_token=True if model_args.use_auth_token else None,
    )
    model = AutoModelForSequenceClassification.from_pretrained(
        model_args.model_name_or_path,
        from_tf=bool(".ckpt" in model_args.model_name_or_path),
        config=config,
        cache_dir=model_args.cache_dir,
        revision=model_args.model_revision,
        use_auth_token=True if model_args.use_auth_token else None,
        ignore_mismatched_sizes=model_args.ignore_mismatched_sizes,
    )

    # Preprocessing the raw_datasets
    sentence1_key, sentence2_key = task_to_keys[data_args.dataset_name]

    # Padding strategy
    if data_args.pad_to_max_length:
        padding = "max_length"
    else:
        # We will pad later, dynamically at batch creation, to the max sequence length in each batch
        padding = False

    # Some models have set the order of the labels to use, so let's make sure we do use it.
    label_to_id = None
    if model.config.label2id != PretrainedConfig(num_labels=num_labels).label2id and not is_regression:
        # Some have all caps in their config, some don't.
        label_name_to_id = {k.lower(): v for k, v in model.config.label2id.items()}
        if list(sorted(label_name_to_id.keys())) == list(sorted(label_list)):
            label_to_id = {i: int(label_name_to_id[label_list[i]]) for i in range(num_labels)}
        else:
            logger.warning(
                "Your model seems to have been trained with labels, but they don't match the dataset: ",
                f"model labels: {list(sorted(label_name_to_id.keys()))}, dataset labels: {list(sorted(label_list))}."
                "\nIgnoring the model labels as a result.",
            )

    if label_to_id is not None:
        model.config.label2id = label_to_id
        model.config.id2label = {id: label for label, id in config.label2id.items()}
    elif not is_regression:
        model.config.label2id = {l: i for i, l in enumerate(label_list)}
        model.config.id2label = {id: label for label, id in config.label2id.items()}

    if data_args.max_seq_length > tokenizer.model_max_length:
        logger.warning(
            f"The max_seq_length passed ({data_args.max_seq_length}) is larger than the maximum length for the"
            f"model ({tokenizer.model_max_length}). Using max_seq_length={tokenizer.model_max_length}."
        )
    max_seq_length = min(data_args.max_seq_length, tokenizer.model_max_length)

    def preprocess_function(examples):
        # Tokenize the texts
        args = (examples[sentence1_key],) if sentence2_key is None else (examples[sentence1_key], examples[sentence2_key])
        result = tokenizer(*args, padding=padding, max_length=max_seq_length, truncation=True)

        # Map labels to IDs (not necessary for GLUE tasks)
        if label_to_id is not None and "label" in examples:
            result["label"] = [(label_to_id[l] if l != -1 else -1) for l in examples["label"]]
        return result

    with training_args.main_process_first(desc="dataset map pre-processing"):
        raw_datasets = raw_datasets.map(
            preprocess_function,
            batched=True,
            load_from_cache_file=not data_args.overwrite_cache,
            desc="Running tokenizer on dataset",
        )
    if training_args.do_train:
        if "train" not in raw_datasets:
            raise ValueError("--do_train requires a train dataset")
        train_dataset = raw_datasets["train"]
        if data_args.max_train_samples is not None:
            max_train_samples = min(len(train_dataset), data_args.max_train_samples)
            train_dataset = train_dataset.select(range(max_train_samples))

    if training_args.do_eval:
        if "validation" not in raw_datasets:
            raise ValueError("--do_eval requires a validation dataset")
        eval_dataset = raw_datasets["validation"]
        if data_args.max_eval_samples is not None:
            max_eval_samples = min(len(eval_dataset), data_args.max_eval_samples)
            eval_dataset = eval_dataset.select(range(max_eval_samples))

    if training_args.do_predict:
        if "test" not in raw_datasets:
            raise ValueError("--do_predict requires a test dataset")
        predict_dataset = raw_datasets["test"]
        if data_args.max_predict_samples is not None:
            max_predict_samples = min(len(predict_dataset), data_args.max_predict_samples)
            predict_dataset = predict_dataset.select(range(max_predict_samples))

    # Log a few random samples from the training set:
    if training_args.do_train:
        for index in random.sample(range(len(train_dataset)), 3):
            logger.info(f"Sample {index} of the training set: {train_dataset[index]}.")

    # Get the metric function
    metric = evaluate.load("glue", data_args.dataset_name)

    # You can define your custom compute_metrics function. It takes an `EvalPrediction` object (a namedtuple with a
    # predictions and label_ids field) and has to return a dictionary string to float.
    def compute_metrics(p: EvalPrediction):
        preds = p.predictions[0] if isinstance(p.predictions, tuple) else p.predictions
        preds = np.squeeze(preds) if is_regression else np.argmax(preds, axis=1)
        result = metric.compute(predictions=preds, references=p.label_ids)
        if len(result) > 1:
            result["combined_score"] = np.mean(list(result.values())).item()
        return result

    # Data collator will default to DataCollatorWithPadding when the tokenizer is passed to Trainer, so we change it if
    # we already did the padding.
    if data_args.pad_to_max_length:
        data_collator = default_data_collator
    elif training_args.fp16:
        data_collator = DataCollatorWithPadding(tokenizer, pad_to_multiple_of=8)
    else:
        data_collator = None

    # Initialize our Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset if training_args.do_train else None,
        eval_dataset=eval_dataset if training_args.do_eval else None,
        compute_metrics=compute_metrics,
        tokenizer=tokenizer,
        data_collator=data_collator,
    )

    # Training
    if training_args.do_train:
        checkpoint = None
        if training_args.resume_from_checkpoint is not None:
            checkpoint = training_args.resume_from_checkpoint
        elif last_checkpoint is not None:
            checkpoint = last_checkpoint
        train_result = trainer.train(resume_from_checkpoint=checkpoint)
        metrics = train_result.metrics
        max_train_samples = data_args.max_train_samples if data_args.max_train_samples is not None else len(train_dataset)
        metrics["train_samples"] = min(max_train_samples, len(train_dataset))

        trainer.save_model()  # Saves the tokenizer too for easy upload

        trainer.log_metrics("train", metrics)
        trainer.save_metrics("train", metrics)
        trainer.save_state()

    # Evaluation
    if training_args.do_eval:
        logger.info("*** Evaluate ***")
        metrics = trainer.evaluate(eval_dataset=eval_dataset)
        max_eval_samples = data_args.max_eval_samples if data_args.max_eval_samples is not None else len(eval_dataset)
        metrics["eval_samples"] = min(max_eval_samples, len(eval_dataset))
        trainer.log_metrics("eval", metrics)
        trainer.save_metrics("eval", metrics)

    if training_args.do_predict:
        logger.info("*** Test ***")

        metrics = trainer.evaluate(eval_dataset=predict_dataset, metric_key_prefix="test")
        max_predict_samples = (
            data_args.max_predict_samples if data_args.max_predict_samples is not None else len(predict_dataset)
        )
        metrics["test_samples"] = min(max_predict_samples, len(predict_dataset))
        trainer.log_metrics("test", metrics)
        trainer.save_metrics("test", metrics)


if __name__ == "__main__":
    main()

baselines/run_glue.py
allenai/tess-diffusion
import logging
import os
import sys
from dataclasses import dataclass, field
from typing import Optional

import datasets
import numpy as np
from datasets import DatasetDict, Dataset, load_dataset
import evaluate
import transformers
from transformers import (
    AutoConfig,
    AutoModelForSeq2SeqLM,
    AutoTokenizer,
    DataCollatorForSeq2Seq,
    HfArgumentParser,
    MBartTokenizer,
    MBartTokenizerFast,
    set_seed,
)
import pdb
from transformers.trainer_utils import get_last_checkpoint
from transformers.utils import check_min_version, send_example_telemetry
from transformers.utils.versions import require_version
from trainer_seq2seq import BaselineSeq2SeqTrainer
from arguments import BaselineSeq2SeqTrainingArguments
from sdlm.data.postprocessors import postprocess_text_for_metric
from sdlm.metrics.metrics import distinct_n_grams

# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
check_min_version("4.25.0")

require_version("datasets>=1.8.0")

logger = logging.getLogger(__name__)


def extract_field(data, field_name):
    results = []
    for item in data:
        results.append(item[field_name])
    return results

def read_wikilarge(data_args):
    raw_datasets = DatasetDict()
    for split in ["train", "dev", "test"]:
        # TODO: change to f"{data_args.dataset_folder}/"
        s1s = open(f"wikilarge/s1.{split}", "r").readlines() 
        s2s = open(f"wikilarge/s2.{split}", "r").readlines()
        data = [{"original": s1, "simplification": s2} for s1, s2 in zip(s1s, s2s)]
        raw_datasets[split] = Dataset.from_list(data)
    return raw_datasets 

def read_diffuseq_datasets(data_args):
    raw_datasets = DatasetDict()
    for split in ["train", "valid", "test"]:
        dataset = load_dataset("json", data_files=f"{data_args.dataset_folder}/{split}.jsonl")["train"]
        data_split = split if split != "valid" else "dev"
        raw_datasets[data_split] = dataset
    return raw_datasets


@dataclass
class ModelArguments:
    """
    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
    """

    model_name_or_path: str = field(
        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
    )
    config_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
    )
    tokenizer_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
    )
    cache_dir: Optional[str] = field(
        default=None,
        metadata={"help": "Where to store the pretrained models downloaded from huggingface.co"},
    )
    use_fast_tokenizer: bool = field(
        default=True,
        metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
    )
    model_revision: str = field(
        default="main",
        metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."},
    )
    use_auth_token: bool = field(
        default=False,
        metadata={
            "help": (
                "Will use the token generated when running `huggingface-cli login` (necessary to use this script "
                "with private models)."
            )
        },
    )
    resize_position_embeddings: Optional[bool] = field(
        default=None,
        metadata={
            "help": (
                "Whether to automatically resize the position embeddings if `max_source_length` exceeds "
                "the model's position embeddings."
            )
        },
    )


@dataclass
class DataTrainingArguments:
    """
    Arguments pertaining to what data we are going to input our model for training and eval.
    """
    dataset_folder: str = field(default=None, metadata={"help": "The dataset folder containing the dataset."})
    lang: Optional[str] = field(default=None, metadata={"help": "Language id for summarization."})

    dataset_name: Optional[str] = field(
        default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."}
    )
    dataset_config_name: Optional[str] = field(
        default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
    )
    text_column: Optional[str] = field(
        default=None,
        metadata={"help": "The name of the column in the datasets containing the full texts (for summarization)."},
    )
    summary_column: Optional[str] = field(
        default=None,
        metadata={"help": "The name of the column in the datasets containing the summaries (for summarization)."},
    )
    train_file: Optional[str] = field(
        default=None, metadata={"help": "The input training data file (a jsonlines or csv file)."}
    )
    validation_file: Optional[str] = field(
        default=None,
        metadata={
            "help": ("An optional input evaluation data file to evaluate the metrics (rouge) on (a jsonlines or csv file).")
        },
    )
    test_file: Optional[str] = field(
        default=None,
        metadata={"help": "An optional input test data file to evaluate the metrics (rouge) on (a jsonlines or csv file)."},
    )
    overwrite_cache: bool = field(default=False, metadata={"help": "Overwrite the cached training and evaluation sets"})
    preprocessing_num_workers: Optional[int] = field(
        default=None,
        metadata={"help": "The number of processes to use for the preprocessing."},
    )
    max_source_length: Optional[int] = field(
        default=1024,
        metadata={
            "help": (
                "The maximum total input sequence length after tokenization. Sequences longer "
                "than this will be truncated, sequences shorter will be padded."
            )
        },
    )
    max_target_length: Optional[int] = field(
        default=128,
        metadata={
            "help": (
                "The maximum total sequence length for target text after tokenization. Sequences longer "
                "than this will be truncated, sequences shorter will be padded."
            )
        },
    )
    val_max_target_length: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "The maximum total sequence length for validation target text after tokenization. Sequences longer "
                "than this will be truncated, sequences shorter will be padded. Will default to `max_target_length`."
                "This argument is also used to override the ``max_length`` param of ``model.generate``, which is used "
                "during ``evaluate`` and ``predict``."
            )
        },
    )
    pad_to_max_length: bool = field(
        default=False,
        metadata={
            "help": (
                "Whether to pad all samples to model maximum sentence length. "
                "If False, will pad the samples dynamically when batching to the maximum length in the batch. More "
                "efficient on GPU but very bad for TPU."
            )
        },
    )
    max_train_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "For debugging purposes or quicker training, truncate the number of training examples to this "
                "value if set."
            )
        },
    )
    max_eval_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
                "value if set."
            )
        },
    )
    max_predict_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "For debugging purposes or quicker training, truncate the number of prediction examples to this "
                "value if set."
            )
        },
    )
    num_beams: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "Number of beams to use for evaluation. This argument will be passed to ``model.generate``, "
                "which is used during ``evaluate`` and ``predict``."
            )
        },
    )
    ignore_pad_token_for_loss: bool = field(
        default=True,
        metadata={"help": "Whether to ignore the tokens corresponding to padded labels in the loss computation or not."},
    )
    source_prefix: Optional[str] = field(
        default="", metadata={"help": "A prefix to add before every source text (useful for T5 models)."}
    )

    forced_bos_token: Optional[str] = field(
        default=None,
        metadata={
            "help": (
                "The token to force as the first generated token after the decoder_start_token_id."
                "Useful for multilingual models like mBART where the first generated token"
                "needs to be the target language token (Usually it is the target language token)"
            )
        },
    )

    def __post_init__(self):
        if self.dataset_name is None:
            raise ValueError("Need a dataset_name")
        if self.val_max_target_length is None:
            self.val_max_target_length = self.max_target_length

simplification_name_mapping = {
    "wikilarge": ("original", "simplification"),
    "wiki_alignment": ("src", "trg"),
    "qqp": ("src", "trg"),
    "qg": ("src", "trg"),
    "cc": ("src", "trg")
}


def main():
    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, BaselineSeq2SeqTrainingArguments))
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
    # information sent is the one passed as arguments along with your Python/PyTorch versions.
    send_example_telemetry("run_summarization", model_args, data_args)

    # Setup logging
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        handlers=[logging.StreamHandler(sys.stdout)],
    )
    log_level = training_args.get_process_log_level()
    logger.setLevel(log_level)
    datasets.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.enable_default_handler()
    transformers.utils.logging.enable_explicit_format()

    # Log on each process the small summary:
    logger.warning(
        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
        + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
    )
    logger.info(f"Training/evaluation parameters {training_args}")

    if data_args.source_prefix is None and model_args.model_name_or_path in [
        "t5-small",
        "t5-base",
        "t5-large",
        "t5-3b",
        "t5-11b",
    ]:
        logger.warning(
            "You're running a t5 model but didn't provide a source prefix, which is the expected, e.g. with "
            "`--source_prefix 'summarize: ' `"
        )

    # Detecting last checkpoint.
    last_checkpoint = None
    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
            raise ValueError(
                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
                "Use --overwrite_output_dir to overcome."
            )
        elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
            logger.info(
                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
            )

    # Set seed before initializing model.
    set_seed(training_args.seed)

    if data_args.dataset_name == "wikilarge":
        raw_datasets = read_wikilarge(data_args)
    elif data_args.dataset_name in ["wiki_alignment", "qqp", "qg", "cc"]:
        raw_datasets = read_diffuseq_datasets(data_args)

    train_dataset = raw_datasets["train"]
    eval_dataset = raw_datasets["dev"]
    predict_dataset = raw_datasets["test"]
    column_names = train_dataset.column_names

    config = AutoConfig.from_pretrained(
        model_args.config_name if model_args.config_name else model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        revision=model_args.model_revision,
        use_auth_token=True if model_args.use_auth_token else None,
    )
    tokenizer = AutoTokenizer.from_pretrained(
        model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        use_fast=model_args.use_fast_tokenizer,
        revision=model_args.model_revision,
        use_auth_token=True if model_args.use_auth_token else None,
    )
    model = AutoModelForSeq2SeqLM.from_pretrained(
        model_args.model_name_or_path,
        from_tf=bool(".ckpt" in model_args.model_name_or_path),
        config=config,
        cache_dir=model_args.cache_dir,
        revision=model_args.model_revision,
        use_auth_token=True if model_args.use_auth_token else None,
    )

    # We resize the embeddings only when necessary to avoid index errors. If you are creating a model from scratch
    # on a small vocab and want a smaller embedding size, remove this test.
    embedding_size = model.get_input_embeddings().weight.shape[0]
    if len(tokenizer) > embedding_size:
        model.resize_token_embeddings(len(tokenizer))

    if model.config.decoder_start_token_id is None and isinstance(tokenizer, (MBartTokenizer, MBartTokenizerFast)):
        if isinstance(tokenizer, MBartTokenizer):
            model.config.decoder_start_token_id = tokenizer.lang_code_to_id[data_args.lang]
        else:
            model.config.decoder_start_token_id = tokenizer.convert_tokens_to_ids(data_args.lang)

    if model.config.decoder_start_token_id is None:
        raise ValueError("Make sure that `config.decoder_start_token_id` is correctly defined")

    if (
        hasattr(model.config, "max_position_embeddings")
        and model.config.max_position_embeddings < data_args.max_source_length
    ):
        if model_args.resize_position_embeddings is None:
            logger.warning(
                "Increasing the model's number of position embedding vectors from"
                f" {model.config.max_position_embeddings} to {data_args.max_source_length}."
            )
            model.resize_position_embeddings(data_args.max_source_length)
        elif model_args.resize_position_embeddings:
            model.resize_position_embeddings(data_args.max_source_length)
        else:
            raise ValueError(
                f"`--max_source_length` is set to {data_args.max_source_length}, but the model only has"
                f" {model.config.max_position_embeddings} position encodings. Consider either reducing"
                f" `--max_source_length` to {model.config.max_position_embeddings} or to automatically resize the"
                " model's position encodings by passing `--resize_position_embeddings`."
            )

    prefix = data_args.source_prefix if data_args.source_prefix is not None else ""

    # Get the column names for input/target.
    dataset_columns = simplification_name_mapping.get(data_args.dataset_name, None)
    text_column, simplification_column = dataset_columns[0], dataset_columns[1]

    # Temporarily set max_target_length for training.
    max_target_length = data_args.max_target_length
    padding = "max_length" if data_args.pad_to_max_length else False

    if training_args.label_smoothing_factor > 0 and not hasattr(model, "prepare_decoder_input_ids_from_labels"):
        logger.warning(
            "label_smoothing is enabled but the `prepare_decoder_input_ids_from_labels` method is not defined for"
            f"`{model.__class__.__name__}`. This will lead to loss being calculated twice and will take up more memory"
        )

    def preprocess_function(examples):
        # remove pairs where at least one record is None

        inputs, targets = [], []
        for i in range(len(examples[text_column])):
            if examples[text_column][i] and examples[simplification_column][i]:
                inputs.append(examples[text_column][i])
                targets.append(examples[simplification_column][i])

        inputs = [prefix + inp for inp in inputs]
        model_inputs = tokenizer(inputs, max_length=data_args.max_source_length, padding=padding, truncation=True)

        # Tokenize targets with the `text_target` keyword argument
        labels = tokenizer(text_target=targets, max_length=max_target_length, padding=padding, truncation=True)

        # If we are padding here, replace all tokenizer.pad_token_id in the labels by -100 when we want to ignore
        # padding in the loss.
        if padding == "max_length" and data_args.ignore_pad_token_for_loss:
            labels["input_ids"] = [
                [(l if l != tokenizer.pad_token_id else -100) for l in label] for label in labels["input_ids"]
            ]

        model_inputs["labels"] = labels["input_ids"]
        return model_inputs

    if training_args.do_train:
        if data_args.max_train_samples is not None:
            max_train_samples = min(len(train_dataset), data_args.max_train_samples)
            train_dataset = train_dataset.select(range(max_train_samples))
        with training_args.main_process_first(desc="train dataset map pre-processing"):
            train_dataset = train_dataset.map(
                preprocess_function,
                batched=True,
                num_proc=data_args.preprocessing_num_workers,
                remove_columns=column_names,
                load_from_cache_file=not data_args.overwrite_cache,
                desc="Running tokenizer on train dataset",
            )

    if training_args.do_eval:
        max_target_length = data_args.val_max_target_length
        if data_args.max_eval_samples is not None:
            max_eval_samples = min(len(eval_dataset), data_args.max_eval_samples)
            eval_dataset = eval_dataset.select(range(max_eval_samples))
        with training_args.main_process_first(desc="validation dataset map pre-processing"):
            eval_dataset = eval_dataset.map(
                preprocess_function,
                batched=True,
                num_proc=data_args.preprocessing_num_workers,
                remove_columns=column_names,
                load_from_cache_file=not data_args.overwrite_cache,
                desc="Running tokenizer on validation dataset",
            )
        eval_input_ids = extract_field(eval_dataset, "input_ids")

    if training_args.do_predict:
        max_target_length = data_args.val_max_target_length
        if data_args.max_predict_samples is not None:
            max_predict_samples = min(len(predict_dataset), data_args.max_predict_samples)
            predict_dataset = predict_dataset.select(range(max_predict_samples))
        with training_args.main_process_first(desc="prediction dataset map pre-processing"):
            predict_dataset = predict_dataset.map(
                preprocess_function,
                batched=True,
                num_proc=data_args.preprocessing_num_workers,
                remove_columns=column_names,
                load_from_cache_file=not data_args.overwrite_cache,
                desc="Running tokenizer on prediction dataset",
            )
        predict_input_ids = extract_field(predict_dataset, "input_ids")

    # Data collator
    label_pad_token_id = -100 if data_args.ignore_pad_token_for_loss else tokenizer.pad_token_id
    data_collator = DataCollatorForSeq2Seq(
        tokenizer,
        model=model,
        label_pad_token_id=label_pad_token_id,
        pad_to_multiple_of=8 if training_args.fp16 else None,
    )

    # Metric
    eval_metrics = {
        "sari": evaluate.load("sari"),
        "bleu": evaluate.load("bleu"),
        "bertscore": evaluate.load("bertscore"),
        "bertscore_them": evaluate.load("bertscore"),
        "rouge": evaluate.load("rouge"),
        "dist":  distinct_n_grams
    }

    def compute_metrics(eval_preds, split="eval"):
        preds, labels = eval_preds
        if isinstance(preds, tuple):
            preds = preds[0]
        decoded_preds_original = tokenizer.batch_decode(preds, skip_special_tokens=True)
        if data_args.ignore_pad_token_for_loss:
            # Replace -100 in the labels as we can't decode them.
            labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
        decoded_labels_original = tokenizer.batch_decode(labels, skip_special_tokens=True)
        if split == "eval":
            sources = tokenizer.batch_decode(eval_input_ids, skip_special_tokens=True)
        else:
            sources = tokenizer.batch_decode(predict_input_ids, skip_special_tokens=True)
        metrics = {}
        for metric_name, metric in eval_metrics.items(): 
            scale = 100 if metric_name != "sari" else 1
            if metric_name == "sari":
                decoded_preds, decoded_labels, sources = postprocess_text_for_metric("sari", decoded_preds_original, decoded_labels_original, sources)
                decoded_labels = [[decoded_label] for decoded_label in decoded_labels]
                key_metrics = metric.compute(sources=sources, predictions=decoded_preds, references=decoded_labels)
            elif metric_name == "bleu":
                decoded_preds, decoded_labels = postprocess_text_for_metric("bleu", decoded_preds_original, decoded_labels_original)
                key_metrics = {"bleu": metric.compute(predictions=decoded_preds, references=decoded_labels)["bleu"]}
            elif metric_name == "bertscore":
                decoded_preds, decoded_labels = postprocess_text_for_metric("bertscore", decoded_preds_original, decoded_labels_original)
                key_metrics = {"bert_score": np.mean(metric.compute(predictions=decoded_preds, references=decoded_labels,  lang="en")['f1'])}
            elif metric_name == "bertscore_them":
                decoded_preds, decoded_labels = postprocess_text_for_metric("bertscore_them", decoded_preds_original, decoded_labels_original)
                key_metrics = {"bert_score_them": np.mean(metric.compute(predictions=decoded_preds, references=decoded_labels, model_type='microsoft/deberta-xlarge-mnli', lang="en")['f1'])}
            elif metric_name == "rouge":
                decoded_preds, decoded_labels = postprocess_text_for_metric("rouge", decoded_preds_original, decoded_labels_original)
                key_metrics = metric.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True)
            elif metric_name == "dist":
                decoded_preds = postprocess_text_for_metric("dist", decoded_preds_original)
                key_metrics = metric(decoded_preds)

            key_metrics = {k: round(v*scale, 2) for k, v in key_metrics.items()}
            metrics.update(key_metrics)
        return metrics

    # Initialize our Trainer
    trainer = BaselineSeq2SeqTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset if training_args.do_train else None,
        eval_dataset=eval_dataset if training_args.do_eval else None,
        tokenizer=tokenizer,
        data_collator=data_collator,
        compute_metrics=compute_metrics if training_args.predict_with_generate else None,
    )

    # Training
    if training_args.do_train:
        checkpoint = None
        if training_args.resume_from_checkpoint is not None:
            checkpoint = training_args.resume_from_checkpoint
        elif last_checkpoint is not None:
            checkpoint = last_checkpoint
        train_result = trainer.train(resume_from_checkpoint=checkpoint)
        trainer.save_model()  # Saves the tokenizer too for easy upload

        metrics = train_result.metrics
        max_train_samples = data_args.max_train_samples if data_args.max_train_samples is not None else len(train_dataset)
        metrics["train_samples"] = min(max_train_samples, len(train_dataset))

        trainer.log_metrics("train", metrics)
        trainer.save_metrics("train", metrics)
        trainer.save_state()

    # Evaluation
    results = {}
    max_length = (
        training_args.generation_max_length
        if training_args.generation_max_length is not None
        else data_args.val_max_target_length
    )
    num_beams = data_args.num_beams if data_args.num_beams is not None else training_args.generation_num_beams
    if training_args.do_eval:
        logger.info("*** Evaluate ***")
        metrics = trainer.evaluate(max_length=max_length, num_beams=num_beams, metric_key_prefix="eval", top_p=training_args.top_p, temperature=training_args.temperature)
        max_eval_samples = data_args.max_eval_samples if data_args.max_eval_samples is not None else len(eval_dataset)
        metrics["eval_samples"] = min(max_eval_samples, len(eval_dataset))

        trainer.log_metrics("eval", metrics)
        trainer.save_metrics("eval", metrics)

    if training_args.do_predict:
        logger.info("*** Predict ***")
        trainer.compute_metrics = lambda x: compute_metrics(x, split="test")
        predict_results = trainer.predict(
            predict_dataset, metric_key_prefix="predict", max_length=max_length, num_beams=num_beams, top_p=training_args.top_p, temperature=training_args.temperature
        )
        metrics = predict_results.metrics
        max_predict_samples = (
            data_args.max_predict_samples if data_args.max_predict_samples is not None else len(predict_dataset)
        )
        metrics["predict_samples"] = min(max_predict_samples, len(predict_dataset))

        trainer.log_metrics("predict", metrics)
        trainer.save_metrics("predict", metrics)

        if trainer.is_world_process_zero():
            if training_args.predict_with_generate:
                predictions = tokenizer.batch_decode(
                    predict_results.predictions, skip_special_tokens=True, clean_up_tokenization_spaces=True
                )
                predictions = [pred.strip() for pred in predictions]
                output_prediction_file = os.path.join(training_args.output_dir, "generated_predictions.txt")
                with open(output_prediction_file, "w") as writer:
                    writer.write("\n".join(predictions))
    return results


if __name__ == "__main__":
    main()

baselines/run_simplification.py
allenai/tess-diffusion
#!/usr/bin/env python
# coding=utf-8
# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Fine-tuning the library models for sequence to sequence.
"""
# You can also adapt this script on your own sequence to sequence task. Pointers for this are left as comments.

import logging
import os
import sys
from dataclasses import dataclass, field
from typing import Optional
import pdb
import datasets
import nltk  # Here to have a nice missing dependency error message early on
import numpy as np
from transformers.trainer_callback import TrainerState
from datasets import load_dataset

import evaluate
import transformers
from filelock import FileLock
from transformers import (
    AutoConfig,
    AutoModelForSeq2SeqLM,
    AutoTokenizer,
    DataCollatorForSeq2Seq,
    HfArgumentParser,
    MBart50Tokenizer,
    MBart50TokenizerFast,
    MBartTokenizer,
    MBartTokenizerFast,
    set_seed,
)
from transformers.trainer_utils import get_last_checkpoint
from transformers.utils import check_min_version, is_offline_mode, send_example_telemetry
from transformers.utils.versions import require_version
from trainer_seq2seq import BaselineSeq2SeqTrainer
from arguments import BaselineSeq2SeqTrainingArguments
from sdlm.data.postprocessors import postprocess_text_for_metric
GENERATION_RESULTS="generations"

# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
check_min_version("4.25.0")

require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/summarization/requirements.txt")

logger = logging.getLogger(__name__)

try:
    nltk.data.find("tokenizers/punkt")
except (LookupError, OSError):
    if is_offline_mode():
        raise LookupError("Offline mode: run this script without TRANSFORMERS_OFFLINE first to download nltk data files")
    with FileLock(".lock") as lock:
        nltk.download("punkt", quiet=True)

# A list of all multilingual tokenizer which require lang attribute.
MULTILINGUAL_TOKENIZERS = [MBartTokenizer, MBartTokenizerFast, MBart50Tokenizer, MBart50TokenizerFast]


@dataclass
class ModelArguments:
    """
    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
    """

    model_name_or_path: str = field(
        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
    )
    config_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
    )
    tokenizer_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
    )
    cache_dir: Optional[str] = field(
        default=None,
        metadata={"help": "Where to store the pretrained models downloaded from huggingface.co"},
    )
    use_fast_tokenizer: bool = field(
        default=True,
        metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
    )
    model_revision: str = field(
        default="main",
        metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."},
    )
    use_auth_token: bool = field(
        default=False,
        metadata={
            "help": (
                "Will use the token generated when running `huggingface-cli login` (necessary to use this script "
                "with private models)."
            )
        },
    )
    resize_position_embeddings: Optional[bool] = field(
        default=None,
        metadata={
            "help": (
                "Whether to automatically resize the position embeddings if `max_source_length` exceeds "
                "the model's position embeddings."
            )
        },
    )


@dataclass
class DataTrainingArguments:
    """
    Arguments pertaining to what data we are going to input our model for training and eval.
    """

    lang: Optional[str] = field(default=None, metadata={"help": "Language id for summarization."})

    dataset_name: Optional[str] = field(
        default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."}
    )
    dataset_config_name: Optional[str] = field(
        default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
    )
    text_column: Optional[str] = field(
        default=None,
        metadata={"help": "The name of the column in the datasets containing the full texts (for summarization)."},
    )
    summary_column: Optional[str] = field(
        default=None,
        metadata={"help": "The name of the column in the datasets containing the summaries (for summarization)."},
    )
    train_file: Optional[str] = field(
        default=None, metadata={"help": "The input training data file (a jsonlines or csv file)."}
    )
    validation_file: Optional[str] = field(
        default=None,
        metadata={
            "help": ("An optional input evaluation data file to evaluate the metrics (rouge) on (a jsonlines or csv file).")
        },
    )
    test_file: Optional[str] = field(
        default=None,
        metadata={"help": "An optional input test data file to evaluate the metrics (rouge) on (a jsonlines or csv file)."},
    )
    overwrite_cache: bool = field(default=False, metadata={"help": "Overwrite the cached training and evaluation sets"})
    preprocessing_num_workers: Optional[int] = field(
        default=None,
        metadata={"help": "The number of processes to use for the preprocessing."},
    )
    max_source_length: Optional[int] = field(
        default=1024,
        metadata={
            "help": (
                "The maximum total input sequence length after tokenization. Sequences longer "
                "than this will be truncated, sequences shorter will be padded."
            )
        },
    )
    max_target_length: Optional[int] = field(
        default=128,
        metadata={
            "help": (
                "The maximum total sequence length for target text after tokenization. Sequences longer "
                "than this will be truncated, sequences shorter will be padded."
            )
        },
    )
    val_max_target_length: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "The maximum total sequence length for validation target text after tokenization. Sequences longer "
                "than this will be truncated, sequences shorter will be padded. Will default to `max_target_length`."
                "This argument is also used to override the ``max_length`` param of ``model.generate``, which is used "
                "during ``evaluate`` and ``predict``."
            )
        },
    )
    pad_to_max_length: bool = field(
        default=False,
        metadata={
            "help": (
                "Whether to pad all samples to model maximum sentence length. "
                "If False, will pad the samples dynamically when batching to the maximum length in the batch. More "
                "efficient on GPU but very bad for TPU."
            )
        },
    )
    max_train_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "For debugging purposes or quicker training, truncate the number of training examples to this "
                "value if set."
            )
        },
    )
    max_eval_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
                "value if set."
            )
        },
    )
    max_predict_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "For debugging purposes or quicker training, truncate the number of prediction examples to this "
                "value if set."
            )
        },
    )
    num_beams: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "Number of beams to use for evaluation. This argument will be passed to ``model.generate``, "
                "which is used during ``evaluate`` and ``predict``."
            )
        },
    )
    ignore_pad_token_for_loss: bool = field(
        default=True,
        metadata={"help": "Whether to ignore the tokens corresponding to padded labels in the loss computation or not."},
    )
    source_prefix: Optional[str] = field(
        default="", metadata={"help": "A prefix to add before every source text (useful for T5 models)."}
    )

    forced_bos_token: Optional[str] = field(
        default=None,
        metadata={
            "help": (
                "The token to force as the first generated token after the decoder_start_token_id."
                "Useful for multilingual models like mBART where the first generated token"
                "needs to be the target language token (Usually it is the target language token)"
            )
        },
    )

    def __post_init__(self):
        if self.dataset_name is None and self.train_file is None and self.validation_file is None:
            raise ValueError("Need either a dataset name or a training/validation file.")
        else:
            if self.train_file is not None:
                extension = self.train_file.split(".")[-1]
                assert extension in ["csv", "json"], "`train_file` should be a csv or a json file."
            if self.validation_file is not None:
                extension = self.validation_file.split(".")[-1]
                assert extension in ["csv", "json"], "`validation_file` should be a csv or a json file."
        if self.val_max_target_length is None:
            self.val_max_target_length = self.max_target_length


summarization_name_mapping = {
    "amazon_reviews_multi": ("review_body", "review_title"),
    "big_patent": ("description", "abstract"),
    "cnn_dailymail": ("article", "highlights"),
    "orange_sum": ("text", "summary"),
    "pn_summary": ("article", "summary"),
    "psc": ("extract_text", "summary_text"),
    "samsum": ("dialogue", "summary"),
    "thaisum": ("body", "summary"),
    "xglue": ("news_body", "news_title"),
    "xsum": ("document", "summary"),
    "wiki_summary": ("article", "highlights"),
    "multi_news": ("document", "summary"),
}


def main():
    # See all possible arguments in src/transformers/training_args.py
    # or by passing the --help flag to this script.
    # We now keep distinct sets of args, for a cleaner separation of concerns.

    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, BaselineSeq2SeqTrainingArguments))
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        # If we pass only one argument to the script and it's the path to a json file,
        # let's parse it to get our arguments.
        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
    # information sent is the one passed as arguments along with your Python/PyTorch versions.
    send_example_telemetry("run_summarization", model_args, data_args)

    # Setup logging
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        handlers=[logging.StreamHandler(sys.stdout)],
    )
    log_level = training_args.get_process_log_level()
    logger.setLevel(log_level)
    datasets.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.enable_default_handler()
    transformers.utils.logging.enable_explicit_format()

    # Log on each process the small summary:
    logger.warning(
        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
        + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
    )
    logger.info(f"Training/evaluation parameters {training_args}")

    if data_args.source_prefix is None and model_args.model_name_or_path in [
        "t5-small",
        "t5-base",
        "t5-large",
        "t5-3b",
        "t5-11b",
    ]:
        logger.warning(
            "You're running a t5 model but didn't provide a source prefix, which is the expected, e.g. with "
            "`--source_prefix 'summarize: ' `"
        )

    # Detecting last checkpoint.
    last_checkpoint = None
    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
            raise ValueError(
                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
                "Use --overwrite_output_dir to overcome."
            )
        elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
            logger.info(
                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
            )

    # Set seed before initializing model.
    set_seed(training_args.seed)

    # Get the datasets: you can either provide your own CSV/JSON training and evaluation files (see below)
    # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/
    # (the dataset will be downloaded automatically from the datasets Hub).
    #
    # For CSV/JSON files this script will use the first column for the full texts and the second column for the
    # summaries (unless you specify column names for this with the `text_column` and `summary_column` arguments).
    #
    # In distributed training, the load_dataset function guarantee that only one local process can concurrently
    # download the dataset.
    if data_args.dataset_name is not None:
        # Downloading and loading a dataset from the hub.
        raw_datasets = load_dataset(
            data_args.dataset_name,
            data_args.dataset_config_name,
            cache_dir=model_args.cache_dir,
            use_auth_token=True if model_args.use_auth_token else None,
        )
    else:
        data_files = {}
        if data_args.train_file is not None:
            data_files["train"] = data_args.train_file
            extension = data_args.train_file.split(".")[-1]
        if data_args.validation_file is not None:
            data_files["validation"] = data_args.validation_file
            extension = data_args.validation_file.split(".")[-1]
        if data_args.test_file is not None:
            data_files["test"] = data_args.test_file
            extension = data_args.test_file.split(".")[-1]
        raw_datasets = load_dataset(
            extension,
            data_files=data_files,
            cache_dir=model_args.cache_dir,
            use_auth_token=True if model_args.use_auth_token else None,
        )
    # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
    # https://huggingface.co/docs/datasets/loading_datasets.html.

    # Load pretrained model and tokenizer
    #
    # Distributed training:
    # The .from_pretrained methods guarantee that only one local process can concurrently
    # download model & vocab.
    config = AutoConfig.from_pretrained(
        model_args.config_name if model_args.config_name else model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        revision=model_args.model_revision,
        use_auth_token=True if model_args.use_auth_token else None,
    )
    tokenizer = AutoTokenizer.from_pretrained(
        model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        use_fast=model_args.use_fast_tokenizer,
        revision=model_args.model_revision,
        use_auth_token=True if model_args.use_auth_token else None,
    )
    model = AutoModelForSeq2SeqLM.from_pretrained(
        model_args.model_name_or_path,
        from_tf=bool(".ckpt" in model_args.model_name_or_path),
        config=config,
        cache_dir=model_args.cache_dir,
        revision=model_args.model_revision,
        use_auth_token=True if model_args.use_auth_token else None,
    )

    # We resize the embeddings only when necessary to avoid index errors. If you are creating a model from scratch
    # on a small vocab and want a smaller embedding size, remove this test.
    embedding_size = model.get_input_embeddings().weight.shape[0]
    if len(tokenizer) > embedding_size:
        model.resize_token_embeddings(len(tokenizer))

    if model.config.decoder_start_token_id is None and isinstance(tokenizer, (MBartTokenizer, MBartTokenizerFast)):
        if isinstance(tokenizer, MBartTokenizer):
            model.config.decoder_start_token_id = tokenizer.lang_code_to_id[data_args.lang]
        else:
            model.config.decoder_start_token_id = tokenizer.convert_tokens_to_ids(data_args.lang)

    if model.config.decoder_start_token_id is None:
        raise ValueError("Make sure that `config.decoder_start_token_id` is correctly defined")

    if (
        hasattr(model.config, "max_position_embeddings")
        and model.config.max_position_embeddings < data_args.max_source_length
    ):
        if model_args.resize_position_embeddings is None:
            logger.warning(
                "Increasing the model's number of position embedding vectors from"
                f" {model.config.max_position_embeddings} to {data_args.max_source_length}."
            )
            model.resize_position_embeddings(data_args.max_source_length)
        elif model_args.resize_position_embeddings:
            model.resize_position_embeddings(data_args.max_source_length)
        else:
            raise ValueError(
                f"`--max_source_length` is set to {data_args.max_source_length}, but the model only has"
                f" {model.config.max_position_embeddings} position encodings. Consider either reducing"
                f" `--max_source_length` to {model.config.max_position_embeddings} or to automatically resize the"
                " model's position encodings by passing `--resize_position_embeddings`."
            )

    prefix = data_args.source_prefix if data_args.source_prefix is not None else ""

    # Preprocessing the datasets.
    # We need to tokenize inputs and targets.
    if training_args.do_train:
        column_names = raw_datasets["train"].column_names
    elif training_args.do_eval:
        column_names = raw_datasets["validation"].column_names
    elif training_args.do_predict:
        column_names = raw_datasets["test"].column_names
    else:
        logger.info("There is nothing to do. Please pass `do_train`, `do_eval` and/or `do_predict`.")
        return

    if isinstance(tokenizer, tuple(MULTILINGUAL_TOKENIZERS)):
        assert (
            data_args.lang is not None
        ), f"{tokenizer.__class__.__name__} is a multilingual tokenizer which requires --lang argument"

        tokenizer.src_lang = data_args.lang
        tokenizer.tgt_lang = data_args.lang

        # For multilingual translation models like mBART-50 and M2M100 we need to force the target language token
        # as the first generated token. We ask the user to explicitly provide this as --forced_bos_token argument.
        forced_bos_token_id = (
            tokenizer.lang_code_to_id[data_args.forced_bos_token] if data_args.forced_bos_token is not None else None
        )
        model.config.forced_bos_token_id = forced_bos_token_id

    # Get the column names for input/target.
    dataset_columns = summarization_name_mapping.get(data_args.dataset_name, None)
    if data_args.text_column is None:
        text_column = dataset_columns[0] if dataset_columns is not None else column_names[0]
    else:
        text_column = data_args.text_column
        if text_column not in column_names:
            raise ValueError(f"--text_column' value '{data_args.text_column}' needs to be one of: {', '.join(column_names)}")
    if data_args.summary_column is None:
        summary_column = dataset_columns[1] if dataset_columns is not None else column_names[1]
    else:
        summary_column = data_args.summary_column
        if summary_column not in column_names:
            raise ValueError(
                f"--summary_column' value '{data_args.summary_column}' needs to be one of: {', '.join(column_names)}"
            )

    # Temporarily set max_target_length for training.
    max_target_length = data_args.max_target_length
    padding = "max_length" if data_args.pad_to_max_length else False

    if training_args.label_smoothing_factor > 0 and not hasattr(model, "prepare_decoder_input_ids_from_labels"):
        logger.warning(
            "label_smoothing is enabled but the `prepare_decoder_input_ids_from_labels` method is not defined for"
            f"`{model.__class__.__name__}`. This will lead to loss being calculated twice and will take up more memory"
        )

    def preprocess_function(examples):
        # remove pairs where at least one record is None

        inputs, targets = [], []
        for i in range(len(examples[text_column])):
            if examples[text_column][i] and examples[summary_column][i]:
                inputs.append(examples[text_column][i])
                targets.append(examples[summary_column][i])

        inputs = [prefix + inp for inp in inputs]
        model_inputs = tokenizer(inputs, max_length=data_args.max_source_length, padding=padding, truncation=True)

        # Tokenize targets with the `text_target` keyword argument
        labels = tokenizer(text_target=targets, max_length=max_target_length, padding=padding, truncation=True)

        # If we are padding here, replace all tokenizer.pad_token_id in the labels by -100 when we want to ignore
        # padding in the loss.
        if padding == "max_length" and data_args.ignore_pad_token_for_loss:
            labels["input_ids"] = [
                [(l if l != tokenizer.pad_token_id else -100) for l in label] for label in labels["input_ids"]
            ]

        model_inputs["labels"] = labels["input_ids"]
        return model_inputs

    if training_args.do_train:
        if "train" not in raw_datasets:
            raise ValueError("--do_train requires a train dataset")
        train_dataset = raw_datasets["train"]
        if data_args.max_train_samples is not None:
            max_train_samples = min(len(train_dataset), data_args.max_train_samples)
            train_dataset = train_dataset.select(range(max_train_samples))
        with training_args.main_process_first(desc="train dataset map pre-processing"):
            train_dataset = train_dataset.map(
                preprocess_function,
                batched=True,
                num_proc=data_args.preprocessing_num_workers,
                remove_columns=column_names,
                load_from_cache_file=not data_args.overwrite_cache,
                desc="Running tokenizer on train dataset",
            )

    if training_args.do_eval:
        max_target_length = data_args.val_max_target_length
        if "validation" not in raw_datasets:
            raise ValueError("--do_eval requires a validation dataset")
        eval_dataset = raw_datasets["validation"]
        if data_args.max_eval_samples is not None:
            max_eval_samples = min(len(eval_dataset), data_args.max_eval_samples)
            eval_dataset = eval_dataset.select(range(max_eval_samples))
        with training_args.main_process_first(desc="validation dataset map pre-processing"):
            eval_dataset = eval_dataset.map(
                preprocess_function,
                batched=True,
                num_proc=data_args.preprocessing_num_workers,
                remove_columns=column_names,
                load_from_cache_file=not data_args.overwrite_cache,
                desc="Running tokenizer on validation dataset",
            )

    if training_args.do_predict:
        max_target_length = data_args.val_max_target_length
        if "test" not in raw_datasets:
            raise ValueError("--do_predict requires a test dataset")
        predict_dataset = raw_datasets["test"]
        if data_args.max_predict_samples is not None:
            max_predict_samples = min(len(predict_dataset), data_args.max_predict_samples)
            predict_dataset = predict_dataset.select(range(max_predict_samples))
        with training_args.main_process_first(desc="prediction dataset map pre-processing"):
            predict_dataset = predict_dataset.map(
                preprocess_function,
                batched=True,
                num_proc=data_args.preprocessing_num_workers,
                remove_columns=column_names,
                load_from_cache_file=not data_args.overwrite_cache,
                desc="Running tokenizer on prediction dataset",
            )

    # Data collator
    label_pad_token_id = -100 if data_args.ignore_pad_token_for_loss else tokenizer.pad_token_id
    data_collator = DataCollatorForSeq2Seq(
        tokenizer,
        model=model,
        label_pad_token_id=label_pad_token_id,
        pad_to_multiple_of=8 if training_args.fp16 else None,
    )

    # Metric
    metric = evaluate.load("rouge")

    def compute_metrics(eval_preds):
        preds, labels = eval_preds
        if isinstance(preds, tuple):
            preds = preds[0]
        decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
        if data_args.ignore_pad_token_for_loss:
            # Replace -100 in the labels as we can't decode them.
            labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
        decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)

        # Some simple post-processing
        decoded_preds, decoded_labels = postprocess_text_for_metric("rouge", decoded_preds, decoded_labels)

        result = metric.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True)
        result = {k: round(v * 100, 4) for k, v in result.items()}
        prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in preds]
        result["gen_len"] = np.mean(prediction_lens)
        return result

    # Initialize our Trainer
    trainer = BaselineSeq2SeqTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset if training_args.do_train else None,
        eval_dataset=eval_dataset if training_args.do_eval else None,
        tokenizer=tokenizer,
        data_collator=data_collator,
        compute_metrics=compute_metrics if training_args.predict_with_generate else None,
    )

    # Training
    if training_args.do_train:
        checkpoint = None
        if training_args.resume_from_checkpoint is not None:
            checkpoint = training_args.resume_from_checkpoint
        elif last_checkpoint is not None:
            checkpoint = last_checkpoint
        train_result = trainer.train(resume_from_checkpoint=checkpoint)
        trainer.save_model()  # Saves the tokenizer too for easy upload

        metrics = train_result.metrics
        max_train_samples = data_args.max_train_samples if data_args.max_train_samples is not None else len(train_dataset)
        metrics["train_samples"] = min(max_train_samples, len(train_dataset))

        trainer.log_metrics("train", metrics)
        trainer.save_metrics("train", metrics)
        trainer.save_state()


    # We will load the best model here to avoid an issue when do_train is not set.
    if training_args.load_states_in_eval_from_model_path and not training_args.do_train:
        trainer.state = TrainerState.load_from_json(os.path.join(model_args.model_name_or_path, "trainer_state.json"))
        if training_args.load_best_model_at_end and trainer.state.best_model_checkpoint is not None:
            checkpoint_path = trainer.state.best_model_checkpoint
        else:
            checkpoint_path = model_args.model_name_or_path
        trainer._load_from_checkpoint(checkpoint_path)
        trainer._load_rng_state(checkpoint_path)

    # Evaluation
    results = {}
    max_length = (
        training_args.generation_max_length
        if training_args.generation_max_length is not None
        else data_args.val_max_target_length
    )
    num_beams = data_args.num_beams if data_args.num_beams is not None else training_args.generation_num_beams
    if training_args.do_eval:
        logger.info("*** Evaluate ***")
        metrics = trainer.evaluate(max_length=max_length, num_beams=num_beams, metric_key_prefix="eval", top_p=training_args.top_p, temperature=training_args.temperature)
        max_eval_samples = data_args.max_eval_samples if data_args.max_eval_samples is not None else len(eval_dataset)
        metrics["eval_samples"] = min(max_eval_samples, len(eval_dataset))

        trainer.log_metrics("eval", metrics)
        trainer.save_metrics("eval", metrics)

    if training_args.do_predict:
        logger.info("*** Predict ***")

        predict_results = trainer.predict(
            predict_dataset, metric_key_prefix="predict", max_length=max_length, num_beams=num_beams, top_p=training_args.top_p, temperature=training_args.temperature
        )
        metrics = predict_results.metrics
        max_predict_samples = (
            data_args.max_predict_samples if data_args.max_predict_samples is not None else len(predict_dataset)
        )
        metrics["predict_samples"] = min(max_predict_samples, len(predict_dataset))

        trainer.log_metrics("predict", metrics)
        trainer.save_metrics("predict", metrics)

        if trainer.is_world_process_zero():
            if training_args.predict_with_generate:
                predictions = tokenizer.batch_decode(predict_results.predictions, skip_special_tokens=True)
                labels = np.where(predict_results.label_ids != -100, predict_results.label_ids, tokenizer.pad_token_id)
                decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
                prediction_results = {"predictions": predictions, "gold_text":decoded_labels}
                trainer.save_metrics(GENERATION_RESULTS + "_predict", prediction_results)

    return results


if __name__ == "__main__":
    main()

baselines/run_summarization.py
allenai/tess-diffusion
from transformers import Seq2SeqTrainer
import shutil
import os
from transformers.utils import logging

logger = logging.get_logger(__name__)

class BaselineSeq2SeqTrainer(Seq2SeqTrainer):
    def annotated_split(self, split):
        return f"{split}_top_p_{self.args.top_p}_temperature_{self.args.temperature}_seed_{self.args.seed}"

    def save_metrics(self, split, metrics, combined=True):
        super().save_metrics(self.annotated_split(split), metrics, combined)

    def log_metrics(self, split, metrics):
        super().log_metrics(self.annotated_split(split), metrics)

    def _rotate_checkpoints(self, use_mtime=False, output_dir=None) -> None:
        if self.args.save_total_limit is None or self.args.save_total_limit <= 0:
            return

        # Check if we should delete older checkpoint(s)
        checkpoints_sorted = self._sorted_checkpoints(use_mtime=use_mtime, output_dir=output_dir)
        if len(checkpoints_sorted) <= self.args.save_total_limit:
            return

        # If save_total_limit=1 with load_best_model_at_end=True, we could end up deleting the last checkpoint, which
        # we don't do to allow resuming.
        save_total_limit = self.args.save_total_limit
        if (
            self.state.best_model_checkpoint is not None
            and self.args.save_total_limit == 1
            and checkpoints_sorted[-1] != self.state.best_model_checkpoint
        ):
            save_total_limit = 2

        number_of_checkpoints_to_delete = max(0, len(checkpoints_sorted) - save_total_limit)
        checkpoints_to_be_deleted = checkpoints_sorted[:number_of_checkpoints_to_delete]
        for checkpoint in checkpoints_to_be_deleted:
            logger.info(f"Deleting older checkpoint [{checkpoint}] due to args.save_total_limit")
            shutil.rmtree(checkpoint)

baselines/trainer_seq2seq.py
allenai/tess-diffusion
# computes the metrics.
import sys
import os
from transformers import HfArgumentParser
from sdlm.arguments import DataTrainingArguments, ModelArguments, TrainingArguments, DiffusionArguments
from run_mlm import get_compute_metrics
from sdlm.trainer import GENERATION_RESULTS
import json
from transformers.trainer_utils import denumpify_detensorize


def save_metrics(metrics, split, training_args):
    path = os.path.join(training_args.output_dir, f"{split}_results.json")
    with open(path, "w") as f:
        json.dump(metrics, f, indent=4, sort_keys=True)


def main():
    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments, DiffusionArguments))
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        # If we pass only one argument to the script and it's the path to a json file,
        # let's parse it to get our arguments.
        # NOTE: keep the diffusion_args to avoid the error when we call the scripts with all the arguments
        # used in running `run_mlm.py`.
        model_args, data_args, training_args, diffusion_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        model_args, data_args, training_args, diffusion_args = parser.parse_args_into_dataclasses()

    compute_metrics = get_compute_metrics(data_args, training_args, model_args)
    # Load generations.
    results = json.load(open(os.path.join(training_args.output_dir, GENERATION_RESULTS + "_eval_results.json")))
    metrics = compute_metrics(results)
    metrics = denumpify_detensorize(metrics)
    save_metrics(metrics, "eval", training_args)


if __name__ == "__main__":
    main()

compute_mlm_metrics.py
allenai/tess-diffusion
"""Evaluates the GPT-2 model results. This script runs on 1 GPU."""
from transformers import AutoTokenizer, AutoModelForCausalLM
import logging
import sys
import os
from sdlm.arguments import DataTrainingArguments, TrainingArguments, ModelArguments
import transformers
from transformers import HfArgumentParser, set_seed
from torch.utils.data import DataLoader
import datasets
from datasets import load_from_disk
from sdlm.data.data_utils import load_data, tokenize_data_new
from sdlm.data.data_collator import SpanInfillingDataCollator
from sdlm.inference.inference_utils import evaluate_generation
import pdb
import json
import torch

logger = logging.getLogger(__name__)


def prepare_inputs(inputs, device):
    return {k: v.to(device) for k, v in inputs.items()}


def main():
    parser = HfArgumentParser((DataTrainingArguments, TrainingArguments, ModelArguments))
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        data_args, training_args, model_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        data_args, training_args, model_args = parser.parse_args_into_dataclasses()

    os.makedirs(training_args.output_dir, exist_ok=True)

    # Setup logging
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        handlers=[logging.StreamHandler(sys.stdout)],
    )

    log_level = training_args.get_process_log_level()
    logger.setLevel(log_level)
    datasets.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.enable_default_handler()
    transformers.utils.logging.enable_explicit_format()

    # Log on each process the small summary:
    logger.warning(
        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
        + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
    )
    # Set the verbosity to info of the Transformers logger (on main process only):
    logger.info(f"Training/evaluation parameters {training_args}")

    # Set seed before initializing model.
    set_seed(training_args.seed)

    tokenizer_kwargs = {
        "cache_dir": model_args.cache_dir,
        "use_fast": model_args.use_fast_tokenizer,
        "revision": model_args.model_revision,
        "use_auth_token": True if model_args.use_auth_token else None,
    }
    tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path, **tokenizer_kwargs)
    roberta_tokenizer = AutoTokenizer.from_pretrained("roberta-large", **tokenizer_kwargs)
    # Running the script requires the tokenizer to have the pad token and since gpt2 tokenizer
    # does not have it, we add the pad_token here. Also, during the generation, they use the
    # eos_token_id as the pad_token_id.
    # We need to modify the bos/eos tokens to help with process function.
    tokenizer.bos_token = roberta_tokenizer.bos_token
    tokenizer.eos_token = roberta_tokenizer.eos_token
    tokenizer.pad_token = tokenizer.eos_token
    # Huggingface requires this to be set.
    tokenizer.padding_side = "left"

    model = AutoModelForCausalLM.from_pretrained(
        model_args.model_name_or_path,
        from_tf=bool(".ckpt" in model_args.model_name_or_path),
        cache_dir=model_args.cache_dir,
        revision=model_args.model_revision,
        use_auth_token=True if model_args.use_auth_token else None,
    ).to(training_args.device)
    model.eval()

    if data_args.tokenized_data_path:
        tokenized_datasets = load_from_disk(data_args.tokenized_data_path)
    else:
        raw_datasets = load_data(data_args, model_args)
        tokenized_datasets = tokenize_data_new(data_args, roberta_tokenizer, raw_datasets, training_args)

    eval_dataset = tokenized_datasets["validation"]
    if data_args.max_eval_samples is not None:
        max_eval_samples = min(len(eval_dataset), data_args.max_eval_samples)
        eval_dataset = eval_dataset.select(range(max_eval_samples))

    pad_to_multiple_of_8 = data_args.line_by_line and training_args.fp16 and not data_args.pad_to_max_length
    data_collator = SpanInfillingDataCollator(
        mode="eval",
        data_args=data_args,
        tokenizer=roberta_tokenizer,
        max_length=data_args.max_seq_length,
        seed=training_args.seed,
        pad_to_multiple_of=8 if pad_to_multiple_of_8 else None,
        eval_context_size=data_args.eval_context_size,
    )

    # Creates the data_loader.
    eval_dataloader = DataLoader(
        eval_dataset,
        collate_fn=data_collator,
        batch_size=training_args.per_device_eval_batch_size,
    )
    all_outputs = []
    all_inputs = []
    all_prefixes = []
    all_masks = []
    with torch.no_grad():
        for i, batch in enumerate(eval_dataloader):
            # De-tokenize with the roberta tokenizer.
            inputs, span_mask = batch["input_ids"], batch["span_mask"]
            if data_args.truncation_length > 0:
                inputs = inputs[:, : -data_args.truncation_length]
                span_mask = span_mask[:, : -data_args.truncation_length]
                max_seq_length = data_args.max_seq_length - data_args.truncation_length
                assert data_args.eval_context_size < max_seq_length
            all_masks.extend(span_mask)
            all_inputs.extend(inputs)
            prefixes_tokens = [input[~mask] for input, mask in zip(inputs, span_mask)]
            prefixes = roberta_tokenizer.batch_decode(prefixes_tokens, skip_special_tokens=True)
            all_prefixes.extend(prefixes)
            # Note that output also include the prefix and we need to remove it here.
            gold_texts = [input[len(prefix) :] for input, prefix in zip(inputs, prefixes_tokens)]
            gold_texts = [tokens.cpu().numpy().tolist() for tokens in gold_texts]
            all_outputs.extend(gold_texts)

    results = {}
    gold_texts = [
        roberta_tokenizer.decode(input[mask], skip_special_tokens=True) for input, mask in zip(all_inputs, all_masks)
    ]
    results = {
        "generated_texts_masked": gold_texts,
        "gold_texts_masked": gold_texts,
        "generated_texts_masked_tokens": all_outputs,
        "prefixes": all_prefixes,
    }
    # Saves the generated results.
    with open(f"{training_args.output_dir}/generated_results.json", "w") as f:
        json.dump(results, f)

    metrics = evaluate_generation(
        results,
        data_args,
        model,
        tokenizer,
        is_conditional_generation=True,
        prefix_lm_eval=True,
        skip_special_tokens=True,
        eval_for_all_metrics=True,
    )
    with open(f"{training_args.output_dir}/metrics.json", "w") as f:
        json.dump(metrics, f)


if __name__ == "__main__":
    main()

gold_text_eval.py
allenai/tess-diffusion
"""Evaluates the GPT-2 model results. This script runs on 1 GPU."""
from transformers import AutoTokenizer, AutoModelForCausalLM
import logging
import sys
import os
from sdlm.arguments import DataTrainingArguments, TrainingArguments, ModelArguments, DiffusionArguments
import transformers
from transformers import HfArgumentParser, set_seed
from torch.utils.data import DataLoader
import datasets
from datasets import load_from_disk
from sdlm.data.data_utils import load_data, tokenize_data_new
from sdlm.data.data_collator import SpanInfillingDataCollator
from sdlm.inference.inference_utils import evaluate_generation
import pdb
import numpy as np
import json
import torch
import pdb

logger = logging.getLogger(__name__)


def prepare_inputs(inputs, device):
    return {k: v.to(device) for k, v in inputs.items()}


def main():
    parser = HfArgumentParser((DataTrainingArguments, TrainingArguments, ModelArguments, DiffusionArguments))
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        data_args, training_args, model_args, diffusion_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        data_args, training_args, model_args, diffusion_args = parser.parse_args_into_dataclasses()

    os.makedirs(training_args.output_dir, exist_ok=True)

    # Setup logging
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        handlers=[logging.StreamHandler(sys.stdout)],
    )

    log_level = training_args.get_process_log_level()
    logger.setLevel(log_level)
    datasets.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.enable_default_handler()
    transformers.utils.logging.enable_explicit_format()

    # Log on each process the small summary:
    logger.warning(
        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
        + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
    )
    # Set the verbosity to info of the Transformers logger (on main process only):
    logger.info(f"Training/evaluation parameters {training_args}")

    # Set seed before initializing model.
    set_seed(training_args.seed)

    tokenizer_kwargs = {
        "cache_dir": model_args.cache_dir,
        "use_fast": model_args.use_fast_tokenizer,
        "revision": model_args.model_revision,
        "use_auth_token": True if model_args.use_auth_token else None,
    }
    tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path, **tokenizer_kwargs)
    roberta_tokenizer = AutoTokenizer.from_pretrained("roberta-large", **tokenizer_kwargs)
    # Running the script requires the tokenizer to have the pad token and since gpt2 tokenizer
    # does not have it, we add the pad_token here. Also, during the generation, they use the
    # eos_token_id as the pad_token_id.
    # We need to modify the bos/eos tokens to help with process function.
    tokenizer.bos_token = roberta_tokenizer.bos_token
    tokenizer.eos_token = roberta_tokenizer.eos_token
    tokenizer.pad_token = tokenizer.eos_token
    # Huggingface requires this to be set.
    tokenizer.padding_side = "left"

    model = AutoModelForCausalLM.from_pretrained(
        model_args.model_name_or_path,
        from_tf=bool(".ckpt" in model_args.model_name_or_path),
        cache_dir=model_args.cache_dir,
        revision=model_args.model_revision,
        use_auth_token=True if model_args.use_auth_token else None,
    ).to(training_args.device)
    model.eval()

    if data_args.tokenized_data_path:
        tokenized_datasets = load_from_disk(data_args.tokenized_data_path)
    else:
        raw_datasets = load_data(data_args, model_args)
        tokenized_datasets = tokenize_data_new(data_args, roberta_tokenizer, raw_datasets, training_args)

    eval_dataset = tokenized_datasets["validation"]
    if data_args.max_eval_samples is not None:
        max_eval_samples = min(len(eval_dataset), data_args.max_eval_samples)
        eval_dataset = eval_dataset.select(range(max_eval_samples))

    pad_to_multiple_of_8 = data_args.line_by_line and training_args.fp16 and not data_args.pad_to_max_length
    data_collator = SpanInfillingDataCollator(
        mode="eval",
        data_args=data_args,
        tokenizer=roberta_tokenizer,
        max_length=data_args.max_seq_length,
        seed=training_args.seed,
        pad_to_multiple_of=8 if pad_to_multiple_of_8 else None,
        eval_context_size=data_args.eval_context_size,
    )

    # Creates the data_loader.
    eval_dataloader = DataLoader(
        eval_dataset,
        collate_fn=data_collator,
        batch_size=training_args.per_device_eval_batch_size,
    )
    all_outputs = []
    all_inputs = []
    all_prefixes = []
    all_masks = []
    all_prefix_tokens = []
    max_seq_length = data_args.max_seq_length
    with torch.no_grad():
        for i, batch in enumerate(eval_dataloader):
            # De-tokenize with the roberta tokenizer.
            inputs, span_mask = batch["input_ids"], batch["span_mask"]
            # Truncate the length if needed.
            if data_args.truncation_length > 0:
                inputs = inputs[:, : -data_args.truncation_length]
                span_mask = span_mask[:, : -data_args.truncation_length]
                max_seq_length = data_args.max_seq_length - data_args.truncation_length
                assert data_args.eval_context_size < max_seq_length

            all_masks.extend(span_mask)
            all_inputs.extend(inputs)
            prefixes = [input[~mask] for input, mask in zip(inputs, span_mask)]
            prefixes = [prefix.cpu().numpy().tolist() for prefix in prefixes]
            all_prefix_tokens.append(prefixes)
            prefixes = roberta_tokenizer.batch_decode(prefixes, skip_special_tokens=True)
            all_prefixes.extend(prefixes)
            prefixes_inputs = tokenizer(prefixes, return_tensors="pt", padding=True)
            prefixes_inputs = prepare_inputs(prefixes_inputs, training_args.device)
            outputs = model.generate(
                input_ids=prefixes_inputs["input_ids"],
                attention_mask=prefixes_inputs["attention_mask"],
                pad_token_id=tokenizer.eos_token_id,
                max_length=max_seq_length,
                min_length=max_seq_length,
                do_sample=True,
                top_p=diffusion_args.top_p,
            )
            # Note that output also include the prefix and we need to remove it here.
            generated_outputs = [output[len(prefix) :] for output, prefix in zip(outputs, prefixes_inputs["input_ids"])]
            generated_outputs = [tokens.cpu().numpy().tolist() for tokens in generated_outputs]
            all_outputs.extend(generated_outputs)
    results = {}
    generated_texts = [tokenizer.decode(output, skip_special_tokens=True) for output in all_outputs]
    gold_texts = [
        roberta_tokenizer.decode(input[mask], skip_special_tokens=True) for input, mask in zip(all_inputs, all_masks)
    ]
    total_texts_marked = [
        prefix + " ***" + generated_text + "***" for prefix, generated_text in zip(all_prefixes, generated_texts)
    ]
    results = {
        "generated_texts_masked": generated_texts,
        "gold_texts_masked": gold_texts,
        "generated_texts_masked_tokens": all_outputs,
        "prefixes": all_prefixes,
        "prefix_tokens": all_prefix_tokens,
    }
    # Saves the generated results.
    with open(f"{training_args.output_dir}/generated_results.json", "w") as f:
        json.dump(results, f)

    metrics = evaluate_generation(
        results,
        data_args,
        model,
        tokenizer,
        is_conditional_generation=True,
        prefix_lm_eval=True,
        skip_special_tokens=True,
        eval_for_all_metrics=True,
    )
    logger.info(metrics)
    for text in total_texts_marked:
        logger.info(text)

    with open(f"{training_args.output_dir}/metrics.json", "w") as f:
        json.dump(metrics, f)


if __name__ == "__main__":
    main()

gpt2_eval.py
allenai/tess-diffusion
""" Finetuning the library models for sequence classification on GLUE."""

import logging
import os
import random
import sys
from dataclasses import dataclass
from transformers.trainer_callback import TrainerState
import datasets
import numpy as np
from datasets import load_dataset
import pdb
import transformers
from transformers import AutoTokenizer, HfArgumentParser, set_seed
from transformers.trainer_utils import get_last_checkpoint
from transformers.utils import check_min_version, send_example_telemetry
from transformers.utils.versions import require_version
from sdlm.arguments import ModelArguments, DiffusionArguments, TrainingArguments
from sdlm.arguments import DataTrainingArguments as BaseDataTrainingArguments
from sdlm.models import RobertaDiffusionConfig, RobertaForDiffusionLM
from sdlm.schedulers import SimplexDDPMScheduler
from sdlm.data.data_utils import split_glue
from sdlm.utils import round_stsb_target, lmap
from sdlm.data.data_collator import DataCollatorForSeq2Seq
from sdlm.trainer import DiffusionTrainer
from sdlm.inference.inference_utils import process_text
from sdlm.metrics.metrics import get_glue_metrics
from sdlm.data.postprocessors import get_post_processor
from transformers.utils import WEIGHTS_NAME

# This is computed with scripts/compute_max_tokens_of_labels.py
MAX_LABEL_LENGTH = 5
check_min_version("4.25.0")

require_version("datasets>=1.8.0")

task_to_keys = {
    "cola": ("sentence", None),
    "mnli": ("premise", "hypothesis"),
    "mrpc": ("sentence1", "sentence2"),
    "qnli": ("question", "sentence"),
    "qqp": ("question1", "question2"),
    "rte": ("sentence1", "sentence2"),
    "sst2": ("sentence", None),
    "stsb": ("sentence1", "sentence2"),
    "wnli": ("sentence1", "sentence2"),
}

task_to_metric = {
    "cola": "matthews_correlation",
    "mnli": "accuracy",
    "mrpc": "combined_score",
    "qnli": "accuracy",
    "qqp": "combined_score",
    "rte": "accuracy",
    "sst2": "accuracy",
    "stsb": "combined_score",
    "wnli": "accuracy",
}

logger = logging.getLogger(__name__)


@dataclass
class DataTrainingArguments(BaseDataTrainingArguments):
    """
    Arguments pertaining to what data we are going to input our model for training and eval.
    """

    def __post_init__(self):
        assert self.dataset_name is not None
        self.dataset_name = self.dataset_name.lower()
        if self.dataset_name not in task_to_keys.keys():
            raise ValueError("Unknown task, you should pick one in " + ",".join(task_to_keys.keys()))


def main():
    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments, DiffusionArguments))
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        # If we pass only one argument to the script and it's the path to a json file,
        # let's parse it to get our arguments.
        model_args, data_args, training_args, diffusion_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        model_args, data_args, training_args, diffusion_args = parser.parse_args_into_dataclasses()

    if training_args.checkpoint_best_model:
        # TODO: ask which one they report and use the one needed here.
        # TODO: test both simplex and logits.
        training_args.metric_for_best_model = "pred_texts_from_simplex_masked_" + task_to_metric[data_args.dataset_name]

    # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
    # information sent is the one passed as arguments along with your Python/PyTorch versions.
    send_example_telemetry("run_glue", model_args, data_args)

    # Setup logging
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        handlers=[logging.StreamHandler(sys.stdout)],
    )

    log_level = training_args.get_process_log_level()
    logger.setLevel(log_level)
    datasets.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.enable_default_handler()
    transformers.utils.logging.enable_explicit_format()

    # Log on each process the small summary:
    logger.warning(
        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
        + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
    )
    logger.info(f"Training/evaluation parameters {training_args}")

    # Detecting last checkpoint.
    last_checkpoint = None
    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
            raise ValueError(
                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
                "Use --overwrite_output_dir to overcome."
            )
        elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
            logger.info(
                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
            )

    # Set seed before initializing model.
    set_seed(training_args.seed)

    # Downloading and loading a dataset from the hub.
    raw_datasets = load_dataset(
        "glue",
        data_args.dataset_name,
        cache_dir=model_args.cache_dir,
        use_auth_token=True if model_args.use_auth_token else None,
    )

    # Split dataset, since test sets of GLUE do not have the labels.
    if data_args.split_glue:
        raw_datasets = split_glue(raw_datasets, data_args.dataset_name, data_args.glue_split_seed) 
    elif data_args.dataset_name == "mnli":
        raw_datasets["validation"] =  raw_datasets["validation_matched"] # mismatched is for reverse, and for normal is matched.
        raw_datasets["test"] = raw_datasets["test_matched"]
      
    # Labels
    is_regression = data_args.dataset_name == "stsb"
    config = RobertaDiffusionConfig.from_pretrained(
        model_args.model_name_or_path,
        self_condition=diffusion_args.self_condition,
        self_condition_zeros_after_softmax=diffusion_args.self_condition_zeros_after_softmax,
        deepmind_conditional=diffusion_args.deepmind_conditional,
        classifier_free_simplex_inputs=diffusion_args.classifier_free_simplex_inputs,
        classifier_free_uncond_input=diffusion_args.classifier_free_uncond_input,
        self_condition_mlp_projection=diffusion_args.self_condition_mlp_projection,
        self_condition_mix_before_weights=diffusion_args.self_condition_mix_before_weights,
        empty_token_be_mask=diffusion_args.empty_token_be_mask,
        cache_dir=model_args.cache_dir,
        revision=model_args.model_revision,
        use_auth_token=True if model_args.use_auth_token else None,
    )

    tokenizer = AutoTokenizer.from_pretrained(
        model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        use_fast=model_args.use_fast_tokenizer,
        revision=model_args.model_revision,
        use_auth_token=True if model_args.use_auth_token else None,
    )
    if model_args.model_name_or_path:
        model = RobertaForDiffusionLM.from_pretrained(
            model_args.model_name_or_path,
            from_tf=bool(".ckpt" in model_args.model_name_or_path),
            config=config,
            cache_dir=model_args.cache_dir,
            revision=model_args.model_revision,
            use_auth_token=True if model_args.use_auth_token else None,
        )
    else:
        logger.info("Training new model from scratch")
        model = RobertaForDiffusionLM.from_config(config)

    # Preprocessing the raw_datasets
    sentence1_key, sentence2_key = task_to_keys[data_args.dataset_name]

    if data_args.max_seq_length > tokenizer.model_max_length:
        logger.warning(
            f"The max_seq_length passed ({data_args.max_seq_length}) is larger than the maximum length for the"
            f"model ({tokenizer.model_max_length}). Using max_seq_length={tokenizer.model_max_length}."
        )
    max_seq_length = min(data_args.max_seq_length, tokenizer.model_max_length)

    def preprocess_function(examples):
        # TODO: here max_length should be max_length minus length of labels.
        # TODO: this is for now, but maybe compute one max_length as a whole.
        # Tokenize the labels.
        targets = [str(round_stsb_target(label)) if is_regression else str(label) for label in examples["label"]]
        labels = tokenizer(text_target=targets, max_length=max_seq_length, padding=False, truncation=True)
        max_label_length = MAX_LABEL_LENGTH  # max([len(label) for label in labels["input_ids"]])

        # Tokenize the texts.
        if data_args.add_t5_tags:
            sentence1_with_tag = [sentence1_key + ": " + sentence_1 for sentence_1 in examples[sentence1_key]]
            if sentence2_key is not None:
                sentence2_with_tag = [sentence2_key + ": " + sentence_2 for sentence_2 in examples[sentence2_key]]
            args = (sentence1_with_tag,) if sentence2_key is None else (sentence1_with_tag, sentence2_with_tag)
        else:
            args = (
                (examples[sentence1_key],) if sentence2_key is None else (examples[sentence1_key], examples[sentence2_key])
            )
        result = tokenizer(*args, padding=False, max_length=max_seq_length - max_label_length, truncation=True)
        result["labels"] = labels["input_ids"]
        return result

    with training_args.main_process_first(desc="dataset map pre-processing"):
        raw_datasets = raw_datasets.map(
            preprocess_function,
            batched=True,
            load_from_cache_file=not data_args.overwrite_cache,
            desc="Running tokenizer on dataset",
        )
    if training_args.do_train:
        if "train" not in raw_datasets:
            raise ValueError("--do_train requires a train dataset")
        train_dataset = raw_datasets["train"]
        if data_args.max_train_samples is not None:
            max_train_samples = min(len(train_dataset), data_args.max_train_samples)
            train_dataset = train_dataset.select(range(max_train_samples))

    if training_args.do_eval:
        if "validation" not in raw_datasets:
            raise ValueError("--do_eval requires a validation dataset")
        eval_dataset = raw_datasets["validation"]
        if data_args.max_eval_samples is not None:
            max_eval_samples = min(len(eval_dataset), data_args.max_eval_samples)
            eval_dataset = eval_dataset.select(range(max_eval_samples))

    def preprocess_logits_for_metrics(logits):
        return logits.argmax(dim=-1)

    if training_args.do_predict or data_args.dataset_name is not None or data_args.test_file is not None:
        if "test" not in raw_datasets:
            raise ValueError("--do_predict requires a test dataset")
        predict_datasets = [raw_datasets["test"]] if data_args.dataset_name != "mnli" else [raw_datasets["test_matched"]]
        if data_args.dataset_name == "mnli":
            predict_datasets.append(raw_datasets["test_mismatched"])

        if data_args.max_predict_samples is not None:
            for i in range(len(predict_datasets)):
                max_predict_samples = min(len(predict_datasets[i]), data_args.max_predict_samples)
                predict_datasets[i] = predict_datasets[i].select(range(max_predict_samples))

    # Log a few random samples from the training set:
    if training_args.do_train:
        for index in random.sample(range(len(train_dataset)), 3):
            logger.info(f"Sample {index} of the training set: {train_dataset[index]}.")

    # Get the metric function
    task_metrics = get_glue_metrics(data_args.dataset_name)

    def postprocess_text(texts):
        # TODO: maybe we need it for others as well.
        return lmap(str.strip, texts)

    # TODO: we maybe need to pad till the sentences, and then predict the tokens we need for the few ones we need.
    def compute_metrics(results):
        post_processor = get_post_processor(data_args.dataset_name)

        # TODO: we need to change the metrics here.
        keys = ["pred_texts_from_simplex_masked", "pred_texts_from_logits_masked"]
        decoded_labels = postprocess_text(process_text(results["gold_texts_masked"]))
        if post_processor is not None:
            decoded_labels = [post_processor(x) for x in decoded_labels]

        metrics = {}
        for key in keys:
            decoded_preds = postprocess_text(process_text(results[key]))
            if post_processor is not None:
                decoded_preds = [post_processor(x) for x in decoded_preds]
            key_metrics = {}
            for metric in task_metrics:
                key_metrics.update(metric(predictions=decoded_preds, targets=decoded_labels))
            if len(key_metrics) > 1:
                key_metrics["combined_score"] = np.mean(list(key_metrics.values())).item()
            key_metrics = {f"{key}_{k}": v for k, v in key_metrics.items()}
            metrics.update(key_metrics)

        return metrics

    # Data collator will default to DataCollatorWithPadding when the tokenizer is passed to Trainer, so we change it if
    # we already did the padding.
    # Data collator. To be consistent with the run_mlm.py we need to add `mode`.
    data_collator = lambda mode: DataCollatorForSeq2Seq(
        tokenizer,
        # Note that if you do not use `pad_to_max_length`, this becomes very slow on multi-gpus.
        padding="max_length" if data_args.pad_to_max_length else True,
        max_length=data_args.max_seq_length,
        pad_to_multiple_of=8 if training_args.fp16 else None,
    )
    # TODO: here we need to make sure we mask to the maximum number of tokens in the labels to not signal the model for the labels.

    noise_scheduler = SimplexDDPMScheduler(
        num_train_timesteps=diffusion_args.num_diffusion_steps,
        beta_schedule=diffusion_args.beta_schedule,
        simplex_value=diffusion_args.simplex_value,
        clip_sample=diffusion_args.clip_sample,
        device=training_args.device,
    )
    inference_noise_scheduler = SimplexDDPMScheduler(
        num_train_timesteps=diffusion_args.num_inference_diffusion_steps,
        beta_schedule=diffusion_args.beta_schedule,
        simplex_value=diffusion_args.simplex_value,
        clip_sample=diffusion_args.clip_sample,
        device=training_args.device,
    )

    # Initialize our Trainer
    trainer = DiffusionTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset if training_args.do_train else None,
        eval_dataset=eval_dataset if training_args.do_eval else None,
        tokenizer=tokenizer,
        data_collator=data_collator,
        compute_metrics=compute_metrics if (training_args.do_eval or training_args.do_predict) else None,
        preprocess_logits_for_metrics=preprocess_logits_for_metrics if (training_args.do_eval or training_args.do_predict) else None,
        noise_scheduler=noise_scheduler,
        diffusion_args=diffusion_args,
        data_args=data_args,
        inference_noise_scheduler=inference_noise_scheduler,
    )

    # Training
    if training_args.do_train:
        checkpoint = None
        if training_args.resume_from_checkpoint is not None:
            checkpoint = training_args.resume_from_checkpoint
        elif last_checkpoint is not None:
            checkpoint = last_checkpoint
        train_result = trainer.train(resume_from_checkpoint=checkpoint)
        metrics = train_result.metrics
        max_train_samples = data_args.max_train_samples if data_args.max_train_samples is not None else len(train_dataset)
        metrics["train_samples"] = min(max_train_samples, len(train_dataset))

        trainer.save_model()  # Saves the tokenizer too for easy upload

        trainer.log_metrics("train", metrics)
        trainer.save_metrics("train", metrics)
        trainer.save_state()

    # We will load the best model here to avoid an issue when do_train is not set.
    if training_args.load_states_in_eval_from_model_path and not training_args.do_train:
        trainer.state = TrainerState.load_from_json(os.path.join(model_args.model_name_or_path, "trainer_state.json"))
        if training_args.load_best_model_at_end and trainer.state.best_model_checkpoint is not None:
            checkpoint_path = trainer.state.best_model_checkpoint
        else:
            checkpoint_path = model_args.model_name_or_path
        trainer._load_from_checkpoint(checkpoint_path)
        trainer._load_rng_state(checkpoint_path)

    # Evaluation
    if training_args.do_eval:
        logger.info("*** Evaluate ***")
        metrics = trainer.evaluate()
        max_eval_samples = data_args.max_eval_samples if data_args.max_eval_samples is not None else len(eval_dataset)
        metrics["eval_samples"] = min(max_eval_samples, len(eval_dataset))
        trainer.log_metrics("eval", metrics)
        trainer.save_metrics("eval", metrics)

    if training_args.do_predict:
        logger.info("*** Test ***")
        for i, predict_dataset in enumerate(predict_datasets):
            metric_key_prefix=f"test_{i}"
            metrics = trainer.evaluate(eval_dataset=predict_dataset, metric_key_prefix=metric_key_prefix)
            max_predict_samples = (
                data_args.max_predict_samples if data_args.max_predict_samples is not None else len(predict_dataset)
            )
            metrics["test_samples"] = min(max_predict_samples, len(predict_dataset))
            trainer.log_metrics(metric_key_prefix, metrics)
            trainer.save_metrics(metric_key_prefix, metrics)


if __name__ == "__main__":
    main()

run_glue.py
allenai/tess-diffusion
import logging
import os
import sys
import pdb
import datasets
from sdlm.data.data_utils import load_data, tokenize_data_new
from datasets import load_from_disk, DatasetDict

import transformers
from transformers import (
    MODEL_FOR_MASKED_LM_MAPPING,
    AutoTokenizer,
    AutoModelForCausalLM,
    HfArgumentParser,
    TrainingArguments,
    set_seed,
)
from transformers.trainer_utils import get_last_checkpoint
from transformers.utils import check_min_version, send_example_telemetry
from transformers.utils.versions import require_version
from sdlm.arguments import DataTrainingArguments, ModelArguments, TrainingArguments, DiffusionArguments
from sdlm.models import RobertaDiffusionConfig, RobertaForDiffusionLM
from sdlm.trainer import DiffusionTrainer
from sdlm.schedulers import SimplexDDPMScheduler
from sdlm.inference.inference_utils import evaluate_generation
from sdlm.data.data_collator import SpanInfillingDataCollator
from sdlm.data.data_utils import split_data_to_train_validation
from transformers.trainer_callback import TrainerState


# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
check_min_version("4.25.0")

require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/language-modeling/requirements.txt")

logger = logging.getLogger(__name__)
MODEL_CONFIG_CLASSES = list(MODEL_FOR_MASKED_LM_MAPPING.keys())
MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)


def get_compute_metrics(data_args, training_args, model_args):
    # Causal language model.
    causal_model = AutoModelForCausalLM.from_pretrained(model_args.autoregressive_eval_model)
    causal_model = causal_model.to(training_args.device)
    causal_tokenizer = AutoTokenizer.from_pretrained(model_args.autoregressive_eval_model)

    is_conditional_generation = data_args.conditional_generation is not None
    prefix_lm_eval = True if data_args.conditional_generation in ["prefix_lm", "ul2", "ul2_with_unconditional", "ul2_variable"] else False
    compute_metrics = lambda results: evaluate_generation(
        results,
        data_args,
        causal_model,
        causal_tokenizer,
        is_conditional_generation,
        prefix_lm_eval=prefix_lm_eval,
        skip_special_tokens=data_args.skip_special_tokens,
        eval_for_all_metrics=training_args.eval_for_all_metrics,
    )
    return compute_metrics


def main():
    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments, DiffusionArguments))
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        # If we pass only one argument to the script and it's the path to a json file,
        # let's parse it to get our arguments.
        model_args, data_args, training_args, diffusion_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        model_args, data_args, training_args, diffusion_args = parser.parse_args_into_dataclasses()

    # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
    # information sent is the one passed as arguments along with your Python/PyTorch versions.
    send_example_telemetry("run_mlm", model_args, data_args)

    # Setup logging
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        handlers=[logging.StreamHandler(sys.stdout)],
    )

    log_level = training_args.get_process_log_level()
    logger.setLevel(log_level)
    datasets.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.enable_default_handler()
    transformers.utils.logging.enable_explicit_format()

    # Log on each process the small summary:
    logger.warning(
        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
        + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
    )
    # Set the verbosity to info of the Transformers logger (on main process only):
    logger.info(f"Training/evaluation parameters {training_args}")

    # Detecting last checkpoint.
    last_checkpoint = None
    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
            raise ValueError(
                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
                "Use --overwrite_output_dir to overcome."
            )
        elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
            logger.info(
                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
            )

    # Set seed before initializing model.
    set_seed(training_args.seed)
    config_kwargs = {
        "cache_dir": model_args.cache_dir,
        "revision": model_args.model_revision,
        "use_auth_token": True if model_args.use_auth_token else None,
    }
    config = RobertaDiffusionConfig.from_pretrained(
        model_args.model_name_or_path,
        self_condition=diffusion_args.self_condition,
        self_condition_zeros_after_softmax=diffusion_args.self_condition_zeros_after_softmax,
        deepmind_conditional=diffusion_args.deepmind_conditional,
        classifier_free_simplex_inputs=diffusion_args.classifier_free_simplex_inputs,
        classifier_free_uncond_input=diffusion_args.classifier_free_uncond_input,
        self_condition_mlp_projection=diffusion_args.self_condition_mlp_projection,
        self_condition_mix_before_weights=diffusion_args.self_condition_mix_before_weights,
        self_condition_mix_logits_before_weights=diffusion_args.self_condition_mix_logits_before_weights,
        empty_token_be_mask=diffusion_args.empty_token_be_mask,
        **config_kwargs,
    )
    tokenizer_kwargs = {
        "cache_dir": model_args.cache_dir,
        "use_fast": model_args.use_fast_tokenizer,
        "revision": model_args.model_revision,
        "use_auth_token": True if model_args.use_auth_token else None,
    }
    if model_args.tokenizer_name:
        tokenizer = AutoTokenizer.from_pretrained(model_args.tokenizer_name, **tokenizer_kwargs)
    elif model_args.model_name_or_path:
        tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path, **tokenizer_kwargs)
    else:
        raise ValueError(
            "You are instantiating a new tokenizer from scratch. This is not supported by this script."
            "You can do it from another script, save it, and load it from here, using --tokenizer_name."
        )

    if model_args.model_name_or_path:
        model = RobertaForDiffusionLM.from_pretrained(
            model_args.model_name_or_path,
            from_tf=bool(".ckpt" in model_args.model_name_or_path),
            config=config,
            cache_dir=model_args.cache_dir,
            revision=model_args.model_revision,
            use_auth_token=True if model_args.use_auth_token else None,
        )
    else:
        logger.info("Training new model from scratch")
        model = RobertaForDiffusionLM.from_config(config)

    # We resize the embeddings only when necessary to avoid index errors. If you are creating a model from scratch
    # on a small vocab and want a smaller embedding size, remove this test.
    vocab_size = model.get_input_embeddings().weight.shape[0]
    if len(tokenizer) > vocab_size:
        model.resize_token_embeddings(len(tokenizer))

    noise_scheduler = SimplexDDPMScheduler(
        num_train_timesteps=diffusion_args.num_diffusion_steps,
        beta_schedule=diffusion_args.beta_schedule,
        simplex_value=diffusion_args.simplex_value,
        clip_sample=diffusion_args.clip_sample,
        device=training_args.device,
    )
    inference_noise_scheduler = SimplexDDPMScheduler(
        num_train_timesteps=diffusion_args.num_inference_diffusion_steps,
        beta_schedule=diffusion_args.beta_schedule,
        simplex_value=diffusion_args.simplex_value,
        clip_sample=diffusion_args.clip_sample,
        device=training_args.device,
    )

    if data_args.tokenized_data_path:
        tokenized_datasets = load_from_disk(data_args.tokenized_data_path)
    else:
        raw_datasets = load_data(data_args, model_args)
        tokenized_datasets = tokenize_data_new(data_args, tokenizer, raw_datasets, training_args)

    # TODO: is this getting the same on each process?
    '''
    if "validation" not in tokenized_datasets.keys():
        tokenized_datasets = split_data_to_train_validation(data_args, tokenized_datasets, training_args.seed)
        """
        train_testvalid = tokenized_datasets["train"].train_test_split(
            test_size=data_args.validation_split_ratio, shuffle=True, seed=training_args.seed
        )
        tokenized_datasets = DatasetDict({"train": train_testvalid["train"], "validation": train_testvalid["test"]})
        """
    '''

    if training_args.do_train:
        if "train" not in tokenized_datasets:
            raise ValueError("--do_train requires a train dataset")
        train_dataset = tokenized_datasets["train"]
        if data_args.max_train_samples is not None:
            max_train_samples = min(len(train_dataset), data_args.max_train_samples)
            train_dataset = train_dataset.select(range(max_train_samples))

    if training_args.do_eval:
        if "validation" not in tokenized_datasets:
            raise ValueError("--do_eval requires a validation dataset")
        eval_dataset = tokenized_datasets["validation"]
        if data_args.max_eval_samples is not None:
            max_eval_samples = min(len(eval_dataset), data_args.max_eval_samples)
            eval_dataset = eval_dataset.select(range(max_eval_samples))

        def preprocess_logits_for_metrics(logits):
            return logits.argmax(dim=-1)

    # Data collator
    # TODO: fix lambda max_seq_length, extra_padding_ratio:
    pad_to_multiple_of_8 = data_args.line_by_line and training_args.fp16 and not data_args.pad_to_max_length
    data_collator = lambda mode: SpanInfillingDataCollator(
        mode=mode,
        data_args=data_args,
        tokenizer=tokenizer,
        max_length=data_args.max_seq_length,
        seed=training_args.seed,
        pad_to_multiple_of=8 if pad_to_multiple_of_8 else None,
        eval_context_size=data_args.eval_context_size,
    )

    if training_args.do_eval:
        compute_metrics = get_compute_metrics(data_args, training_args, model_args)

    # Initialize our Trainer
    trainer = DiffusionTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset if training_args.do_train else None,
        eval_dataset=eval_dataset if training_args.do_eval else None,
        tokenizer=tokenizer,
        data_collator=data_collator,
        compute_metrics=compute_metrics if training_args.do_eval and not training_args.without_compute_metrics else None,
        preprocess_logits_for_metrics=preprocess_logits_for_metrics if training_args.do_eval else None,
        noise_scheduler=noise_scheduler,
        diffusion_args=diffusion_args,
        data_args=data_args,
        inference_noise_scheduler=inference_noise_scheduler,
    )

    # Training
    if training_args.do_train:
        checkpoint = None
        if training_args.resume_from_checkpoint is not None:
            checkpoint = training_args.resume_from_checkpoint
        elif last_checkpoint is not None:
            checkpoint = last_checkpoint
        train_result = trainer.train(resume_from_checkpoint=checkpoint)
        trainer.save_model()  # Saves the tokenizer too for easy upload
        metrics = train_result.metrics

        max_train_samples = data_args.max_train_samples if data_args.max_train_samples is not None else len(train_dataset)
        metrics["train_samples"] = min(max_train_samples, len(train_dataset))

        trainer.log_metrics("train", metrics)
        trainer.save_metrics("train", metrics)
        trainer.save_state()

    # Evaluation
    if training_args.do_eval:
        if training_args.load_states_in_eval_from_model_path:
            trainer._load_from_checkpoint(model_args.model_name_or_path)
            trainer.state = TrainerState.load_from_json(os.path.join(model_args.model_name_or_path, "trainer_state.json"))
            trainer._load_rng_state(model_args.model_name_or_path)

        # np.save("weights.npy", model.vocab_to_hidden_dim_embed.weight.data.numpy())

        logger.info("*** Evaluate ***")
        metrics = trainer.evaluate()
        max_eval_samples = data_args.max_eval_samples if data_args.max_eval_samples is not None else len(eval_dataset)
        metrics["eval_samples"] = min(max_eval_samples, len(eval_dataset))
        trainer.log_metrics("eval", metrics)
        trainer.save_metrics("eval", metrics)


if __name__ == "__main__":
    main()

run_mlm.py
allenai/tess-diffusion
"""
Fine-tuning the library models for sequence to sequence.
"""

import logging
import os
import sys
import json 
import datasets
import nltk
import numpy as np 
from datasets import DatasetDict, Dataset, load_dataset
from transformers.trainer_callback import TrainerState
import evaluate
import transformers
from filelock import FileLock
from transformers import AutoTokenizer, HfArgumentParser, set_seed
from transformers.trainer_utils import get_last_checkpoint
from transformers.utils import check_min_version, is_offline_mode, send_example_telemetry
from transformers.utils.versions import require_version
from sdlm.arguments import ModelArguments, DataTrainingArguments, Seq2SeqTrainingArguments, DiffusionArguments
from sdlm.models import RobertaDiffusionConfig, RobertaForDiffusionLM
from sdlm.schedulers import SimplexDDPMScheduler
import pdb
from sdlm.trainer import DiffusionTrainer
from sdlm.data.data_collator import DataCollatorForSeq2Seq
from sdlm.inference.inference_utils import process_text
from sdlm.metrics.metrics import distinct_n_grams
from sdlm.data.postprocessors import postprocess_text_for_metric


# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
check_min_version("4.25.0")

require_version("datasets>=1.8.0")

logger = logging.getLogger(__name__)

try:
    nltk.data.find("tokenizers/punkt")
except (LookupError, OSError):
    if is_offline_mode():
        raise LookupError("Offline mode: run this script without TRANSFORMERS_OFFLINE first to download nltk data files")
    with FileLock(".lock") as lock:
        nltk.download("punkt", quiet=True)

def read_wikilarge(data_args):
    raw_datasets = DatasetDict()
    for split in ["train", "dev", "test"]:
        # TODO: change to f"{data_args.dataset_folder}/"
        s1s = open(f"wikilarge/s1.{split}", "r").readlines() 
        s2s = open(f"wikilarge/s2.{split}", "r").readlines()
        data = [{"original": s1, "simplification": s2} for s1, s2 in zip(s1s, s2s)]
        raw_datasets[split] = Dataset.from_list(data)
    return raw_datasets 

def read_diffuseq_datasets(data_args):
    raw_datasets = DatasetDict()
    for split in ["train", "valid", "test"]:
        dataset = load_dataset("json", data_files=f"{data_args.dataset_folder}/{split}.jsonl")["train"]
        data_split = split if split != "valid" else "dev"
        raw_datasets[data_split] = dataset
    return raw_datasets

simplification_name_mapping = {
    "wikilarge": ("original", "simplification"),
    "wiki_alignment": ("src", "trg"),
    "qqp": ("src", "trg"),
    "qg": ("src", "trg"),
    "cc": ("src", "trg")
}


def main():
    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, Seq2SeqTrainingArguments, DiffusionArguments))
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        model_args, data_args, training_args, diffusion_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        model_args, data_args, training_args, diffusion_args = parser.parse_args_into_dataclasses()

    assert data_args.max_target_length + data_args.max_source_length <= data_args.max_seq_length

    # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
    # information sent is the one passed as arguments along with your Python/PyTorch versions.
    send_example_telemetry("run_summarization", model_args, data_args)

    # Setup logging
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        handlers=[logging.StreamHandler(sys.stdout)],
    )
    log_level = training_args.get_process_log_level()
    logger.setLevel(log_level)
    datasets.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.enable_default_handler()
    transformers.utils.logging.enable_explicit_format()

    # Log on each process the small summary:
    logger.warning(
        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
        + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
    )
    logger.info(f"Training/evaluation parameters {training_args}")

    # Detecting last checkpoint.
    last_checkpoint = None
    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
            raise ValueError(
                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
                "Use --overwrite_output_dir to overcome."
            )
        elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
            logger.info(
                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
            )

    # Set seed before initializing model.
    set_seed(training_args.seed)

    if data_args.dataset_name == "wikilarge":
        raw_datasets = read_wikilarge(data_args)
    elif data_args.dataset_name in ["wiki_alignment", "qqp", "qg", "cc"]:
        raw_datasets = read_diffuseq_datasets(data_args)
    train_dataset = raw_datasets["train"]
    eval_dataset = raw_datasets["dev"]
    test_dataset = raw_datasets["test"]
    column_names = train_dataset.column_names
    config = RobertaDiffusionConfig.from_pretrained(
        model_args.model_name_or_path,
        self_condition=diffusion_args.self_condition,
        self_condition_zeros_after_softmax=diffusion_args.self_condition_zeros_after_softmax,
        deepmind_conditional=diffusion_args.deepmind_conditional,
        classifier_free_simplex_inputs=diffusion_args.classifier_free_simplex_inputs,
        classifier_free_uncond_input=diffusion_args.classifier_free_uncond_input,
        self_condition_mlp_projection=diffusion_args.self_condition_mlp_projection,
        self_condition_mix_before_weights=diffusion_args.self_condition_mix_before_weights,
        self_condition_mix_logits_before_weights=diffusion_args.self_condition_mix_logits_before_weights,
        empty_token_be_mask=diffusion_args.empty_token_be_mask,
        cache_dir=model_args.cache_dir,
        revision=model_args.model_revision,
        use_auth_token=True if model_args.use_auth_token else None,
    )
    tokenizer = AutoTokenizer.from_pretrained(
        model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        use_fast=model_args.use_fast_tokenizer,
        revision=model_args.model_revision,
        use_auth_token=True if model_args.use_auth_token else None,
    )
    if model_args.model_name_or_path:
        model = RobertaForDiffusionLM.from_pretrained(
            model_args.model_name_or_path,
            from_tf=bool(".ckpt" in model_args.model_name_or_path),
            config=config,
            cache_dir=model_args.cache_dir,
            revision=model_args.model_revision,
            use_auth_token=True if model_args.use_auth_token else None,
        )
    else:
        logger.info("Training new model from scratch")
        model = RobertaForDiffusionLM.from_config(config)

    # We resize the embeddings only when necessary to avoid index errors. If you are creating a model from scratch
    # on a small vocab and want a smaller embedding size, remove this test.
    vocab_size = model.get_input_embeddings().weight.shape[0]
    if len(tokenizer) > vocab_size:
        model.resize_token_embeddings(len(tokenizer))

    total_seq2seq_length = data_args.max_source_length + data_args.max_target_length
    if hasattr(model.config, "max_position_embeddings") and model.config.max_position_embeddings < total_seq2seq_length:
        if model_args.resize_position_embeddings is None:
            logger.warning(
                "Increasing the model's number of position embedding vectors from"
                f" {model.config.max_position_embeddings} to {total_seq2seq_length}."
            )
            # position_ids starts from `padding_idx + 1` (padding_index=1) and we therefore requires
            # 2 more position embeddings.
            model.resize_position_embeddings(total_seq2seq_length + 2)
        elif model_args.resize_position_embeddings:
            model.resize_position_embeddings(total_seq2seq_length + 2)
        else:
            raise ValueError(
                f"`max_source_length`+`max_target_length` is set to {total_seq2seq_length}, but the model only has"
                f" {model.config.max_position_embeddings} position encodings. Consider either reducing"
                f" `max_source_length`+`max_target_length` to {model.config.max_position_embeddings} or to automatically resize the"
                " model's position encodings by passing `--resize_position_embeddings`."
            )

    # We need to tokenize inputs and targets.
    # Get the column names for input/target.
    dataset_columns = simplification_name_mapping.get(data_args.dataset_name, None)
    assert dataset_columns is not None, "You need to provide the columns names."
    text_column, simplification_column = dataset_columns[0], dataset_columns[1]

    # Temporarily set max_target_length for training.
    max_target_length = data_args.max_target_length

    """
    if training_args.label_smoothing_factor > 0 and not hasattr(model, "prepare_decoder_input_ids_from_labels"):
        logger.warning(
            "label_smoothing is enabled but the `prepare_decoder_input_ids_from_labels` method is not defined for"
            f"`{model.__class__.__name__}`. This will lead to loss being calculated twice and will take up more memory"
        )
    """

    def preprocess_function(examples):
        # remove pairs where at least one record is None

        inputs, targets = [], []
        for i in range(len(examples[text_column])):
            if examples[text_column][i] and examples[simplification_column][i]:
                inputs.append(examples[text_column][i])
                targets.append(examples[simplification_column][i])
        # TODO: we need to process first the target, then cut the inputs to the max_length-target length to use the
        # maximum number of tokens.
        model_inputs = tokenizer(inputs, max_length=data_args.max_source_length, padding=False, truncation=True)
        # Tokenize targets with the `text_target` keyword argument
        labels = tokenizer(text_target=targets, max_length=max_target_length, padding=False, truncation=True)
        model_inputs["labels"] = labels["input_ids"]
        return model_inputs

    if training_args.do_train:
        if data_args.max_train_samples is not None:
            max_train_samples = min(len(train_dataset), data_args.max_train_samples)
            train_dataset = train_dataset.select(range(max_train_samples))
        with training_args.main_process_first(desc="train dataset map pre-processing"):
            train_dataset = train_dataset.map(
                preprocess_function,
                batched=True,
                num_proc=data_args.preprocessing_num_workers,
                remove_columns=column_names,
                load_from_cache_file=not data_args.overwrite_cache,
                desc="Running tokenizer on train dataset",
            )

    if training_args.do_eval:
        max_target_length = data_args.val_max_target_length
        if data_args.max_eval_samples is not None:
            max_eval_samples = min(len(eval_dataset), data_args.max_eval_samples)
            eval_dataset = eval_dataset.select(range(max_eval_samples))
        with training_args.main_process_first(desc="validation dataset map pre-processing"):
            eval_dataset = eval_dataset.map(
                preprocess_function,
                batched=True,
                num_proc=data_args.preprocessing_num_workers,
                remove_columns=column_names,
                load_from_cache_file=not data_args.overwrite_cache,
                desc="Running tokenizer on validation dataset",
            )

    if training_args.do_predict:
        max_target_length = data_args.val_max_target_length
        test_dataset = raw_datasets["test"]
        if data_args.max_predict_samples is not None:
            max_predict_samples = min(len(test_dataset), data_args.max_predict_samples)
            test_dataset = test_dataset.select(range(max_predict_samples))
        with training_args.main_process_first(desc="prediction dataset map pre-processing"):
            test_dataset = test_dataset.map(
                preprocess_function,
                batched=True,
                num_proc=data_args.preprocessing_num_workers,
                remove_columns=column_names,
                load_from_cache_file=not data_args.overwrite_cache,
                desc="Running tokenizer on prediction dataset",
            )

    def preprocess_logits_for_metrics(logits):
        return logits.argmax(dim=-1)

    # TODO: we may want to add predict back.

    # Data collator. To be consistent with the run_mlm.py we need to add `mode`.
    data_collator = lambda mode: DataCollatorForSeq2Seq(
        tokenizer,
        # Note that if you do not use `pad_to_max_length`, this becomes very slow on multi-gpus.
        padding="max_length" if data_args.pad_to_max_length else True,
        max_length=data_args.max_seq_length,
        pad_to_multiple_of=8 if training_args.fp16 else None,
    )
    noise_scheduler = SimplexDDPMScheduler(
        num_train_timesteps=diffusion_args.num_diffusion_steps,
        beta_schedule=diffusion_args.beta_schedule,
        simplex_value=diffusion_args.simplex_value,
        clip_sample=diffusion_args.clip_sample,
        device=training_args.device,
    )
    inference_noise_scheduler = SimplexDDPMScheduler(
        num_train_timesteps=diffusion_args.num_inference_diffusion_steps,
        beta_schedule=diffusion_args.beta_schedule,
        simplex_value=diffusion_args.simplex_value,
        clip_sample=diffusion_args.clip_sample,
        device=training_args.device,
    )

    # Metric
    eval_metrics = {
        "sari": evaluate.load("sari"),
        "bleu": evaluate.load("bleu"),
        "bertscore": evaluate.load("bertscore"),
        "bertscore_them": evaluate.load("bertscore"),
        "rouge": evaluate.load("rouge"),
        "dist":  distinct_n_grams
    }

    def compute_metrics(results):
        keys = ["pred_texts_from_simplex_masked", "pred_texts_from_logits_masked"]
        metrics = {}
        for key in keys:
            decoded_preds_original = process_text(results[key]) if not data_args.skip_special_tokens else results[key]
            decoded_labels_original = process_text(results["gold_texts_masked"]) if not data_args.skip_special_tokens else results["gold_texts_masked"]
            sources = results["prefixes"]
            for metric_name, metric in eval_metrics.items():
                scale = 100 if metric_name != "sari" else 1
                if metric_name == "sari":
                    decoded_preds, decoded_labels, sources = postprocess_text_for_metric("sari", decoded_preds_original, decoded_labels_original, sources)
                    decoded_labels = [[decoded_label] for decoded_label in decoded_labels]
                    key_metrics = metric.compute(sources=sources, predictions=decoded_preds, references=decoded_labels)
                elif metric_name == "bleu":
                    decoded_preds, decoded_labels = postprocess_text_for_metric("bleu", decoded_preds_original, decoded_labels_original)
                    key_metrics = {"bleu": metric.compute(predictions=decoded_preds, references=decoded_labels)["bleu"]}
                elif metric_name == "bertscore":
                    decoded_preds, decoded_labels = postprocess_text_for_metric("bertscore", decoded_preds_original, decoded_labels_original)
                    key_metrics = {"bert_score": np.mean(metric.compute(predictions=decoded_preds, references=decoded_labels,  lang="en")['f1'])}
                elif metric_name == "bertscore_them":
                    decoded_preds, decoded_labels = postprocess_text_for_metric("bertscore_them", decoded_preds_original, decoded_labels_original)
                    key_metrics = {"bert_score_them": np.mean(metric.compute(predictions=decoded_preds, references=decoded_labels, model_type='microsoft/deberta-xlarge-mnli', lang="en")['f1'])}
                elif metric_name == "rouge":
                    decoded_preds, decoded_labels = postprocess_text_for_metric("rouge", decoded_preds_original, decoded_labels_original)
                    key_metrics = metric.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True)
                elif metric_name == "dist":
                    decoded_preds = postprocess_text_for_metric("dist", decoded_preds_original)
                    key_metrics = metric(decoded_preds)

                key_metrics = {k: round(v*scale, 2) for k, v in key_metrics.items()}
                key_metrics = {f"{key}_{k}": v for k, v in key_metrics.items()}
                metrics.update(key_metrics)
        return metrics

    # Initialize our Trainer
    trainer = DiffusionTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset if training_args.do_train else None,
        eval_dataset=eval_dataset if training_args.do_eval else None,
        tokenizer=tokenizer,
        data_collator=data_collator,
        compute_metrics=compute_metrics if (training_args.do_eval or training_args.do_predict) else None,
        preprocess_logits_for_metrics=preprocess_logits_for_metrics if (training_args.do_eval or training_args.do_predict) else None,
        noise_scheduler=noise_scheduler,
        diffusion_args=diffusion_args,
        data_args=data_args,
        inference_noise_scheduler=inference_noise_scheduler,
    )

    # Training
    if training_args.do_train:
        checkpoint = None
        if training_args.resume_from_checkpoint is not None:
            checkpoint = training_args.resume_from_checkpoint
        elif last_checkpoint is not None:
            checkpoint = last_checkpoint
        train_result = trainer.train(resume_from_checkpoint=checkpoint)
        trainer.save_model()  # Saves the tokenizer too for easy upload

        metrics = train_result.metrics
        max_train_samples = data_args.max_train_samples if data_args.max_train_samples is not None else len(train_dataset)
        metrics["train_samples"] = min(max_train_samples, len(train_dataset))

        trainer.log_metrics("train", metrics)
        trainer.save_metrics("train", metrics)
        trainer.save_state()


    # We will load the best model here to avoid an issue when do_train is not set.
    if training_args.load_states_in_eval_from_model_path and not training_args.do_train:
        trainer.state = TrainerState.load_from_json(os.path.join(model_args.model_name_or_path, "trainer_state.json"))
        if training_args.load_best_model_at_end and trainer.state.best_model_checkpoint is not None:
            checkpoint_path = trainer.state.best_model_checkpoint
        else:
            checkpoint_path = model_args.model_name_or_path
        trainer._load_from_checkpoint(checkpoint_path)
        trainer._load_rng_state(checkpoint_path)

    # Evaluation
    results = {}
    num_beams = data_args.num_beams if data_args.num_beams is not None else training_args.generation_num_beams
    if training_args.do_eval:
        logger.info("*** Evaluate ***")
        # TODO: num_beans should be added for ours as well.
        # metrics = trainer.evaluate(max_length=max_length, num_beams=num_beams, metric_key_prefix="eval")
        metrics = trainer.evaluate()
        max_eval_samples = data_args.max_eval_samples if data_args.max_eval_samples is not None else len(eval_dataset)
        metrics["eval_samples"] = min(max_eval_samples, len(eval_dataset))

        trainer.log_metrics("eval", metrics)
        trainer.save_metrics("eval", metrics)

    if training_args.do_predict:
        logger.info("*** Test ***")
        # TODO: num_beans should be added for ours as well.
        # metrics = trainer.evaluate(max_length=max_length, num_beams=num_beams, metric_key_prefix="eval")
        metrics = trainer.evaluate(test_dataset, metric_key_prefix="test")
        max_test_samples = data_args.max_predict_samples if data_args.max_predict_samples is not None else len(test_dataset)
        metrics["test_samples"] = min(max_test_samples, len(test_dataset))
        trainer.log_metrics(f"test", metrics)
        trainer.save_metrics(f"test", metrics)

    # TODO: we may want to add predict part back.
    return results


if __name__ == "__main__":
    main()

run_simplification.py
allenai/tess-diffusion
"""
Fine-tuning the library models for sequence to sequence.
"""

import logging
import os
import sys

import datasets
import nltk

import evaluate
import transformers
from transformers.trainer_callback import TrainerState
from transformers import AutoTokenizer, HfArgumentParser, set_seed
from transformers.trainer_utils import get_last_checkpoint
from transformers.utils import check_min_version, send_example_telemetry
from transformers.utils.versions import require_version
from sdlm.data.data_utils import load_data
from sdlm.arguments import ModelArguments, DataTrainingArguments, Seq2SeqTrainingArguments, DiffusionArguments
from sdlm.models import RobertaDiffusionConfig, RobertaForDiffusionLM
from sdlm.schedulers import SimplexDDPMScheduler
import pdb
from sdlm.trainer import DiffusionTrainer
from sdlm.data.data_collator import DataCollatorForSeq2Seq
from sdlm.inference.inference_utils import process_text
from sdlm.data.postprocessors import postprocess_text_for_metric

# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
check_min_version("4.25.0")
require_version("datasets>=1.8.0")
logger = logging.getLogger(__name__)

summarization_name_mapping = {
    "amazon_reviews_multi": ("review_body", "review_title"),
    "big_patent": ("description", "abstract"),
    "cnn_dailymail": ("article", "highlights"),
    "orange_sum": ("text", "summary"),
    "pn_summary": ("article", "summary"),
    "psc": ("extract_text", "summary_text"),
    "samsum": ("dialogue", "summary"),
    "thaisum": ("body", "summary"),
    "xglue": ("news_body", "news_title"),
    "xsum": ("document", "summary"),
    "wiki_summary": ("article", "highlights"),
    "multi_news": ("document", "summary"),
}


def main():
    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, Seq2SeqTrainingArguments, DiffusionArguments))
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        model_args, data_args, training_args, diffusion_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        model_args, data_args, training_args, diffusion_args = parser.parse_args_into_dataclasses()

    assert data_args.max_target_length + data_args.max_source_length <= data_args.max_seq_length

    # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
    # information sent is the one passed as arguments along with your Python/PyTorch versions.
    send_example_telemetry("run_summarization", model_args, data_args)

    # Setup logging
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        handlers=[logging.StreamHandler(sys.stdout)],
    )
    log_level = training_args.get_process_log_level()
    logger.setLevel(log_level)
    datasets.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.enable_default_handler()
    transformers.utils.logging.enable_explicit_format()

    # Log on each process the small summary:
    logger.warning(
        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
        + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
    )
    logger.info(f"Training/evaluation parameters {training_args}")

    # Detecting last checkpoint.
    last_checkpoint = None
    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
            raise ValueError(
                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
                "Use --overwrite_output_dir to overcome."
            )
        elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
            logger.info(
                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
            )

    # Set seed before initializing model.
    set_seed(training_args.seed)

    raw_datasets = load_data(data_args, model_args)

    config = RobertaDiffusionConfig.from_pretrained(
        model_args.model_name_or_path,
        self_condition=diffusion_args.self_condition,
        self_condition_zeros_after_softmax=diffusion_args.self_condition_zeros_after_softmax,
        deepmind_conditional=diffusion_args.deepmind_conditional,
        classifier_free_simplex_inputs=diffusion_args.classifier_free_simplex_inputs,
        classifier_free_uncond_input=diffusion_args.classifier_free_uncond_input,
        self_condition_mlp_projection=diffusion_args.self_condition_mlp_projection,
        self_condition_mix_before_weights=diffusion_args.self_condition_mix_before_weights,
        self_condition_mix_logits_before_weights=diffusion_args.self_condition_mix_logits_before_weights,
        empty_token_be_mask=diffusion_args.empty_token_be_mask,
        cache_dir=model_args.cache_dir,
        revision=model_args.model_revision,
        use_auth_token=True if model_args.use_auth_token else None,
    )
    tokenizer = AutoTokenizer.from_pretrained(
        model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        use_fast=model_args.use_fast_tokenizer,
        revision=model_args.model_revision,
        use_auth_token=True if model_args.use_auth_token else None,
    )
    if model_args.model_name_or_path:
        model = RobertaForDiffusionLM.from_pretrained(
            model_args.model_name_or_path,
            from_tf=bool(".ckpt" in model_args.model_name_or_path),
            config=config,
            cache_dir=model_args.cache_dir,
            revision=model_args.model_revision,
            use_auth_token=True if model_args.use_auth_token else None,
        )
    else:
        logger.info("Training new model from scratch")
        model = RobertaForDiffusionLM.from_config(config)

    # We resize the embeddings only when necessary to avoid index errors. If you are creating a model from scratch
    # on a small vocab and want a smaller embedding size, remove this test.
    vocab_size = model.get_input_embeddings().weight.shape[0]
    if len(tokenizer) > vocab_size:
        model.resize_token_embeddings(len(tokenizer))

    total_seq2seq_length = data_args.max_source_length + data_args.max_target_length
    if hasattr(model.config, "max_position_embeddings") and model.config.max_position_embeddings < total_seq2seq_length:
        if model_args.resize_position_embeddings is None:
            logger.warning(
                "Increasing the model's number of position embedding vectors from"
                f" {model.config.max_position_embeddings} to {total_seq2seq_length}."
            )
            # position_ids starts from `padding_idx + 1` (padding_index=1) and we therefore requires
            # 2 more position embeddings.
            model.resize_position_embeddings(
                total_seq2seq_length + 2, with_alternatation=model_args.resize_position_embeddings_alternatively
            )
        elif model_args.resize_position_embeddings:
            model.resize_position_embeddings(
                total_seq2seq_length + 2, with_alternatation=model_args.resize_position_embeddings_alternatively
            )
        else:
            raise ValueError(
                f"`max_source_length`+`max_target_length` is set to {total_seq2seq_length}, but the model only has"
                f" {model.config.max_position_embeddings} position encodings. Consider either reducing"
                f" `max_source_length`+`max_target_length` to {model.config.max_position_embeddings} or to automatically resize the"
                " model's position encodings by passing `--resize_position_embeddings`."
            )

    # Preprocessing the datasets.
    # We need to tokenize inputs and targets.
    if training_args.do_train:
        column_names = raw_datasets["train"].column_names
    elif training_args.do_eval:
        column_names = raw_datasets["validation"].column_names
    elif training_args.do_predict:
        column_names = raw_datasets["test"].column_names
    else:
        logger.info("There is nothing to do. Please pass `do_train`, `do_eval` and/or `do_predict`.")
        return

    # Get the column names for input/target.
    dataset_columns = summarization_name_mapping.get(data_args.dataset_name, None)
    assert dataset_columns is not None, "You need to provide the columns names."
    text_column, summary_column = dataset_columns[0], dataset_columns[1]

    # Temporarily set max_target_length for training.
    max_target_length = data_args.max_target_length

    """
    if training_args.label_smoothing_factor > 0 and not hasattr(model, "prepare_decoder_input_ids_from_labels"):
        logger.warning(
            "label_smoothing is enabled but the `prepare_decoder_input_ids_from_labels` method is not defined for"
            f"`{model.__class__.__name__}`. This will lead to loss being calculated twice and will take up more memory"
        )
    """

    def preprocess_function(examples):
        # remove pairs where at least one record is None

        inputs, targets = [], []
        for i in range(len(examples[text_column])):
            if examples[text_column][i] and examples[summary_column][i]:
                inputs.append(examples[text_column][i])
                targets.append(examples[summary_column][i])
        # TODO: we need to process first the target, then cut the inputs to the max_length-target length to use the
        # maximum number of tokens.
        model_inputs = tokenizer(inputs, max_length=data_args.max_source_length, padding=False, truncation=True)
        # Tokenize targets with the `text_target` keyword argument
        labels = tokenizer(text_target=targets, max_length=max_target_length, padding=False, truncation=True)
        model_inputs["labels"] = labels["input_ids"]
        return model_inputs

    if training_args.do_train:
        if "train" not in raw_datasets:
            raise ValueError("--do_train requires a train dataset")
        train_dataset = raw_datasets["train"]
        if data_args.max_train_samples is not None:
            max_train_samples = min(len(train_dataset), data_args.max_train_samples)
            train_dataset = train_dataset.select(range(max_train_samples))
        with training_args.main_process_first(desc="train dataset map pre-processing"):
            train_dataset = train_dataset.map(
                preprocess_function,
                batched=True,
                num_proc=data_args.preprocessing_num_workers,
                remove_columns=column_names,
                load_from_cache_file=not data_args.overwrite_cache,
                desc="Running tokenizer on train dataset",
            )

    if training_args.do_eval:
        max_target_length = data_args.val_max_target_length
        if "validation" not in raw_datasets:
            raise ValueError("--do_eval requires a validation dataset")
        eval_dataset = raw_datasets["validation"]
        if data_args.max_eval_samples is not None:
            max_eval_samples = min(len(eval_dataset), data_args.max_eval_samples)
            eval_dataset = eval_dataset.select(range(max_eval_samples))
        with training_args.main_process_first(desc="validation dataset map pre-processing"):
            eval_dataset = eval_dataset.map(
                preprocess_function,
                batched=True,
                num_proc=data_args.preprocessing_num_workers,
                remove_columns=column_names,
                load_from_cache_file=not data_args.overwrite_cache,
                desc="Running tokenizer on validation dataset",
            )
    def preprocess_logits_for_metrics(logits):
        return logits.argmax(dim=-1)

    if training_args.do_predict:
        max_target_length = data_args.val_max_target_length
        if "test" not in raw_datasets:
            raise ValueError("--do_predict requires a test dataset")
        test_dataset = raw_datasets["test"]
        if data_args.max_predict_samples is not None:
            max_predict_samples = min(len(test_dataset), data_args.max_predict_samples)
            test_dataset = test_dataset.select(range(max_predict_samples))
        with training_args.main_process_first(desc="prediction dataset map pre-processing"):
            test_dataset = test_dataset.map(
                preprocess_function,
                batched=True,
                num_proc=data_args.preprocessing_num_workers,
                remove_columns=column_names,
                load_from_cache_file=not data_args.overwrite_cache,
                desc="Running tokenizer on prediction dataset",
            )

    # TODO: we may want to add predict back.

    # Data collator. To be consistent with the run_mlm.py we need to add `mode`.
    data_collator = lambda mode: DataCollatorForSeq2Seq(
        tokenizer,
        # Note that if you do not use `pad_to_max_length`, this becomes very slow on multi-gpus.
        padding="max_length" if data_args.pad_to_max_length else True,
        max_length=data_args.max_seq_length,
        pad_to_multiple_of=8 if training_args.fp16 else None,
    )
    noise_scheduler = SimplexDDPMScheduler(
        num_train_timesteps=diffusion_args.num_diffusion_steps,
        beta_schedule=diffusion_args.beta_schedule,
        simplex_value=diffusion_args.simplex_value,
        clip_sample=diffusion_args.clip_sample,
        device=training_args.device,
    )
    inference_noise_scheduler = SimplexDDPMScheduler(
        num_train_timesteps=diffusion_args.num_inference_diffusion_steps,
        beta_schedule=diffusion_args.beta_schedule,
        simplex_value=diffusion_args.simplex_value,
        clip_sample=diffusion_args.clip_sample,
        device=training_args.device,
    )

    # Metric
    metric = evaluate.load("rouge")

    def compute_metrics(results):
        keys = ["pred_texts_from_simplex_masked", "pred_texts_from_logits_masked"]
        metrics = {}
        for key in keys:
            decoded_preds = process_text(results[key]) if not data_args.skip_special_tokens else results[key]
            # Note that since decoded_labels is getting updated after post-process, we
            # need to compute it here for each key.
            decoded_labels = process_text(results["gold_texts_masked"]) if not data_args.skip_special_tokens else  results["gold_texts_masked"]
            decoded_preds, decoded_labels = postprocess_text_for_metric("rouge", decoded_preds, decoded_labels)
            key_metrics = metric.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True)
            key_metrics = {k: round(v * 100, 4) for k, v in key_metrics.items()}
            key_metrics = {f"{key}_{k}": v for k, v in key_metrics.items()}
            metrics.update(key_metrics)
        return metrics

    # Initialize our Trainer
    trainer = DiffusionTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset if training_args.do_train else None,
        eval_dataset=eval_dataset if training_args.do_eval else None,
        tokenizer=tokenizer,
        data_collator=data_collator,
        compute_metrics=compute_metrics if (training_args.do_eval or training_args.do_predict) else None,
        preprocess_logits_for_metrics=preprocess_logits_for_metrics if (training_args.do_eval or training_args.do_predict) else None,
        noise_scheduler=noise_scheduler,
        diffusion_args=diffusion_args,
        data_args=data_args,
        inference_noise_scheduler=inference_noise_scheduler,
    )

    # Training
    if training_args.do_train:
        checkpoint = None
        if training_args.resume_from_checkpoint is not None:
            checkpoint = training_args.resume_from_checkpoint
        elif last_checkpoint is not None:
            checkpoint = last_checkpoint
        train_result = trainer.train(resume_from_checkpoint=checkpoint)
        trainer.save_model()  # Saves the tokenizer too for easy upload

        metrics = train_result.metrics
        max_train_samples = data_args.max_train_samples if data_args.max_train_samples is not None else len(train_dataset)
        metrics["train_samples"] = min(max_train_samples, len(train_dataset))

        trainer.log_metrics("train", metrics)
        trainer.save_metrics("train", metrics)
        trainer.save_state()

    # We will load the best model here to avoid an issue when do_train is not set.
    if training_args.load_states_in_eval_from_model_path and not training_args.do_train:
        trainer.state = TrainerState.load_from_json(os.path.join(model_args.model_name_or_path, "trainer_state.json"))
        if training_args.load_best_model_at_end and trainer.state.best_model_checkpoint is not None:
            checkpoint_path = trainer.state.best_model_checkpoint
        else:
            checkpoint_path = model_args.model_name_or_path
        trainer._load_from_checkpoint(checkpoint_path)
        trainer._load_rng_state(checkpoint_path)

    # Evaluation
    results = {}
    max_length = (
        training_args.generation_max_length
        if training_args.generation_max_length is not None
        else data_args.val_max_target_length
    )
    num_beams = data_args.num_beams if data_args.num_beams is not None else training_args.generation_num_beams
    if training_args.do_eval:
        logger.info("*** Evaluate ***")
        # TODO: num_beans should be added for ours as well.
        # metrics = trainer.evaluate(max_length=max_length, num_beams=num_beams, metric_key_prefix="eval")
        metrics = trainer.evaluate()
        max_eval_samples = data_args.max_eval_samples if data_args.max_eval_samples is not None else len(eval_dataset)
        metrics["eval_samples"] = min(max_eval_samples, len(eval_dataset))
        trainer.log_metrics("eval", metrics)
        trainer.save_metrics("eval", metrics)

    if training_args.do_predict:
        logger.info("*** Test ***")
        metrics = trainer.evaluate(test_dataset, metric_key_prefix="test")
        max_predict_samples = data_args.max_predict_samples if data_args.max_predict_samples is not None else len(test_dataset)
        metrics["test_samples"] = min(max_predict_samples, len(test_dataset))
        trainer.log_metrics("test", metrics)
        trainer.save_metrics("test", metrics)

    # TODO: we may want to add predict part back.
    return results


if __name__ == "__main__":
    main()

run_summarization.py
allenai/tess-diffusion
"""
Fine-tuning the library models for sequence to sequence.
"""

import logging
import os
import sys

import datasets
import numpy as np
from sdlm.data.data_utils import load_data

import evaluate
import transformers
from transformers import AutoTokenizer, HfArgumentParser, set_seed
from transformers.trainer_utils import get_last_checkpoint
from transformers.utils import check_min_version, send_example_telemetry
from transformers.utils.versions import require_version
from sdlm.data.data_utils import load_data
from sdlm.arguments import ModelArguments, DataTrainingArguments, Seq2SeqTrainingArguments, DiffusionArguments
from sdlm.models import XLMRobertaForDiffusionLM, XLMRobertaDiffusionConfig
from sdlm.models import RobertaDiffusionConfig, RobertaForDiffusionLM
from sdlm.schedulers import SimplexDDPMScheduler
import pdb
from sdlm.trainer import DiffusionTrainer
from sdlm.data.data_collator import DataCollatorForSeq2Seq
from sdlm.inference.inference_utils import process_text
import torch
from sdlm.data.postprocessors import postprocess_text_for_metric

# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
check_min_version("4.25.0")

require_version("datasets>=1.8.0")

logger = logging.getLogger(__name__)


def main():
    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, Seq2SeqTrainingArguments, DiffusionArguments))
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        model_args, data_args, training_args, diffusion_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        model_args, data_args, training_args, diffusion_args = parser.parse_args_into_dataclasses()

    assert data_args.max_target_length + data_args.max_source_length <= data_args.max_seq_length
    # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
    # information sent is the one passed as arguments along with your Python/PyTorch versions.
    send_example_telemetry("run_translation", model_args, data_args)

    # Setup logging
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        handlers=[logging.StreamHandler(sys.stdout)],
    )
    log_level = training_args.get_process_log_level()
    logger.setLevel(log_level)
    datasets.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.enable_default_handler()
    transformers.utils.logging.enable_explicit_format()

    # Log on each process the small summary:
    logger.warning(
        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
        + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
    )
    logger.info(f"Training/evaluation parameters {training_args}")

    # Detecting last checkpoint.
    last_checkpoint = None
    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
            raise ValueError(
                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
                "Use --overwrite_output_dir to overcome."
            )
        elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
            logger.info(
                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
            )

    # Set seed before initializing model.
    set_seed(training_args.seed)

    raw_datasets = load_data(data_args, model_args)

    config = RobertaDiffusionConfig.from_pretrained(
        model_args.model_name_or_path,
        self_condition=diffusion_args.self_condition,
        self_condition_zeros_after_softmax=diffusion_args.self_condition_zeros_after_softmax,
        deepmind_conditional=diffusion_args.deepmind_conditional,
        classifier_free_simplex_inputs=diffusion_args.classifier_free_simplex_inputs,
        classifier_free_uncond_input=diffusion_args.classifier_free_uncond_input,
        self_condition_mlp_projection=diffusion_args.self_condition_mlp_projection,
        self_condition_mix_before_weights=diffusion_args.self_condition_mix_before_weights,
        self_condition_mix_logits_before_weights=diffusion_args.self_condition_mix_logits_before_weights,
        empty_token_be_mask=diffusion_args.empty_token_be_mask,
        cache_dir=model_args.cache_dir,
        revision=model_args.model_revision,
        use_auth_token=True if model_args.use_auth_token else None,
    )
    tokenizer = AutoTokenizer.from_pretrained(
        model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        use_fast=model_args.use_fast_tokenizer,
        revision=model_args.model_revision,
        use_auth_token=True if model_args.use_auth_token else None,
    )
    if model_args.model_name_or_path:
        model = RobertaForDiffusionLM.from_pretrained(
            model_args.model_name_or_path,
            from_tf=bool(".ckpt" in model_args.model_name_or_path),
            config=config,
            cache_dir=model_args.cache_dir,
            revision=model_args.model_revision,
            use_auth_token=True if model_args.use_auth_token else None,
        )
    else:
        logger.info("Training new model from scratch")
        model = RobertaForDiffusionLM.from_config(config)

    # We resize the embeddings only when necessary to avoid index errors. If you are creating a model from scratch
    # on a small vocab and want a smaller embedding size, remove this test.
    vocab_size = model.get_input_embeddings().weight.shape[0]
    if len(tokenizer) > vocab_size:
        model.resize_token_embeddings(len(tokenizer))

    # TODO: is this needed?
    """
    if model.config.decoder_start_token_id is None:
        raise ValueError("Make sure that `config.decoder_start_token_id` is correctly defined")
    """

    # Preprocessing the datasets.
    # We need to tokenize inputs and targets.
    if training_args.do_train:
        column_names = raw_datasets["train"].column_names
    elif training_args.do_eval:
        column_names = raw_datasets["validation"].column_names
    elif training_args.do_predict:
        column_names = raw_datasets["test"].column_names
    else:
        logger.info("There is nothing to do. Please pass `do_train`, `do_eval` and/or `do_predict`.")
        return

    # Get the language codes for input/target.
    source_lang = data_args.source_lang.split("_")[0]
    target_lang = data_args.target_lang.split("_")[0]

    # Temporarily set max_target_length for training.
    max_target_length = data_args.max_target_length
    # padding = "max_length" if data_args.pad_to_max_length else False

    """
    if training_args.label_smoothing_factor > 0 and not hasattr(model, "prepare_decoder_input_ids_from_labels"):
        logger.warning(
            "label_smoothing is enabled but the `prepare_decoder_input_ids_from_labels` method is not defined for"
            f"`{model.__class__.__name__}`. This will lead to loss being calculated twice and will take up more memory"
        )
    """

    def preprocess_function(examples):
        inputs = [ex[source_lang] for ex in examples["translation"]]
        targets = [ex[target_lang] for ex in examples["translation"]]
        model_inputs = tokenizer(inputs, max_length=data_args.max_source_length, padding=False, truncation=True)
        # Tokenize targets with the `text_target` keyword argument
        labels = tokenizer(text_target=targets, max_length=max_target_length, padding=False, truncation=True)
        model_inputs["labels"] = labels["input_ids"]
        return model_inputs

    if training_args.do_train:
        if "train" not in raw_datasets:
            raise ValueError("--do_train requires a train dataset")
        train_dataset = raw_datasets["train"]
        if data_args.max_train_samples is not None:
            max_train_samples = min(len(train_dataset), data_args.max_train_samples)
            train_dataset = train_dataset.select(range(max_train_samples))
        with training_args.main_process_first(desc="train dataset map pre-processing"):
            train_dataset = train_dataset.map(
                preprocess_function,
                batched=True,
                num_proc=data_args.preprocessing_num_workers,
                remove_columns=column_names,
                load_from_cache_file=not data_args.overwrite_cache,
                desc="Running tokenizer on train dataset",
            )

    if training_args.do_eval:
        max_target_length = data_args.val_max_target_length
        if "validation" not in raw_datasets:
            raise ValueError("--do_eval requires a validation dataset")
        eval_dataset = raw_datasets["validation"]
        if data_args.max_eval_samples is not None:
            max_eval_samples = min(len(eval_dataset), data_args.max_eval_samples)
            eval_dataset = eval_dataset.select(range(max_eval_samples))
        with training_args.main_process_first(desc="validation dataset map pre-processing"):
            eval_dataset = eval_dataset.map(
                preprocess_function,
                batched=True,
                num_proc=data_args.preprocessing_num_workers,
                remove_columns=column_names,
                load_from_cache_file=not data_args.overwrite_cache,
                desc="Running tokenizer on validation dataset",
            )

    def preprocess_logits_for_metrics(logits):
        return logits.argmax(dim=-1)

    # TODO: we may want to add predict back.

    # Data collator. To be consistent with the run_mlm.py we need to add `mode`.
    data_collator = lambda mode: DataCollatorForSeq2Seq(
        tokenizer,
        # Note that if you do not use `pad_to_max_length`, this becomes very slow on multi-gpus.
        padding="max_length" if data_args.pad_to_max_length else True,
        max_length=data_args.max_seq_length,
        pad_to_multiple_of=8 if training_args.fp16 else None,
    )
    noise_scheduler = SimplexDDPMScheduler(
        num_train_timesteps=diffusion_args.num_diffusion_steps,
        beta_schedule=diffusion_args.beta_schedule,
        simplex_value=diffusion_args.simplex_value,
        clip_sample=diffusion_args.clip_sample,
        device=training_args.device,
    )
    inference_noise_scheduler = SimplexDDPMScheduler(
        num_train_timesteps=diffusion_args.num_inference_diffusion_steps,
        beta_schedule=diffusion_args.beta_schedule,
        simplex_value=diffusion_args.simplex_value,
        clip_sample=diffusion_args.clip_sample,
        device=training_args.device,
    )

    # Metric
    metric = evaluate.load("sacrebleu")

    def compute_metrics(results):
        keys = ["pred_texts_from_simplex_masked", "pred_texts_from_logits_masked"]
        metrics = {}
        for key in keys:
            decoded_preds = process_text(results[key])
            # Note that since decoded_labels is getting updated after post-process, we
            # need to compute it here for each key.
            decoded_labels = process_text(results["gold_texts_masked"])
            decoded_preds, decoded_labels = postprocess_text_for_metric("bleu", decoded_preds, decoded_labels)
            key_metrics = metric.compute(predictions=decoded_preds, references=decoded_labels)
            key_metrics = {"bleu": key_metrics["score"]}
            key_metrics = {k: round(v, 4) for k, v in key_metrics.items()}
            key_metrics = {f"{key}_{k}": v for k, v in key_metrics.items()}
            metrics.update(key_metrics)
        return metrics

    # Initialize our Trainer
    trainer = DiffusionTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset if training_args.do_train else None,
        eval_dataset=eval_dataset if training_args.do_eval else None,
        tokenizer=tokenizer,
        data_collator=data_collator,
        compute_metrics=compute_metrics if (training_args.do_eval or training_args.do_predict) else None,
        preprocess_logits_for_metrics=preprocess_logits_for_metrics if (training_args.do_eval or training_args.do_predict) else None,
        noise_scheduler=noise_scheduler,
        diffusion_args=diffusion_args,
        data_args=data_args,
        inference_noise_scheduler=inference_noise_scheduler,
    )

    # Training
    if training_args.do_train:
        checkpoint = None
        if training_args.resume_from_checkpoint is not None:
            checkpoint = training_args.resume_from_checkpoint
        elif last_checkpoint is not None:
            checkpoint = last_checkpoint
        train_result = trainer.train(resume_from_checkpoint=checkpoint)
        trainer.save_model()  # Saves the tokenizer too for easy upload

        metrics = train_result.metrics
        max_train_samples = data_args.max_train_samples if data_args.max_train_samples is not None else len(train_dataset)
        metrics["train_samples"] = min(max_train_samples, len(train_dataset))

        trainer.log_metrics("train", metrics)
        trainer.save_metrics("train", metrics)
        trainer.save_state()

    # Evaluation
    results = {}
    max_length = (
        training_args.generation_max_length
        if training_args.generation_max_length is not None
        else data_args.val_max_target_length
    )
    num_beams = data_args.num_beams if data_args.num_beams is not None else training_args.generation_num_beams
    if training_args.do_eval:
        logger.info("*** Evaluate ***")
        # TODO: num_beans should be added for ours as well.
        # metrics = trainer.evaluate(max_length=max_length, num_beams=num_beams, metric_key_prefix="eval")
        metrics = trainer.evaluate()
        max_eval_samples = data_args.max_eval_samples if data_args.max_eval_samples is not None else len(eval_dataset)
        metrics["eval_samples"] = min(max_eval_samples, len(eval_dataset))
        trainer.log_metrics("eval", metrics)
        trainer.save_metrics("eval", metrics)
    
    if training_args.do_predict:
        logger.info("*** Test ***")
        # TODO: num_beans should be added for ours as well.
        # metrics = trainer.evaluate(max_length=max_length, num_beams=num_beams, metric_key_prefix="eval")
        metrics = trainer.evaluate(test_dataset, metric_key_prefix="test")
        max_predict_samples = data_args.max_predict_samples if data_args.max_predict_samples is not None else len(test_dataset)
        metrics["test_samples"] = min(max_predict_samples, len(test_dataset))
        trainer.log_metrics("test", metrics)
        trainer.save_metrics("test", metrics)
    return results

if __name__ == "__main__":
    main()

run_translation.py
allenai/tess-diffusion
'''
SSD LM Inference script from https://colab.research.google.com/drive/1vNKqvzzJQp3k89QPuns5ibsq-VNC9wGN?usp=sharing
'''

import os
import sys
import torch
import transformers
import accelerate
import numpy as np
from termcolor import colored
import time
import json
import random
import math
import logging
from tqdm.auto import tqdm
from argparse import Namespace
from huggingface_hub.file_download import hf_hub_download
from accelerate import Accelerator
from transformers import (
    CONFIG_MAPPING,
    MODEL_MAPPING,
    AutoConfig,
    AutoModel,
    AutoModelForMaskedLM,
    AutoModelForSequenceClassification,
    AutoTokenizer,
    set_seed,
)

#gen_len = 450
block_size = 25
diffusion_steps = 100

for gen_len in [100, 200, 300]:
    print(f'------------------{gen_len}----------------------------------')
    # if doing unconstrained generation
    args = Namespace()
    args.model_name_or_path = "xhan77/ssdlm"
    args.max_seq_length = gen_len + 50
    args.one_hot_value = 5
    args.decoding_block_size = block_size
    args.decode_total_gen_len = gen_len # should be divisible by decode_depth
    args.decode_depth = math.ceil(gen_len / block_size)
    args.decode_log_interval = 100
    args.total_t = diffusion_steps
    args.projection_top_p = 0.95
    args.seed = 2022
    args.decode_ctr_lr = 0.0 # set to 0 for unconstrained generation, large value for controlled generation
    args.use_slow_tokenizer = False

    # a few helper functions
    def get_time_variables(t, total_t, device): # cosine schedule

        def ft(small_t, big_t, s=1e-4):
            return torch.cos((small_t / big_t + s) / (1 + s) * math.pi / 2) ** 2

        alpha_t_bar = ft(t, total_t) / ft(torch.zeros(t.shape, device=device), total_t)
        alpha_t_minus_bar = ft(t-1, total_t) / ft(torch.zeros(t.shape, device=device), total_t)
        beta_t = 1 - (alpha_t_bar / alpha_t_minus_bar)
        beta_t_til = (1 - alpha_t_minus_bar) / (1 - alpha_t_bar) * beta_t
        alpha_t = 1 - beta_t
        return alpha_t_bar, alpha_t_minus_bar, beta_t, beta_t_til, alpha_t


    def apply_controlling_drift(args, perturbed_inputs_diralpha):
        if args.decode_ctr_lr <= 0:
            args.ctr_loss = -1
            return perturbed_inputs_diralpha

        if args.ctr_model is None:
            args.ctr_model = AutoModelForSequenceClassification.from_pretrained(args.ctr_model_name).to(args.accelerator.device)
        optimizing_label_index = args.ctr_opt_label_idx

        for ctr_i in range(1):
            with torch.enable_grad():
                perturbed_inputs_diralpha_4ctr = perturbed_inputs_diralpha.clone()
                perturbed_inputs_diralpha_4ctr.requires_grad_()
                perturbed_inputs_simplex_4ctr = torch.nn.functional.softmax(perturbed_inputs_diralpha_4ctr, dim=-1)
                perturbed_inputs_embeds_4ctr = torch.nn.functional.linear(perturbed_inputs_simplex_4ctr, args.ctr_model.get_input_embeddings().weight.t())
                ctr_loss = -torch.nn.functional.log_softmax(args.ctr_model(inputs_embeds=perturbed_inputs_embeds_4ctr).logits, dim=-1)[:,optimizing_label_index].mean()
                args.ctr_loss = ctr_loss
                ctr_delta = -torch.autograd.grad(ctr_loss, perturbed_inputs_diralpha_4ctr)[0]
            perturbed_inputs_diralpha = perturbed_inputs_diralpha + args.decode_ctr_lr * ctr_delta # we use a fixed balancing factor in this work, which can be improved in the future
        
        return perturbed_inputs_diralpha


    def logits_sampling_projection(logits, top_p, one_hot_value):
        assert len(logits.size()) == 3
        very_low_value = -10000

        probs = torch.nn.functional.softmax(logits, dim=-1)
        sorted_probs, indices = torch.sort(probs, dim=-1, descending=True)
        cum_sum_probs = torch.cumsum(sorted_probs, dim=-1)
        nucleus = cum_sum_probs < top_p
        nucleus = torch.cat([nucleus.new_ones(nucleus.shape[:-1] + (1,)), nucleus[..., :-1]], dim=-1)
        valid_indices = nucleus.scatter(2, indices, nucleus)

        filtered_logits = logits.masked_fill(valid_indices == 0, -float('Inf'))
        m = torch.distributions.categorical.Categorical(logits=filtered_logits)
        selected = m.sample()
        return 2 * one_hot_value * torch.nn.functional.one_hot(selected, logits.size(2)) - one_hot_value


    def decode(args, batch_input_ids, dec_depth, total_t, model_embedding_lut, embedding_sum_layer, timestep_layer, model, tokenizer):
        batch_size = 1 # for the demo
        if args.decode_truncate_len > 0:
            diffusion_input_ids = batch_input_ids[:, args.context_size:-args.decode_truncate_len]
        else:
            diffusion_input_ids = batch_input_ids[:, args.context_size:]
        
        assert (args.max_seq_length - args.context_size - args.decode_truncate_len) % dec_depth == 0, "check whether the total generation length is divisible by the depth of decoding"
        unit_seq_len = int((args.max_seq_length - args.context_size - args.decode_truncate_len) / dec_depth)
        if args.context_size > 0:
            unit_context_input_ids = batch_input_ids[:, :args.context_size].clone()
        else:
            unit_context_input_ids = None
        history_decode_ids = None

        start_time = time.time()
        for i in range(dec_depth):
            unit_noise = args.one_hot_value * torch.normal(0, 1, size=(batch_size, unit_seq_len, args.vocab_size), device=args.accelerator.device)
            xt = unit_noise

            if unit_context_input_ids is not None:
                context_inputs_embeds = model_embedding_lut(unit_context_input_ids)
            else:
                context_inputs_embeds = None

            t_range = list(range(1, total_t+1))
            t_range.reverse()
            progress_bar = tqdm(range(len(t_range)), disable=not args.accelerator.is_local_main_process)
            
            for t in t_range:
                selected_t = torch.tensor([t], device=args.accelerator.device, dtype=torch.float).repeat(batch_size)
                alpha_t_bar, alpha_t_minus_bar, beta_t, beta_t_til, alpha_t = get_time_variables(selected_t, total_t, args.accelerator.device)
                zt = args.one_hot_value * torch.normal(0, 1, size=(batch_size, unit_seq_len, args.vocab_size), device=args.accelerator.device)
                
                perturbed_inputs_diralpha = xt
                perturbed_inputs_simplex = torch.nn.functional.softmax(perturbed_inputs_diralpha, dim=-1)

                perturbed_inputs_embeds = embedding_sum_layer(perturbed_inputs_simplex)
                t_progress = selected_t / total_t
                timestep_embeds = timestep_layer(t_progress.view(batch_size,1,1).repeat(1,unit_seq_len,1))

                diffusion_embeds = perturbed_inputs_embeds + timestep_embeds
                if context_inputs_embeds is not None:
                    diffusion_embeds = torch.cat((context_inputs_embeds, diffusion_embeds), dim=1)
                outputs = model(inputs_embeds=diffusion_embeds, output_hidden_states=False)
                equivalent_score = outputs.logits
                if unit_context_input_ids is not None:
                    equivalent_score = equivalent_score[:, unit_context_input_ids.size(1):].contiguous()

                # controlled generation if the balancing factor > 0
                equivalent_score = apply_controlling_drift(args, equivalent_score)

                projected_logits = logits_sampling_projection(equivalent_score, top_p=args.projection_top_p, one_hot_value=args.one_hot_value)

                xt = torch.sqrt(alpha_t_minus_bar).view(-1, 1, 1) * projected_logits
                xt = xt + torch.sqrt(1 - alpha_t_minus_bar).view(-1, 1, 1) * zt

                progress_bar.update(1)

                if t % args.decode_log_interval == 0 or t == 1:
                    simplex = torch.nn.functional.softmax(xt, dim=-1)
                    #logger.info(f"noise coef at t: {torch.sqrt(1 - alpha_t_bar).item()}")

                    if unit_context_input_ids is not None:
                        context_sequences = tokenizer.batch_decode(unit_context_input_ids.detach())
                        # logger.info(f"context: {context_sequences}")
                    
                    real_token_ids_list = torch.argmax(simplex, dim=-1).view(batch_size, unit_seq_len)
                    sampled_sequences = tokenizer.batch_decode(real_token_ids_list.clone().detach())
                    # logger.info(f"t={t} (argmax w_t-1): {colored(str(sampled_sequences), 'red')}")

                    simplex = equivalent_score
                    real_token_ids_list = torch.argmax(simplex, dim=-1).view(batch_size, unit_seq_len)
                    sampled_sequences = tokenizer.batch_decode(real_token_ids_list.clone().detach())
                    # logger.info(f"t={t} (argmax w_logits): {colored(str(sampled_sequences), 'blue')}")

                    #alt_i = 1 # look at the second best candidate; note that the whole sequence is not meaningful; each token can be considered as a substitution for the corresponding token in the argmax sequence
                    #alt_real_token_ids_list = torch.topk(simplex, alt_i+1, dim=-1).indices[:, :, alt_i].view(batch_size, unit_seq_len)
                    #alt_sampled_sequences = tokenizer.batch_decode(alt_real_token_ids_list.clone().detach().to('cpu'))
                    # logger.info(f"t={t} (argsecondmax w_logits): {alt_sampled_sequences}")

                    # logger.info(f"ctr loss: {args.ctr_loss}")
            
            unit_context_input_ids = torch.cat((unit_context_input_ids, real_token_ids_list), dim=1)
            if history_decode_ids is None:
                history_decode_ids = real_token_ids_list
            else:
                history_decode_ids = torch.cat((history_decode_ids, real_token_ids_list), dim=1)

        if args.context_size > 0:
            init_context_input_ids = batch_input_ids[:, :args.context_size].clone()
            context_sequences = tokenizer.batch_decode(init_context_input_ids.detach().to('cpu'))
        else:
            init_context_input_ids = None
            context_sequences = None
        gold_sequences = tokenizer.batch_decode(diffusion_input_ids.clone().detach().to('cpu'))
        sampled_sequences = tokenizer.batch_decode(history_decode_ids.clone().detach().to('cpu'))
        logger.info(f"Time taken: {time.time() - start_time}")
        logger.info(f"context: {context_sequences}")
        logger.info(f"gold: {colored(str(gold_sequences), 'yellow')}")
        logger.info(f"generation: {colored(str(sampled_sequences), 'red')}")

        return history_decode_ids, init_context_input_ids, diffusion_input_ids, sampled_sequences, context_sequences, gold_sequences

    accelerator = Accelerator()
    accelerate.utils.set_seed(args.seed, device_specific=True)
    logger = logging.getLogger(__name__)
    logger.setLevel(logging.INFO)
    stdout_handler = logging.StreamHandler(stream=sys.stdout)
    handlers = [stdout_handler]
    logging.basicConfig(
        level=logging.DEBUG, 
        format='[%(asctime)s] {%(filename)s:%(lineno)d} %(levelname)s - %(message)s',
        handlers=handlers
    )

    config = AutoConfig.from_pretrained(args.model_name_or_path, max_position_embeddings = 100000)
    tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path, use_fast=not args.use_slow_tokenizer)
    model = AutoModelForMaskedLM.from_pretrained(args.model_name_or_path, from_tf=False, config=config, ignore_mismatched_sizes=True)

    model.resize_token_embeddings(len(tokenizer))
    vocab_size = model.get_input_embeddings().weight.size(0)
    hidden_size = model.get_input_embeddings().weight.size(1)

    embedding_sum_layer = torch.nn.Linear(vocab_size, hidden_size, bias=False)
    _stdict = torch.load(os.path.join(hf_hub_download(args.model_name_or_path, "embed_sum_layer.pt")))
    _stdict = dict((_k[len("module."):], _stdict[_k]) if _k.startswith("module.") else (_k, _stdict[_k]) for _k in _stdict)
    embedding_sum_layer.load_state_dict(_stdict)

    timestep_layer = torch.nn.Linear(1, hidden_size, bias=True)
    _stdict = torch.load(os.path.join(hf_hub_download(args.model_name_or_path, "timestep_layer.pt")))
    _stdict = dict((_k[len("module."):], _stdict[_k]) if _k.startswith("module.") else (_k, _stdict[_k]) for _k in _stdict)
    timestep_layer.load_state_dict(_stdict)

    model, embedding_sum_layer, timestep_layer = accelerator.prepare(model, embedding_sum_layer, timestep_layer)

    # a bit more preparation before decoding
    model.eval()
    model_embedding_lut = accelerator.unwrap_model(model).get_input_embeddings()
    args.vocab_size = vocab_size
    args.accelerator = accelerator
    args.ctr_model = None
    args.orig_decode_truncate_len = args.max_seq_length - args.decode_total_gen_len

    # ENTER YOUR PROMPT HERE!
    prompt = "A man of innumerable personalities and powers vs. the most powerful artificial intelligence in this universe: Legion vs. Nimrod! With Nightcrawler in Orchis clutches, David Haller and his allies will have to confront the mastermind who"


    if prompt[0] != ' ': # prepend a space to the prompt if necessary
        prompt = f" {prompt}" # can use ' ' or '\n\n' as prefix to the prompt
    input_ids = torch.LongTensor(tokenizer.encode(prompt, add_special_tokens=False)).to(args.accelerator.device)
    args.context_size = len(input_ids)
    assert args.max_seq_length - args.decode_total_gen_len - args.context_size > 0, "check the length of the prompt"
    args.decode_truncate_len = args.orig_decode_truncate_len - args.context_size
    input_ids = input_ids.unsqueeze(0)

    # start sampling from SSD-LM
    with torch.inference_mode():
        for _ in range(5):
            history_decode_ids, context_input_ids, diffusion_input_ids, sampled_sequences, context_sequences, gold_sequences = \
                decode(args, input_ids, args.decode_depth, args.total_t, model_embedding_lut, embedding_sum_layer, timestep_layer, model, tokenizer)

scripts/benchmarking_scripts/ssdlm_inference_benchmark.py
allenai/tess-diffusion
import torch
import time
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

model = AutoModelForSeq2SeqLM.from_pretrained('facebook/bart-large', max_position_embeddings=100000, ignore_mismatched_sizes=True)
tokenizer = AutoTokenizer.from_pretrained('facebook/bart-large')

prompt = "A man of innumerable personalities and powers vs. the most powerful artificial intelligence in this universe: Legion vs. Nimrod! With Nightcrawler in Orchis clutches, David Haller and his allies will have to confront the mastermind who"

inputs = tokenizer(prompt, return_tensors="pt")

model = model.to('cuda')
inputs = inputs.to('cuda')

for lengths in [100, 200, 300]:
    print(f'-----------------{lengths}---------------------------')
    for _ in range(5):
        with torch.inference_mode():
            time_start = time.time()
            outputs = model.generate(**inputs, min_length=lengths, max_length=lengths)
            print(f"Time taken: {time.time() - time_start}")
            print(f"Output: {tokenizer.decode(outputs[0])}")
    print(f'-----------------{lengths} end---------------------------')
    
scripts/benchmarking_scripts/test_hf_autoregressive_generation.py
allenai/tess-diffusion
import logging
import os
import sys
import time

import torch.nn.functional as F
from transformers import (
    MODEL_FOR_MASKED_LM_MAPPING,
    AutoTokenizer,
    HfArgumentParser,
    set_seed,
)
from sdlm.arguments import ModelArguments, DiffusionArguments
from sdlm.models import RobertaDiffusionConfig, RobertaForDiffusionLM
from sdlm.schedulers import SimplexDDPMScheduler

import torch
from sdlm.pipelines.simplex_ddpm import SimplexDDPMPipeline



logger = logging.getLogger(__name__)
MODEL_CONFIG_CLASSES = list(MODEL_FOR_MASKED_LM_MAPPING.keys())
MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)

def main():
    parser = HfArgumentParser((ModelArguments, DiffusionArguments))
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        # If we pass only one argument to the script and it's the path to a json file,
        # let's parse it to get our arguments.
        model_args, diffusion_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        model_args, diffusion_args = parser.parse_args_into_dataclasses()


    # Set seed before initializing model.
    set_seed(42)
    config_kwargs = {
        "cache_dir": model_args.cache_dir,
        "revision": model_args.model_revision,
        "use_auth_token": True if model_args.use_auth_token else None,
    }
    config = RobertaDiffusionConfig.from_pretrained(
        model_args.model_name_or_path,
        self_condition=diffusion_args.self_condition,
        self_condition_zeros_after_softmax=diffusion_args.self_condition_zeros_after_softmax,
        deepmind_conditional=diffusion_args.deepmind_conditional,
        classifier_free_simplex_inputs=diffusion_args.classifier_free_simplex_inputs,
        classifier_free_uncond_input=diffusion_args.classifier_free_uncond_input,
        self_condition_mlp_projection=diffusion_args.self_condition_mlp_projection,
        self_condition_mix_before_weights=diffusion_args.self_condition_mix_before_weights,
        self_condition_mix_logits_before_weights=diffusion_args.self_condition_mix_logits_before_weights,
        empty_token_be_mask=diffusion_args.empty_token_be_mask,
        **config_kwargs,
    )
    tokenizer_kwargs = {
        "cache_dir": model_args.cache_dir,
        "use_fast": model_args.use_fast_tokenizer,
        "revision": model_args.model_revision,
        "use_auth_token": True if model_args.use_auth_token else None,
    }
    if model_args.tokenizer_name:
        tokenizer = AutoTokenizer.from_pretrained(model_args.tokenizer_name, **tokenizer_kwargs)
    elif model_args.model_name_or_path:
        tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path, **tokenizer_kwargs)
    else:
        raise ValueError(
            "You are instantiating a new tokenizer from scratch. This is not supported by this script."
            "You can do it from another script, save it, and load it from here, using --tokenizer_name."
        )

    if model_args.model_name_or_path:
        model = RobertaForDiffusionLM.from_pretrained(
            model_args.model_name_or_path,
            from_tf=bool(".ckpt" in model_args.model_name_or_path),
            config=config,
            cache_dir=model_args.cache_dir,
            revision=model_args.model_revision,
            use_auth_token=True if model_args.use_auth_token else None,
            ignore_mismatched_sizes=True
        )
    else:
        raise RuntimeError("You need to load a pretrained model")

    # We resize the xs only when necessary to avoid index errors. If you are creating a model from scratch
    # on a small vocab and want a smaller embedding size, remove this test.
    vocab_size = model.get_input_embeddings().weight.shape[0]
    if len(tokenizer) > vocab_size:
        model.resize_token_embeddings(len(tokenizer))

    # for some insane reason some of the model is not correctly loaded using from_pretrained...
    # state_dict = torch.load(os.path.join(model_args.model_name_or_path, "pytorch_model.bin"), map_location="cpu")
    # for some insane reason the word embeddings dont get loaded
    # model.roberta.embeddings.word_embeddings.weight = torch.nn.Parameter(state_dict['roberta.embeddings.word_embeddings.weight'])
    # model.tie_weights()
    # print([k for k in state_dict if torch.any(state_dict[k] != model.state_dict()[k])])

    def generate(
        inputs,
        simplex_value=5.0,
        top_p=.99,
        temperature=1.0,
        diffusion_steps=2500,
        beta_schedule="squaredcos_improved_ddpm",
        clip_sample=False,
        guidance_scale=1.0,
        generated_sequence_length=256,
        sleep_time=0.0
    ):
        generated_sequence_length = int(generated_sequence_length)
        tokenized_input = tokenizer([inputs], add_special_tokens=False, return_tensors='pt').input_ids
        tokenized_input_len = tokenized_input.shape[1]
        tokenized_input = torch.cat(
            [tokenized_input, torch.ones((1, generated_sequence_length))],
            axis=-1
        ).long()
        span_mask = torch.cat(
            [torch.zeros((1, tokenized_input_len)), torch.ones((1, generated_sequence_length))],
            axis=-1
        ).bool()
        inputs = {
            'input_ids': tokenized_input.cuda(),
            'span_mask': span_mask.cuda()
        }

        model.eval()
        
        pipeline = SimplexDDPMPipeline(
            model=model.cuda(),
            scheduler=SimplexDDPMScheduler(
                num_train_timesteps=diffusion_steps,
                beta_schedule=beta_schedule,
                simplex_value=simplex_value,
                clip_sample=clip_sample,
                device=torch.device("cuda", 0),
            ),
            simplex_value=simplex_value,
            top_p=top_p,
            sampling_type="top_p",  # currently only this is supported
            is_conditional_generation=True,
            tokenizer=tokenizer,
            classifier_free_uncond_input='empty_token',
            temperature=temperature,
            guidance_softmax_combination=True
        )
        pipeline_args = {
            "batch_size": 1,
            "seq_length": generated_sequence_length,
            "batch": inputs,
            "guidance_scale": guidance_scale,
        }

        time_start = time.time()
        print(tokenizer.decode(pipeline(**pipeline_args).simplex.argmax(-1)[0]))
        print(f"Time taken: {time.time() - time_start}")

    
    for lengths in [100, 200, 300]:
        print(f'-----------------{lengths}---------------------------')
        for _ in range(5):
            generate(
                "A man of innumerable personalities and powers vs. the most powerful artificial intelligence in this universe: Legion vs. Nimrod! With Nightcrawler in Orchis clutches, David Haller and his allies will have to confront the mastermind who",
                simplex_value=5.0,
                top_p=.99,
                temperature=1.0,
                diffusion_steps=100,
                beta_schedule="squaredcos_improved_ddpm",
                clip_sample=False,
                guidance_scale=1.0,
                generated_sequence_length=lengths,
                sleep_time=0.0,
            )
        print(f'-----------------{lengths} end---------------------------')
        

if __name__ == "__main__":
    main()

scripts/benchmarking_scripts/test_tess_inference_speed.py
allenai/tess-diffusion
from datasets import DatasetDict, Dataset
import datasets
from transformers import AutoTokenizer
import numpy as np
from datasets import load_dataset
import pdb

raw_datasets = load_dataset("cnn_dailymail", '3.0.0')

pdb.set_trace()
tokenizer = AutoTokenizer.from_pretrained("roberta-large", use_fast=True)
for split in ["validation", "train", "test"]:
    original_lengths = []
    simplification_lengths = []
    data = raw_datasets[split]
    for d in data:
        token_ids = tokenizer(d["article"], padding=False, truncation=False)["input_ids"]
        original_lengths.append(len(token_ids))
        token_ids = tokenizer(d["highlights"], padding=False, truncation=False)["input_ids"]
        simplification_lengths.append(len(token_ids))

    print("original ", np.mean(original_lengths), " ", np.std(original_lengths), " ", np.max(original_lengths))
    print(
        "summary ",
        np.mean(simplification_lengths),
        " ",
        np.std(simplification_lengths),
        " ",
        np.max(simplification_lengths),
    )

scripts/compute_cnn_dailymail_length.py
allenai/tess-diffusion
# Computes the maximum length in a dataset.
# Usage: python scripts/compute_max_data_length.py --dataset_name xsum --dataset_config "3.0.0" --model_name_or_path "roberta-large"

from sdlm.data.data_utils import load_data
from transformers import AutoTokenizer, HfArgumentParser
from sdlm.arguments import ModelArguments, DataTrainingArguments, Seq2SeqTrainingArguments, DiffusionArguments
import sys
import pdb
import logging
import os
import numpy as np

logger = logging.getLogger(__name__)

summarization_name_mapping = {
    "amazon_reviews_multi": ("review_body", "review_title"),
    "big_patent": ("description", "abstract"),
    "cnn_dailymail": ("article", "highlights"),
    "orange_sum": ("text", "summary"),
    "pn_summary": ("article", "summary"),
    "psc": ("extract_text", "summary_text"),
    "samsum": ("dialogue", "summary"),
    "thaisum": ("body", "summary"),
    "xglue": ("news_body", "news_title"),
    "xsum": ("document", "summary"),
    "wiki_summary": ("article", "highlights"),
    "multi_news": ("document", "summary"),
}


def compute_max_length(dataset, prefix):
    data_lengths = []
    target_lengths = []
    # Compute the length.
    for x in dataset:
        data_length = len(x["input_ids"])
        target_length = len(x["labels"])
        data_lengths.append(data_length)
        target_lengths.append(target_length)

    np.save(prefix + "_target_lengths.npy", target_lengths)
    np.save(prefix + "_data_lengths.npy", data_lengths)
    print("Target length ", np.mean(target_lengths), np.std(target_lengths), np.max(target_lengths), np.min(target_lengths))
    print("Data length ", np.mean(data_lengths), np.std(data_lengths), np.max(data_lengths), np.min(data_lengths))


def main():
    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, Seq2SeqTrainingArguments, DiffusionArguments))
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        model_args, data_args, training_args, _ = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        model_args, data_args, training_args, _ = parser.parse_args_into_dataclasses()

    tokenizer = AutoTokenizer.from_pretrained(
        model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        use_fast=model_args.use_fast_tokenizer,
        revision=model_args.model_revision,
        use_auth_token=True if model_args.use_auth_token else None,
    )
    max_target_length = data_args.max_target_length

    def preprocess_function(examples):
        # remove pairs where at least one record is None
        inputs, targets = [], []
        for i in range(len(examples[text_column])):
            if examples[text_column][i] and examples[summary_column][i]:
                inputs.append(examples[text_column][i])
                targets.append(examples[summary_column][i])
        # TODO: we need to process first the target, then cut the inputs to the max_length-target length to use the
        # maximum number of tokens.
        model_inputs = tokenizer(inputs, max_length=data_args.max_source_length, padding=False, truncation=True)
        # Tokenize targets with the `text_target` keyword argument
        labels = tokenizer(text_target=targets, max_length=max_target_length, padding=False, truncation=True)
        model_inputs["labels"] = labels["input_ids"]
        return model_inputs

    raw_datasets = load_data(data_args, model_args)
    # Preprocessing the datasets.
    # We need to tokenize inputs and targets.
    column_names = raw_datasets["train"].column_names

    # Get the column names for input/target.
    dataset_columns = summarization_name_mapping.get(data_args.dataset_name, None)
    assert dataset_columns is not None, "You need to provide the columns names."
    text_column, summary_column = dataset_columns[0], dataset_columns[1]

    train_dataset = raw_datasets["train"]
    validation_dataset = raw_datasets["validation"]
    test_dataset = raw_datasets["test"]
    with training_args.main_process_first(desc="train dataset map pre-processing"):
        train_dataset = train_dataset.map(
            preprocess_function,
            batched=True,
            num_proc=data_args.preprocessing_num_workers,
            remove_columns=column_names,
            load_from_cache_file=not data_args.overwrite_cache,
            desc="Running tokenizer on train dataset",
        )
        validation_dataset = validation_dataset.map(
            preprocess_function,
            batched=True,
            num_proc=data_args.preprocessing_num_workers,
            remove_columns=column_names,
            load_from_cache_file=not data_args.overwrite_cache,
            desc="Running tokenizer on train dataset",
        )
        test_dataset = test_dataset.map(
            preprocess_function,
            batched=True,
            num_proc=data_args.preprocessing_num_workers,
            remove_columns=column_names,
            load_from_cache_file=not data_args.overwrite_cache,
            desc="Running tokenizer on train dataset",
        )

    compute_max_length(train_dataset, prefix="train")
    compute_max_length(validation_dataset, prefix="validation")
    compute_max_length(test_dataset, prefix="test")
    # Max data_length across train/validation/test is 1024 and max target_length is 120.


if __name__ == "__main__":
    main()

scripts/compute_max_data_length.py
allenai/tess-diffusion
# Computes maximum numeber of tokens in the glue labels.
# from datasets import load_dataset
import pdb
from transformers import AutoTokenizer
import numpy as np

GLUE = ["mrpc", "qnli", "qqp", "rte", "sst2", "stsb", "wnli", "cola", "mnli"]
labels = {
    "cola": ["0", "1"],
    "mnli": ["0", "1", "2"],
    "mrpc": ["0", "1"],
    "rte": ["0", "1"],
    "sst2": ["0", "1"],
    "wnli": ["0", "1"],
    "stsb": [str(np.round(label, decimals=1)) for label in np.arange(0, 5.2, 0.2)],
    "qqp": ["0", "1"],
    "qnli": ["0", "1"],
}
total_max_label_length = -1
max_seq_length = 128
tokenizer_name = "roberta-large"
tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, use_fast=True)

for data in GLUE:
    is_regression = data == "stsb"
    # raw_datasets = load_dataset("glue", data)
    targets = labels[data]
    targets_tokens = tokenizer(text_target=targets, max_length=max_seq_length, padding=False, truncation=True)
    max_label_length = max([len(tokens) for tokens in targets_tokens["input_ids"]])
    total_max_label_length = max(total_max_label_length, max_label_length)

print(total_max_label_length)

scripts/compute_max_tokens_of_labels.py
allenai/tess-diffusion
# buckets the summaries and computes the rouge for each bucket.
import json
import pdb
from collections import Counter

import evaluate
import numpy as np
from transformers import AutoTokenizer

from sdlm.data.postprocessors import postprocess_text_for_metric


def extract_results(indices, lengths, data, x):
    extracted_indices = []
    for i, value in enumerate(indices):
        if value == x:
            extracted_indices.append(i)
    extracted_data = {}
    for k, v in data.items():
        extracted_data[k] = np.array(v)[extracted_indices]
    return extracted_data


def compute_metrics(results):
    keys = ["pred_texts_from_logits_masked"]
    metrics = {}
    for key in keys:
        decoded_preds = process_text(results[key]) if not skip_special_tokens else results[key]
        # Note that since decoded_labels is getting updated after post-process, we
        # need to compute it here for each key.
        decoded_labels = (
            process_text(results["gold_texts_masked"]) if not skip_special_tokens else results["gold_texts_masked"]
        )
        decoded_preds, decoded_labels = postprocess_text_for_metric("rouge", decoded_preds, decoded_labels)
        key_metrics = metric.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True)
        key_metrics = {k: round(v * 100, 4) for k, v in key_metrics.items()}
        key_metrics = {f"{key}_{k}": v for k, v in key_metrics.items()}
        metrics.update(key_metrics)
    return metrics


skip_special_tokens = True
generation_path = "${LOCAL_DIR}/outputs/paper_experiments/cnn_dailymail_results/ours_lr_3e-5_max_steps_120000_model_roberta-base/generated_test_top_p_None_temperature_1.0_seed_42_results.json"
data = json.load(open(generation_path))
tokenizer = AutoTokenizer.from_pretrained("roberta-base")
targets = data["gold_texts_masked"]

lengths = []
for target in targets:
    tokens = tokenizer(target)["input_ids"]
    lengths.append(len(tokens))

pred_lengths = []
for pred in data["pred_texts_from_logits_masked"]:
    tokens = tokenizer(pred)["input_ids"]
    pred_lengths.append(len(tokens))

pdb.set_trace()
bins = np.arange(0, 150, 25)
indices = np.digitize(lengths, bins)
counts = Counter(indices)
print(counts)
metric = evaluate.load("rouge")
unique_indices = np.unique(indices)
bucket_metrics = {}
for x in unique_indices:
    extracted_data = extract_results(indices, lengths, data, x)
    bucket_results = compute_metrics(extracted_data)
    bucket_metrics[str(x)] = bucket_results

bucket_metrics["target_lengths"] = lengths
bucket_metrics["pred_lengths"] = pred_lengths

with open("our_model_cnn_bucket_results.json", "w") as outfile:
    json.dump(bucket_metrics, outfile)

scripts/compute_rouge_for_buckets.py
allenai/tess-diffusion
# buckets the summaries and computes the rouge for each bucket.
import json
import pdb
from collections import Counter

import evaluate
import numpy as np
from transformers import AutoTokenizer

from sdlm.data.postprocessors import postprocess_text_for_metric


def extract_results(indices, lengths, data, x):
    extracted_indices = []
    for i, value in enumerate(indices):
        if value == x:
            extracted_indices.append(i)
    extracted_data = {}
    for k, v in data.items():
        extracted_data[k] = np.array(v)[extracted_indices]
    return extracted_data


def compute_metrics(results):
    keys = ["predictions"]
    metrics = {}
    for key in keys:
        decoded_preds = results[key]
        decoded_labels = results["gold_text"]
        decoded_preds, decoded_labels = postprocess_text_for_metric("rouge", decoded_preds, decoded_labels)
        key_metrics = metric.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True)
        key_metrics = {k: round(v * 100, 4) for k, v in key_metrics.items()}
        key_metrics = {f"{key}_{k}": v for k, v in key_metrics.items()}
        metrics.update(key_metrics)
    return metrics


skip_special_tokens = True
generation_path = "${LOCAL_DIR}/outputs/paper_experiments/cnn_dailymail_results/baseline_bart_base_lr_3e-5_max_steps_120000/checkpoint-120000/generations/generations_predict_top_p_None_temperature_1.0_seed_42_results.json"
data = json.load(open(generation_path))
tokenizer = AutoTokenizer.from_pretrained("facebook/bart-base")
targets = data["gold_text"]

lengths = []
for target in targets:
    tokens = tokenizer(target)["input_ids"]
    lengths.append(len(tokens))

pred_lengths = []
for pred in data["predictions"]:
    tokens = tokenizer(pred)["input_ids"]
    pdb.set_trace()
    pred_lengths.append(len(tokens))

bins = np.arange(0, 150, 25)
indices = np.digitize(lengths, bins)
counts = Counter(indices)
print(counts)
metric = evaluate.load("rouge")
unique_indices = np.unique(indices)
bucket_metrics = {}
for x in unique_indices:
    extracted_data = extract_results(indices, lengths, data, x)
    bucket_results = compute_metrics(extracted_data)
    bucket_metrics[str(x)] = bucket_results

bucket_metrics["target_lengths"] = lengths
bucket_metrics["pred_lengths"] = pred_lengths
with open("baseline_model_cnn_bucket_results.json", "w") as outfile:
    json.dump(bucket_metrics, outfile)

scripts/compute_rouge_for_buckets_for_bart.py
allenai/tess-diffusion
import datasets
import numpy as np
from datasets import Dataset, DatasetDict, load_dataset
from transformers import AutoTokenizer

dataset_folder = "${LOCAL_DIR}/simplex-diffusion/datasets/qqp/"
raw_datasets = DatasetDict()
for split in ["train", "valid", "test"]:
    dataset = load_dataset("json", data_files=f"{dataset_folder}/{split}.jsonl")["train"]
    data_split = split if split != "valid" else "dev"
    raw_datasets[data_split] = dataset

tokenizer = AutoTokenizer.from_pretrained("roberta-large", use_fast=True)
for split in ["train", "dev", "test"]:
    original_lengths = []
    simplification_lengths = []
    data = raw_datasets[split]
    for d in data:
        token_ids = tokenizer(d["src"], padding=False, truncation=False)["input_ids"]
        original_lengths.append(len(token_ids))
        token_ids = tokenizer(d["trg"], padding=False, truncation=False)["input_ids"]
        simplification_lengths.append(len(token_ids))

    print("src ", np.mean(original_lengths), " ", np.std(original_lengths), " ", np.max(original_lengths))
    print(
        "trg ",
        np.mean(simplification_lengths),
        " ",
        np.std(simplification_lengths),
        " ",
        np.max(simplification_lengths),
    )

scripts/compute_wiki_alignment_length.py
allenai/tess-diffusion
from datasets import DatasetDict, Dataset
import datasets
from transformers import AutoTokenizer
import numpy as np

raw_datasets = datasets.DatasetDict()
for split in ["train", "dev", "test"]:
    s1s = open(f"wikilarge/s1.{split}", "r").readlines()
    s2s = open(f"wikilarge/s2.{split}", "r").readlines()
    data = [{"original": s1, "simplification": s2} for s1, s2 in zip(s1s, s2s)]
    raw_datasets[split] = Dataset.from_list(data)

tokenizer = AutoTokenizer.from_pretrained("roberta-large", use_fast=True)
for split in ["train", "dev", "test"]:
    original_lengths = []
    simplification_lengths = []
    data = raw_datasets[split]
    for d in data:
        token_ids = tokenizer(d["original"], padding=False, truncation=False)["input_ids"]
        original_lengths.append(len(token_ids))
        token_ids = tokenizer(d["simplification"], padding=False, truncation=False)["input_ids"]
        simplification_lengths.append(len(token_ids))

    print("original ", np.mean(original_lengths), " ", np.std(original_lengths), " ", np.max(original_lengths))
    print(
        "simplification ",
        np.mean(simplification_lengths),
        " ",
        np.std(simplification_lengths),
        " ",
        np.max(simplification_lengths),
    )

scripts/compute_wikilarge_length.py
allenai/tess-diffusion
import json 

data = json.load(open('our_model_cnn_bucket_results.json', 'r'))


import matplotlib.pyplot as plt

# Uncomment the following line if you use ipython notebook
# %matplotlib inline

width = 0.35       # the width of the bars

men_means = {'G1': 20, 'G2': 35, 'G3': 30, 'G4': 35, 'G5': 27}
men_std = {'G1': 2, 'G2': 3, 'G3': 4, 'G4': 1, 'G5': 2}

rects1 = plt.bar(men_means.keys(), men_means.values(), -width, align='edge',
                yerr=men_std.values(), color='r', label='Men')

women_means = {'G1': 25, 'G2': 32, 'G3': 34, 'G4': 20, 'G5': 25}
women_std = {'G1': 3, 'G2': 5, 'G3': 2, 'G4': 3, 'G5': 3}

rects2 = plt.bar(women_means.keys(), women_means.values(), +width, align='edge',
                yerr=women_std.values(), color='y', label='Women')

# add some text for labels, title and axes ticks
plt.xlabel('Groups')
plt.ylabel('Scores')
plt.title('Scores by group and gender')
plt.legend()

def autolabel(rects):
    """
    Attach a text label above each bar displaying its height
    """
    for rect in rects:
        height = rect.get_height()
        plt.text(rect.get_x() + rect.get_width()/2., 1.05*height,
                '%d' % int(height),
                ha='center', va='bottom')

autolabel(rects1)
autolabel(rects2)
plt.savefig("figure.pdf")
plt.show()

scripts/plot_bucketize_rouge.py
allenai/tess-diffusion
import json
import os

import matplotlib.pyplot as plt
import numpy as np

top_ps = [0.9, 0.95, 0.99]  # [0.1, 0.5, 0.7, 0.9, 0.99]
temperatures = [1.0, 2.0, 4.0, 10.0]  # 0.1, 0.5

for k in ["pred_texts_from_logits_masked_muave"]:  # , "pred_texts_from_simplex_masked_muave"]:
    for top_p in top_ps:
        muaves = []
        actual_ts = []
        for temperature in temperatures:
            path = f"${LOCAL_DIR}/outputs/paper_experiments/tune_temperatures/ul2_self_condition_addition_context_25_top_p_{top_p}_temperature_{temperature}_truncation_56/checkpoint-15000/eval_results.json"
            # path=f"${LOCAL_DIR}/outputs/paper_experiments/tune_temperatures/ul2_self_condition_context_25_top_p_{top_p}_temperature_{temperature}_truncation_56/checkpoint-15000/eval_results.json"
            # path=f"${LOCAL_DIR}/outputs/paper_experiments/tune_temperatures/ul2_context_25_top_p_{top_p}_temperature_{temperature}_truncation_56/checkpoint-19000/eval_results.json"
            # path = f"${LOCAL_DIR}/outputs/paper_experiments/tune_temperature/ul2_self_condition_addition_context_25_generations_{top_p}_temperature_{temperature}/eval_results.json"
            results = json.load(open(path))
            if k in results:
                muave = results[k]
                muaves.append(np.round(muave * 100, 2))
                actual_ts.append(temperature)
                print(muaves)
                print(actual_ts)
        print("P ", top_p)
        print("Temperatures ", temperatures)
        print("MUAVEs       ", muaves)
        # Plot it.
        # plt.plot(temperature, muaves)
        # plt.title(f"top_p_{top_p}")
        # plt.xlabel(f"temperature")
        # plt.ylabel("muave")
        # plt.savefig(f"temperature_top_p_{top_p}.png")

    print("=" * 100)

scripts/plot_muave_based_on_top_p_temperatures.py
allenai/tess-diffusion
import json

import matplotlib.pyplot as plt

top_ps = [0.0, 0.1, 0.2, 0.5, 0.7, 0.9, 0.95, 0.99]
muaves = []
for top_p in top_ps:
    path = (
        f"${LOCAL_DIR}/outputs/paper_experiments/tune_top_p/ul2_length_50_context_25_generations_{top_p}/eval_results.json"
    )
    metrics = json.load(open(path))
    muave = metrics["pred_texts_from_simplex_masked_muave"]
    muaves.append(muave)

    print(muaves)
    print(top_ps)

# Plot
plt.plot(top_ps, muaves)
plt.xlabel("top-p")
plt.ylabel("muave")
plt.savefig("image.png")

scripts/plot_muave_based_to_top_p.py
allenai/tess-diffusion
import json
import os
import pdb

import numpy as np

# Reads glue values from a folder and computes the average.

task_to_metric = {
    "cola": ["matthews_correlation"],
    "mnli": ["accuracy"],
    "mrpc": ["accuracy", "f1", "combined_score"],
    "qnli": ["accuracy"],
    "qqp": ["accuracy", "f1", "combined_score"],
    "rte": ["accuracy"],
    "sst2": ["accuracy"],
    "stsb": ["pearson", "spearmanr", "combined_score"],
    "wnli": ["accuracy"],
}

glue_ordered = ["mnli", "qnli", "qqp", "rte", "sst2", "mrpc", "cola", "stsb", "wnli"]  # "mrpc"
small_datasets = ["cola", "mrpc", "rte", "stsb", "wnli"]


def read_values(paths, is_baseline=False, tasks=None):
    results = {}
    ALL_TASKS = tasks if tasks is not None else task_to_metric
    ORDERED_TASKS = glue_ordered if tasks is None else tasks
    for task in ALL_TASKS:
        results[task] = {}
        path = paths[task]
        data = json.load(open(path))
        for metric in task_to_metric[task]:
            if is_baseline:
                scale = 100
                results[task][metric] = np.round(data["test_" + metric] * scale, 2)
            else:
                scale = 1
                results[task][metric] = np.round(data["eval_pred_texts_from_logits_masked_" + metric] * scale, 2)
    print(results)

    # Computes average.
    all_scores = []
    for task in results:
        metric = task_to_metric[task][0] if len(task_to_metric[task]) == 1 else "combined_score"
        all_scores.append(results[task][metric])
    avg_results = np.round(np.mean(all_scores), 2)
    print("Average", avg_results)

    # Show results in the format of latex.
    table_row = []
    for task in ORDERED_TASKS:
        task_results = []
        for metric in task_to_metric[task]:
            if metric == "combined_score":
                continue
            task_results.append(results[task][metric])
        task_results = [str(t) for t in task_results]
        table_row.append("/".join(task_results))

    table_row.append(str(avg_results))
    print("&".join(table_row))


"""
output_dir = "${LOCAL_DIR}/outputs/simplex_new/glue_roberta_large_baseline_tuned/"
paths={task:os.path.join(output_dir, task, "test_results.json") for task in task_to_metric.keys()}
read_values(paths)
"""

"""
# Read glue values.
output_dir = "${LOCAL_DIR}/outputs/paper_experiments/ours_glue_self_condition_mean"
dirs = {
    "cola": "cola_steps_10_wd_0.01",
    "mnli": "mnli_steps_10_wd_0.01",
    "mrpc": "mrpc_steps_10_wd_0.01",
    "qnli": "qnli_steps_10_wd_0.01",
    "qqp": "qqp_steps_10_wd_0.01",
    "sst2": "sst2_steps_10_wd_0.01",
    "stsb": "stsb_steps_10_wd_0.01",
    "mnli": "mnli_steps_10_wd_0.01",
    "mrpc": "mrpc_steps_10_wd_0.01",
    "rte": "rte_steps_10_wd_0.01",
    "sst2": "sst2_steps_10_wd_0.01",
    "wnli": "wnli_steps_10_wd_0.01",
}
paths = {}
for task in dirs:
    paths[task] = os.path.join(output_dir, dirs[task], "test_results.json")
read_values(paths)
"""

# Read the GLUE baseline results.
paths = {}
for task in task_to_metric.keys():
    """
    if task in small_datasets:
       path_task = f"${LOCAL_DIR}/outputs/paper_experiments/glue_results/ours_self_condition_mean_mix_before_weights_{task}_steps_10_no_wd_max_16k_steps/"
    else:
       path_task = f"${LOCAL_DIR}/outputs/paper_experiments/glue_results/ours_self_condition_mean_mix_before_weights_{task}_steps_10_no_wd_max_steps_set"
    """
    # path_task = f"${LOCAL_DIR}/outputs/paper_experiments/glue_results/ours_self_condition_mean_mix_before_weights_{task}_steps_10_no_wd_max_16k_steps/"

    # **** this is selected *****
    # path_task = f"${LOCAL_DIR}/outputs/paper_experiments/glue_results/ours_self_condition_mean_mix_before_weights_{task}_steps_10_no_wd_max_steps_set"
    # **** baseline *****
    # path_task = f"${LOCAL_DIR}/outputs/paper_experiments/glue_results/baseline_{task}"

    # glue on all data.
    path_task = f"${LOCAL_DIR}/outputs/paper_experiments/glue_results/ours_self_condition_mean_mix_before_weights_{task}_steps_10_no_wd_max_steps_set_all_eval_data"

    # path_task=f"${LOCAL_DIR}/outputs/paper_experiments/glue_results/ours_self_condition_mean_mix_before_weights_{task}_steps_10_no_wd/"
    # path_task = f"${LOCAL_DIR}/outputs/paper_experiments/glue_results/ours_self_condition_mean_mix_before_weights_{task}_steps_10_wd_0.01/"
    # path_task =f"${LOCAL_DIR}/outputs/paper_experiments/glue_results/baseline_{task}"
    paths[task] = os.path.join(
        path_task, "eval_top_p_None_temperature_1.0_seed_42_guidance_scale_1.0_results.json"
    )  # f"eval_0_top_p_None_temperature_1.0_results.json") #"test_results.json")
print(glue_ordered)
read_values(paths, is_baseline=False)  # , tasks=small_datasets)

scripts/read_glue_values.py
allenai/tess-diffusion
import json
import pdb

import numpy as np

# path = "${LOCAL_DIR}/outputs/paper_experiments/gold_text_with_gpt_large_256"
path = "${LOCAL_DIR}/outputs/paper_experiments/gold_truncation__context_size_25"

name_to_short = {
    "generated_texts_masked_perplexity": "PPL",
    "generated_texts_masked_dist-1": "Dist-1",
    "generated_texts_masked_dist-2": "Dist-2",
    "generated_texts_masked_dist-3": "Dist-3",
    "generated_texts_masked_muave": "MAUVE",
    "generated_texts_masked_repetition": "Repetition",
    "generated_texts_masked_zipf_minus_a": "ZIPF-a",
    "generated_texts_masked_zipf_minus_r": "ZIPF-r",
    "generated_texts_masked_zipf_p": "ZIPF-p",
}

ordered_key = ["MAUVE", "PPL", "Dist-1", "Dist-2", "Dist-3", "ZIPF-a", "Repetition"]
metrics = json.load(open(f"{path}/metrics.json"))
metrics = {
    name_to_short[k]: np.round(100 * v, 2) if not name_to_short[k] in ["PPL", "ZIPF-a", "Repetition"] else np.round(v, 2)
    for k, v in metrics.items()
}

values = []
for k in ordered_key:
    values.append(metrics[k])
values = [str(v) for v in values]
values = "&".join(values)
print(ordered_key)
print(values)
print("_" * 100)

scripts/read_gold_metrics.py
allenai/tess-diffusion
import json
import os
import pdb

import numpy as np

# path = "${LOCAL_DIR}/outputs/paper_experiments/gpt2_evals_200/"
path = "${LOCAL_DIR}/outputs/paper_experiments/gpt2_evals_256/"

name_to_short = {
    "generated_texts_masked_perplexity": "PPL",
    "generated_texts_masked_dist-1": "Dist-1",
    "generated_texts_masked_dist-2": "Dist-2",
    "generated_texts_masked_dist-3": "Dist-3",
    "generated_texts_masked_muave": "MAUVE",
    "generated_texts_masked_repetition": "Repetition",
    "generated_texts_masked_zipf_minus_a": "ZIPF-a",
    "generated_texts_masked_zipf_minus_r": "ZIPF-r",
    "generated_texts_masked_zipf_p": "ZIPF-p",
}

ordered_key = ["MAUVE", "PPL", "Dist-1", "Dist-2", "Dist-3", "ZIPF-a", "Repetition"]
for name in ["gpt2_large_top_p", "gpt2_xl_top_p", "gpt2_medium_top_p"]:
    for top_p in [0.95, 0.99, 0.9]:
        print(f"{name}_{top_p}")
        json_filepath = f"{path}/{name}_{top_p}/metrics.json"
        if os.path.isfile(json_filepath):
            metrics = json.load(open(json_filepath))
        else:
            metrics = np.load(f"{path}/{name}_{top_p}/metrics.npy", allow_pickle=True).item()

        metrics = {
            name_to_short[k]: np.round(100 * v, 2)
            if not name_to_short[k] in ["PPL", "ZIPF-a", "Repetition"]
            else np.round(v, 2)
            for k, v in metrics.items()
        }

        values = []
        for k in ordered_key:
            values.append(metrics[k])
        values = [str(v) for v in values]
        values = "&".join(values)
        print(ordered_key)
        print(values)
    print("_" * 100)

scripts/read_gpt_metrics.py
allenai/tess-diffusion
import json
import pdb

import numpy as np


def print_values(path, prefix="logits"):
    name_to_short = {
        "perplexity": "PPL",
        "dist-1": "Dist-1",
        "dist-2": "Dist-2",
        "dist-3": "Dist-3",
        "muave": "MAUVE",
        "repetition": "Repetition",
        "zipf_minus_a": "ZIPF-a",
        "zipf_minus_r": "ZIPF-r",
        "zipf_p": "ZIPF-p",
    }

    if prefix == "logits":
        prefix = "pred_texts_from_logits_masked_"
    else:
        prefix = "pred_texts_from_simplex_masked_"
    ordered_key = ["MAUVE", "PPL", "Dist-1", "Dist-2", "Dist-3", "ZIPF-a", "Repetition"]
    metrics = json.load(open(f"{path}"))
    results = {}
    for k, v in metrics.items():
        for x in name_to_short.keys():
            results[name_to_short[x]] = (
                np.round(100 * metrics[prefix + x], 2)
                if not x in ["perplexity", "zipf_minus_a", "repetition"]  # ["PPL", "ZIPF-a", "Repetition"]
                else np.round(metrics[prefix + x], 2)
            )
    values = []
    for k in ordered_key:
        values.append(results[k])
    values = [str(v) for v in values]
    values = "&".join(values)
    print(ordered_key)
    print(values)
    print("_" * 100)


def read_ours():
    path = "${LOCAL_DIR}/outputs/paper_experiments/ours_eval/self_condition_logits_addition_guidance_5_0.99_56_1000/eval_results.json"
    print_values(path)


def read_ours_temperatures_top_p():
    for temperature in [1]:  # [1, 2, 4]:
        for top_p in [0.9, 0.95, 0.99]:  # 0.8,
            print(f"Top-p {top_p} temperature {temperature}")
            # path = f"${LOCAL_DIR}/outputs/paper_experiments/self_condition/tune_length_25_context_25_truncation_56/ul2_self_condition_mean_top_p_{top_p}_temperature_1/checkpoint-13000/eval_results.json"
            # path = f"${LOCAL_DIR}/outputs/paper_experiments/self_condition/tune_length_25_context_25_truncation_56/ul2_self_condition_with_max_top_p_{top_p}_temperature_1/checkpoint-13000/eval_results.json"
            path = f"${LOCAL_DIR}/outputs/paper_experiments/self_condition/tune_length_25_context_25_truncation_56/ul2_self_condition_with_addition_top_p_{top_p}_temperature_1/checkpoint-15000/eval_results.json"
            # ${LOCAL_DIR}/outputs/paper_experiments/tune_length_25_context_25_truncation_206/ul2_self_condition_with_addition_top_p_{top_p}_temperature_{temperature}/checkpoint-28000/eval_results.json"
            print_values(path, prefix="logits")


read_ours_temperatures_top_p()

scripts/read_our_results.py
allenai/tess-diffusion
import json

import numpy as np


def read_results(path, is_baseline, key=None):
    prefix = "test_pred_texts_from_logits_masked_" if not is_baseline else "predict_"
    # metrics = ["sari", "bleu", "bert_score",  "rouge1", "rouge2", "rougelsum", "dist-1", "dist-4"] # "bert_score_them",
    metrics = ["bleu", "bert_score_them", "rougeLsum", "dist-1", "dist-4"]
    # metrics = ["rouge1", "rouge2", "rougeLsum"]
    # metrics=  ["sari", "bleu", "bert_score_them", "rougeLsum"]
    results = json.load(open(path, "r"))
    values = []
    results_metrics = []
    for metric in metrics:
        if f"{prefix}{metric}" in results:
            values.append(np.round(results[f"{prefix}{metric}"], 2))
            results_metrics.append(metric)

    values = [str(v) for v in values]
    # print("&".join(results_metrics))
    if key is not None:
        values = [key] + values
        print("&".join(values), r"\\")
    else:
        print("&".join(values), r"\\")


is_baseline = False
# path="${LOCAL_DIR}/outputs/paper_experiments/simplification_results/ours_lr_2e-5_no_wd/test_top_p_None_temperature_1.0_results.json"
# path = "${LOCAL_DIR}/outputs/paper_experiments/wiki_alignment_tune_lr/ours_lr_3e-5_no_wd/test_top_p_None_temperature_1.0_results.json"
# path="${LOCAL_DIR}/outputs/paper_experiments/wiki_alignment_tune_lr/roberta_base_lr_3e-5_no_wd/predict_top_p_None_temperature_1.0_results.json"
# path="${LOCAL_DIR}/outputs/paper_experiments/qqp_tune_steps/ours_lr_3e-5_steps_90000/checkpoint-90000/test_top_p_None_temperature_1.0_results.json"

# path="${LOCAL_DIR}/outputs/paper_experiments/wiki_alignment_self_condition_ablations/lr_3e-5_no_wd_steps_60000_no_self_condition/test_results.json"
# path="${LOCAL_DIR}/outputs/paper_experiments/wiki_alignment_self_condition_ablations/lr_3e-5_no_wd_steps_60000_self_condition_logits/test_results.json"
# path="${LOCAL_DIR}/outputs/paper_experiments/wiki_alignment_self_condition_ablations/lr_3e-5_no_wd_steps_60000_self_condition_logits_mean_ours_mix_before_weights/test_results.json"
# path="${LOCAL_DIR}/outputs/paper_experiments/qqp_tune_steps/bart_base_lr_3e-5_no_wd_max_steps_90000/predict_top_p_None_temperature_1.0_results.json"

# ablation results.
# path="${LOCAL_DIR}/outputs/paper_experiments/wiki_alignment_self_condition_ablations/lr_3e-5_no_wd_steps_60000_no_self_condition/checkpoint-60000/test_top_p_None_temperature_1.0_results.json"
# path="${LOCAL_DIR}/outputs/paper_experiments/wiki_alignment_self_condition_ablations/lr_3e-5_no_wd_steps_60000_self_condition_logits/checkpoint-60000/test_top_p_None_temperature_1.0_results.json"
# path="${LOCAL_DIR}/outputs/paper_experiments/wiki_alignment_self_condition_ablations/lr_3e-5_no_wd_steps_60000_self_condition_logits_mean_ours_mix_before_weights/checkpoint-60000/test_top_p_None_temperature_1.0_results.json"
# path="${LOCAL_DIR}/outputs/paper_experiments/qg_tune_steps/ours_lr_3e-5_steps_120000/test_top_p_None_temperature_1.0_seed_42_results.json"
# path="${LOCAL_DIR}/outputs/paper_experiments/simplification_results/baseline_lr_2e-5_no_wd/predict_top_p_None_temperature_1.0_results.json"
# path="${LOCAL_DIR}/outputs/paper_experiments/cnn_dailymail_results/ours_lr_3e-5_max_steps_120000_model_roberta-base/test_top_p_None_temperature_1.0_seed_42_results.json"
# path="${LOCAL_DIR}/outputs/paper_experiments/cnn_dailymail_results/baseline_bart_base_lr_3e-5_max_steps_120000/predict_top_p_None_temperature_1.0_seed_42_results.json"

# guidance scores.
"""
paths = {}
for t in ["1.5", "10.0", "16.0", "20.0", "40.0", "100.0",  "500.0"]: #["2.0", "4.0", "8.0", "10.0", "16.0"]:
   # paths[t]=f"${LOCAL_DIR}/outputs/paper_experiments/wiki_alignment_with_guidance/ours_lr_3e-5_no_wd_guidance_2.0_inputs_empty_token_repeat/checkpoint-80000/examples_500/test_top_p_None_temperature_1.0_seed_42_guidance_scale_{t}_results.json"
   # paths[t]=f"${LOCAL_DIR}/outputs/paper_experiments/wiki_alignment_with_guidance/ours_lr_3e-5_no_wd_guidance_2.0_inputs_empty_token_repeat/checkpoint-80000/examples_500/softmax/test_top_p_None_temperature_1.0_seed_42_guidance_scale_{t}_results.json"
   
   # paths[t] = f"${LOCAL_DIR}/outputs/paper_experiments/wiki_alignment_with_guidance/ours_lr_3e-5_no_wd_guidance_2.0_mask_token/checkpoint-80000/examples_500/test_top_p_None_temperature_1.0_seed_42_guidance_scale_{t}_results.json"
   # paths[t] = f"${LOCAL_DIR}/outputs/paper_experiments/wiki_alignment_with_guidance/ours_lr_3e-5_no_wd_guidance_2.0_mask_token/checkpoint-80000/examples_500/softmax/test_top_p_None_temperature_1.0_seed_42_guidance_scale_{t}_results.json"

   #paths[t] = f"${LOCAL_DIR}/outputs/paper_experiments/wiki_alignment_with_guidance/ours_lr_3e-5_no_wd_guidance_2.0_mask_token_inputs/checkpoint-80000/examples_500/test_top_p_None_temperature_1.0_seed_42_guidance_scale_{t}_results.json"
   # paths[t] = f"${LOCAL_DIR}/outputs/paper_experiments/wiki_alignment_with_guidance/ours_lr_3e-5_no_wd_guidance_2.0_mask_token_inputs/checkpoint-80000/examples_500/softmax/test_top_p_None_temperature_1.0_seed_42_guidance_scale_{t}_results.json"
   
   # paths[t] = f"${LOCAL_DIR}/outputs/paper_experiments/wiki_alignment_with_guidance/ours_lr_3e-5_no_wd_guidance_2.0/checkpoint-80000/examples_500/test_top_p_None_temperature_1.0_seed_42_guidance_scale_{t}_results.json"
   # paths[t] = f"${LOCAL_DIR}/outputs/paper_experiments/wiki_alignment_with_guidance/ours_lr_3e-5_no_wd_guidance_2.0/checkpoint-80000/examples_500/softmax/test_top_p_None_temperature_1.0_seed_42_guidance_scale_{t}_results.json"
   # paths[t]=f"${LOCAL_DIR}/outputs/paper_experiments/wiki_alignment_tune_lr/ours_lr_3e-5_no_wd/checkpoint-80000/examples_500/test_top_p_None_temperature_1.0_seed_42_guidance_scale_1.0_results.json"
   paths[t] =  f"${LOCAL_DIR}/outputs/paper_experiments/wiki_alignment_with_guidance/ours_lr_3e-5_no_wd_guidance_2.0/checkpoint-80000/examples_500/softmax/test_top_p_None_temperature_1.0_seed_42_guidance_scale_{t}_results.json"

for k, path in paths.items():
   #print(k) 
   read_results(path, is_baseline, key=k)
"""
"""
# qqp ablation.
paths = {}
top_p="0.99"
paths.update({"no-self": f"${LOCAL_DIR}/outputs/paper_experiments/qqp_self_condition_ablations/ours_lr_3e-5_steps_60000_no_self/checkpoint-60000/test_top_p_{top_p}_temperature_1.0_seed_42_guidance_scale_1.0_results.json"})
paths.update({"logits": f"${LOCAL_DIR}/outputs/paper_experiments/qqp_self_condition_ablations/ours_lr_3e-5_steps_60000_self_logits/checkpoint-60000/test_top_p_{top_p}_temperature_1.0_seed_42_guidance_scale_1.0_results.json"})
paths.update({"ours": f"${LOCAL_DIR}/outputs/paper_experiments/qqp_self_condition_ablations/ours_lr_3e-5_steps_60000_self_condition_logits_mean_mix_before_weights/checkpoint-60000/test_top_p_{top_p}_temperature_1.0_seed_42_guidance_scale_1.0_results.json"})
        
for k, path in paths.items():
   print(k) 
   read_results(path, is_baseline)
"""

# qqp less iterations.
# path="${LOCAL_DIR}/outputs/paper_experiments/qqp_tune_steps/ours_lr_3e-5_steps_90000/checkpoint-90000/ablation_inference_steps/steps_100/test_top_p_None_temperature_1.0_seed_42_guidance_scale_1.0_results.json"
# read_results(path, False)


# newsela less iterations.
# path="${LOCAL_DIR}/outputs/paper_experiments/wiki_alignment_tune_lr/ours_lr_3e-5_no_wd/checkpoint-80000/inference_steps_ablation_all_data/step_100/test_top_p_None_temperature_1.0_seed_42_guidance_scale_1.0_results.json"
# read_results(path, False)


# qg dataset.
# path="${LOCAL_DIR}/outputs/paper_experiments/qg_tune_steps/ours_lr_3e-5_steps_120000/checkpoint-120000/inference_steps_ablations/steps_100/test_top_p_None_temperature_1.0_seed_42_guidance_scale_1.0_results.json"
# read_results(path, False)

# text summarization.
# path="${LOCAL_DIR}/outputs/paper_experiments/cnn_dailymail_results/ours_lr_3e-5_max_steps_120000_model_roberta-base/checkpoint-120000/inference_steps_ablation_on_all_data/step_100/test_top_p_None_temperature_1.0_seed_42_guidance_scale_1.0_results.json"
# read_results(path, False)


# path="${LOCAL_DIR}/outputs/paper_experiments/cnn_dailymail_results_trained_from_ul2_variable_length_256/ours_lr_3e-5_max_steps_120000_model_roberta-base/test_top_p_None_temperature_1.0_seed_42_guidance_scale_1.0_results.json"
# read_results(path, False)


# path="${LOCAL_DIR}/outputs/paper_experiments/wiki_alignment_tuned_from_ul2_variable_len_256_checkpoint/ours_lr_3e-5_no_wd/test_top_p_None_temperature_1.0_seed_42_guidance_scale_1.0_results.json"
# read_results(path, False)

# path="${LOCAL_DIR}/outputs/paper_experiments/qg_ul2_variable_len_256_checkpoint/ours_lr_3e-5_steps_120000/test_top_p_None_temperature_1.0_seed_42_guidance_scale_1.0_results.json"
# read_results(path, False)


# path="${LOCAL_DIR}/outputs/paper_experiments/qqp_tune_from_ul2_variable_len_256_checkpoint/ours_lr_3e-5_steps_90000/test_top_p_None_temperature_1.0_seed_42_guidance_scale_1.0_results.json"
# path="${LOCAL_DIR}/outputs/paper_experiments/wiki_alignment_tuned_from_ul2_variable_len_256_checkpoint/ours_lr_3e-5_no_wd_max_steps_60000/test_top_p_None_temperature_1.0_seed_42_guidance_scale_1.0_results.json"
# path="${LOCAL_DIR}/outputs/paper_experiments/qqp_tune_from_ul2_variable_len_256_checkpoint/ours_lr_3e-5_steps_70000/test_top_p_None_temperature_1.0_seed_42_guidance_scale_1.0_results.json"
path = "${LOCAL_DIR}/outputs/paper_experiments/qg_ul2_variable_len_256_checkpoint/ours_lr_3e-5_steps_100000/test_top_p_None_temperature_1.0_seed_42_guidance_scale_1.0_results.json"
read_results(path, False)

scripts/read_simplification.py
allenai/tess-diffusion
import json
import pdb

import numpy as np

# path = "${LOCAL_DIR}/ssd-lm/logging/ssd_dbs25/ctx25_trunc56_depth7_ctrlr0.0_step1000_topp0.9_ssd_gen_sampling_metrics.json"
path = "${LOCAL_DIR}/ssdlm-baseline/logging/ssd_dbs25/ctx25_trunc206_depth1_ctrlr0.0_step1000_topp0.99_ssd_gen_sampling_metrics.json"

name_to_short = {
    "perplexity": "PPL",
    "dist-1": "Dist-1",
    "dist-2": "Dist-2",
    "dist-3": "Dist-3",
    "muave": "MAUVE",
    "repetition": "Repetition",
    "zipf_minus_a": "ZIPF-a",
    "zipf_minus_r": "ZIPF-r",
    "zipf_p": "ZIPF-p",
}

ordered_key = ["MAUVE", "PPL", "Dist-1", "Dist-2", "Dist-3", "ZIPF-a", "Repetition"]
metrics = json.load(open(f"{path}"))
metrics = {
    name_to_short[k]: np.round(100 * v, 2) if not name_to_short[k] in ["PPL", "ZIPF-a", "Repetition"] else np.round(v, 2)
    for k, v in metrics.items()
}

values = []
for k in ordered_key:
    values.append(metrics[k])
values = [str(v) for v in values]
values = "&".join(values)
print(ordered_key)
print(values)
print("_" * 100)

scripts/read_ssdlm_metrics.py
allenai/tess-diffusion
import csv
import json
import os
import pdb


def map_predictions(predictions, dataset):
    label_to_word = {"0": "entailment", "1": "not_entailment"}
    label_to_word_threeway = {"0": "entailment", "1": "neutral", "2": "contradiction"}
    if dataset == "qnli":
        predictions = [label_to_word[prediction] for prediction in predictions]
    elif dataset == "rte":
        predictions = [label_to_word[prediction] for prediction in predictions]
    elif dataset == "mnli":
        predictions = [label_to_word_threeway[prediction] for prediction in predictions]
    return predictions


dataset_to_filename = {
    "cola": "CoLA",
    "mrpc": "MRPC",
    "sst2": "SST-2",
    "rte": "RTE",
    "qqp": "QQP",
    "qnli": "QNLI",
    "wnli": "WNLI",
    "stsb": "STS-B",
    "mnli": "MNLI-m",
}
OUTPUT_DIR = "glue_results"


def write_to_a_file(predictions, filename, dataset):
    predictions = map_predictions(predictions, dataset)
    label_file = f"{OUTPUT_DIR}/{filename}.tsv"
    with open(label_file, "w", encoding="utf8", newline="") as tsv_file:
        tsv_writer = csv.writer(tsv_file, delimiter="\t", lineterminator="\n")
        tsv_writer.writerow(["index", "prediction"])
        for index, label in enumerate(predictions):
            tsv_writer.writerow([index, label])


if not os.path.isdir(OUTPUT_DIR):
    os.mkdir(OUTPUT_DIR)

datasets = ["qnli", "rte", "stsb", "wnli", "qqp", "mrpc", "sst2", "mnli", "cola"]
for dataset in datasets:
    path = f"${LOCAL_DIR}/outputs/paper_experiments/glue_results/ours_self_condition_mean_mix_before_weights_{dataset}_steps_10_no_wd_max_steps_set/all_data_eval/generated_test_top_p_None_temperature_1.0_seed_42_guidance_scale_1.0_results.json"
    f = open(path)
    predictions = json.load(f)["pred_texts_from_logits_masked"]
    filename = dataset_to_filename[dataset]
    write_to_a_file(predictions, filename, dataset)
    pdb.set_trace()

scripts/write_glue.py
allenai/tess-diffusion
"""Arguments used in training/inference/data processing."""
from dataclasses import dataclass, field
from typing import Optional

from transformers import MODEL_MAPPING, SchedulerType
from transformers import TrainingArguments as HFTrainingArguments

MODEL_CONFIG_CLASSES = list(MODEL_MAPPING.keys())
MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)


@dataclass
class ModelArguments:
    """
    Arguments pertaining to which model/config/tokenizer we are going to fine-tune, or train from scratch.
    """

    model_name_or_path: Optional[str] = field(
        default=None,
        metadata={
            "help": ("The model checkpoint for weights initialization. Don't set if you want to train a model from scratch.")
        },
    )
    model_type: Optional[str] = field(
        default=None,
        metadata={"help": "If training from scratch, pass a model type from the list: " + ", ".join(MODEL_TYPES)},
    )
    config_overrides: Optional[str] = field(
        default=None,
        metadata={
            "help": (
                "Override some existing default config settings when a model is trained from scratch. Example: "
                "n_embd=10,resid_pdrop=0.2,scale_attn_weights=false,summary_type=cls_index"
            )
        },
    )
    tokenizer_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
    )
    cache_dir: Optional[str] = field(
        default=None,
        metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"},
    )
    use_fast_tokenizer: bool = field(
        default=True,
        metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
    )
    model_revision: str = field(
        default="main",
        metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."},
    )
    use_auth_token: bool = field(
        default=False,
        metadata={
            "help": (
                "Will use the token generated when running `huggingface-cli login` (necessary to use this script "
                "with private models)."
            )
        },
    )
    autoregressive_eval_model: str = field(
        default="EleutherAI/gpt-neo-1.3B",
        metadata={"help": "The autoregressive model used to measure the evaluation perplexity."},
    )
    resize_position_embeddings: Optional[bool] = field(
        default=None,
        metadata={
            "help": (
                "Whether to automatically resize the position embeddings if `max_source_length` exceeds "
                "the model's position embeddings."
            )
        },
    )
    resize_position_embeddings_alternatively: Optional[bool] = field(
        default=False,
        metadata={
            "help": "If set, resizes the position embedding alternatively, and copies from the original for the uncovered part."
        },
    )

    def __post_init__(self):
        if self.config_overrides is not None and (self.model_name_or_path is not None):
            raise ValueError("--config_overrides can't be used in combination with --model_name_or_path")


@dataclass
class TrainingArguments(HFTrainingArguments):
    lr_scheduler_type: SchedulerType = field(
        default="linear",
        metadata={
            "help": (
                "The scheduler type to use. It can be `linear`, `cosine`,"
                "`cosine_with_restarts`, `polynomial`, `constant`, and `constant_with_warmup`"
            )
        },
    )
    output_dir: Optional[str] = field(default=None, metadata={"help": "Where to store the final model."})
    checkpointing_steps: int = field(default=1000, metadata={"help": "Specifies the checkpoint step."})
    resume_from_checkpoint: Optional[str] = field(
        default=None, metadata={"help": "If the training should continue from a checkpoint folder."}
    )
    log_generated_texts: bool = field(default=True, metadata={"help": "If set, logs generated texts."})
    checkpoint_best_model: bool = field(
        default=False,
        metadata={
            "help": "If set, for `run_glue.py` it sets the metrics name" "to save the best model in each checkpoint step."
        },
    )
    eval_for_all_metrics: bool = field(default=False, metadata={"help": "If set, evaluates on all metrics in run_mlm.py"})
    load_states_in_eval_from_model_path: bool = field(
        default=True,
        metadata={
            "help": "In case of only using --do_eval without --do_train, use it to load the states before eval."
            "keep this to true, it causes otherwise an issue with huggingface when doing only --do_eval."
            "This parameter when running baselines does not have any impact and is not needed."
        },
    )
    without_compute_metrics: bool = field(
        default=False,
        metadata={
            "help": "If set, does not compute the metrics. we are observing MAUVE is very slow"
            "on multi-gpu setting and we do this to compute the metrics separately."
            "If using this option, you can call `compute_mlm_metrics.py` to compute them on 1 GPU later on."
        },
    )
    compute_eval_loss_with_simplex: bool = field(
        default=False, metadata={"help": "If set, computes the evaluation loss from the simplex values."}
    )
    ssdlm_optimizer: bool = field(default=False, metadata={"help": "If set, uses ssdlm optimizer."})


@dataclass
class Seq2SeqTrainingArguments(TrainingArguments):
    """
    Args:
        sortish_sampler (`bool`, *optional*, defaults to `False`):
            Whether to use a *sortish sampler* or not. Only possible if the underlying datasets are *Seq2SeqDataset*
            for now but will become generally available in the near future.
            It sorts the inputs according to lengths in order to minimize the padding size, with a bit of randomness
            for the training set.
        generation_max_length (`int`, *optional*):
            The `max_length` to use on each evaluation loop when `predict_with_generate=True`. Will default to the
            `max_length` value of the model configuration.
        generation_num_beams (`int`, *optional*):
            The `num_beams` to use on each evaluation loop when `predict_with_generate=True`. Will default to the
            `num_beams` value of the model configuration.
    """

    sortish_sampler: bool = field(default=False, metadata={"help": "Whether to use SortishSampler or not."})
    generation_max_length: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "The `max_length` to use on each evaluation loop when `predict_with_generate=True`. Will default "
                "to the `max_length` value of the model configuration."
            )
        },
    )
    generation_num_beams: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "The `num_beams` to use on each evaluation loop when `predict_with_generate=True`. Will default "
                "to the `num_beams` value of the model configuration."
            )
        },
    )


@dataclass
class DataTrainingArguments:
    """
    Arguments pertaining to what data we are going to input our model for training and eval.
    """
    split_glue: bool = field(default=True, metadata={"help": "If set to true split the glue dev/train to make the test set"
        "otherwises uses the original splits."})
    glue_split_seed: int = field(default=42, metadata={"help": "Seed to split the glue data."})
    tokenized_data_path: Optional[str] = field(default=None, metadata={"help": "If set, reads a tokenized train data."})
    dataset_name: Optional[str] = field(
        default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."}
    )
    dataset_config_name: Optional[str] = field(
        default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
    )
    dataset_folder: str = field(default=None, metadata={"help": "The dataset folder containing the dataset."})
    train_file: Optional[str] = field(default=None, metadata={"help": "The input training data file (a text file)."})
    validation_file: Optional[str] = field(
        default=None,
        metadata={"help": "An optional input evaluation data file to evaluate the perplexity on (a text file)."},
    )
    test_file: Optional[str] = field(default=None, metadata={"help": "A text file containing the test data."})
    overwrite_cache: bool = field(default=False, metadata={"help": "Overwrite the cached training and evaluation sets"})
    validation_split_ratio: Optional[float] = field(
        default=0.01,
        metadata={"help": "The ratio(< 1.0) of the train set used as validation set in case there's no validation split."},
    )
    max_seq_length: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "The maximum total input sequence length after tokenization. Sequences longer "
                "than this will be truncated."
            )
        },
    )
    preprocessing_num_workers: Optional[int] = field(
        default=None,
        metadata={"help": "The number of processes to use for the preprocessing."},
    )
    line_by_line: bool = field(
        default=False,
        metadata={"help": "Whether distinct lines of text in the dataset are to be handled as distinct sequences."},
    )
    pad_to_max_length: bool = field(
        default=False,
        metadata={
            "help": (
                "Whether to pad all samples to `max_seq_length`. "
                "If False, will pad the samples dynamically when batching to the maximum length in the batch."
            )
        },
    )
    max_train_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "For debugging purposes or quicker training, truncate the number of training examples to this "
                "value if set."
            )
        },
    )
    max_eval_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
                "value if set."
            )
        },
    )
    mask_ratio: float = field(default=0.15, metadata={"help": "Defines the ratio of mask tokens. A number between 0 and 1."})
    mean_mask_span_length: int = field(default=3, metadata={"help": "Defines the average mask length."})
    extra_padding_ratio: float = field(
        default=0.0,
        metadata={
            "help": (
                "Defines the ratio for the extra padding"
                "which are added only to the training data, in case of `span_infilling` uniformly."
            )
        },
    )
    conditional_generation: Optional[str] = field(
        default=None,
        metadata={
            "help": "It can be `span_infilling`, `prefix_lm`, `ul2`, or `ul2_with_unconditional`, `seq2seq`."
            "In case of `span_infilling`: It trains/evals on filling spans like T5. In `prefix_lm`: it trains/evals"
            "on completing the prefixes like GPT2. In `ul2`, it trains on a mixture of span_infilling, agressive"
            "span_infilling, or prefix_lm and evals on prefix_lm with masking half of the sequence. In case of"
            "`ul2_with_unconditional`: it uses ul2 with also including unconditional generation during training."
            "`seq2seq` is used for translation or summarization tasks. `ul2_variable`: is ul2 for the different"
            "T5 mask_ratio till half of the sequence."
        },
    )
    eval_context_size: Optional[int] = field(
        default=None,
        metadata={
            "help": "By default we consider the half of sequence as prompt when evaluating for `conditional_generation` of"
            "`ul2` and `prefix_lm`. If this parameter is set, it specifies the context size during the evaluation."
        },
    )
    # TODO: later fix masking length with truncation.
    truncation_length: Optional[int] = field(
        default=0,
        metadata={
            "help": "If set, we will truncate the tokens from the end for the given length."
            "Note we still compute masking length based on original data length!"
        },
    )
    skip_special_tokens: bool = field(
        default=True,
        metadata={
            "help": "If training line by line set this to False to generate end token and cut. Also, in case you want to consider generation till </s> and cut the rest."
        },
    )
    # Parameters used in seq2seq training for summarization.
    """
    test_file: Optional[str] = field(
        default=None,
        metadata={"help": "An optional input test data file to evaluate the metrics (rouge) on (a jsonlines or csv file)."},
    )
    """
    max_source_length: Optional[int] = field(
        default=1024,
        metadata={
            "help": (
                "The maximum total input sequence length after tokenization. Sequences longer "
                "than this will be truncated, sequences shorter will be padded."
            )
        },
    )
    max_target_length: Optional[int] = field(
        default=128,
        metadata={
            "help": (
                "The maximum total sequence length for target text after tokenization. Sequences longer "
                "than this will be truncated, sequences shorter will be padded."
            )
        },
    )
    val_max_target_length: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "The maximum total sequence length for validation target text after tokenization. Sequences longer "
                "than this will be truncated, sequences shorter will be padded. Will default to `max_target_length`."
                "This argument is also used to override the ``max_length`` param of ``model.generate``, which is used "
                "during ``evaluate`` and ``predict``."
            )
        },
    )
    max_predict_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "For debugging purposes or quicker training, truncate the number of prediction examples to this "
                "value if set."
            )
        },
    )
    num_beams: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "Number of beams to use for evaluation. This argument will be passed to ``model.generate``, "
                "which is used during ``evaluate`` and ``predict``."
            )
        },
    )
    # Translation arguments.
    source_lang: str = field(default=None, metadata={"help": "Source language id for translation."})
    target_lang: str = field(default=None, metadata={"help": "Target language id for translation."})
    add_t5_tags: bool = field(
        default=False, metadata={"help": "In case of GLUE, it adds tags to the sentences like `sentence1:` ... ."}
    )

    def __post_init__(self):
        if (
            not self.tokenized_data_path
            and self.dataset_name is None
            and (self.train_file is None and self.validation_file is None)
        ):
            raise ValueError(
                "Need either a task (only used for the `run_glue.py`), a dataset name or a training/validation file or a tokenized dataset path."
            )
        else:
            if self.train_file is not None:
                extension = self.train_file.split(".")[-1]
                if extension not in ["csv", "json", "txt"]:
                    raise ValueError("`train_file` should be a csv, a json or a txt file.")
            if self.validation_file is not None:
                extension = self.validation_file.split(".")[-1]
                if extension not in ["csv", "json", "txt"]:
                    raise ValueError("`validation_file` should be a csv, a json or a txt file.")

        if self.val_max_target_length is None:
            self.val_max_target_length = self.max_target_length

        if self.conditional_generation is not None:
            assert self.conditional_generation in [
                "span_infilling",
                "ul2",
                "ul2_with_unconditional",
                "prefix_lm",
                "seq2seq",
                "ul2_variable",
            ]


@dataclass
class DiffusionArguments:
    """Defines the diffusion related parameters."""

    simplex_value: float = field(
        default=5.0,
        metadata={
            "help": (
                "We map the token ids to a vector of vocabulary size, where for tokens not"
                "equal to the token id `-simplex_value` is selected, and `simplex_value` otherwise."
            )
        },
    )
    num_diffusion_steps: int = field(default=2500, metadata={"help": "Defines the number of diffusion steps."})
    num_inference_diffusion_steps: int = field(default=2500, metadata={"help": "Number of inference diffusion steps."})
    beta_schedule: str = field(
        default="squaredcos_improved_ddpm",
        metadata={
            "help": (
                "The beta schedule, a mapping from a beta range to a sequence of betas for stepping the model."
                "Choose from `linear`, `scaled_linear`, or `squaredcos_cap_v2`, `squaredcos_improved_ddpm`."
                "`squaredcos_improved_ddpm` model is proposed in eqn.17 in Improved ddpm"
                "(https://arxiv.org/pdf/2102.09672.pdf)"
            )
        },
    )
    sampling_type: str = field(default="top_p", metadata={"help": "Sampling type used during the logit projection."})
    top_p: Optional[float] = field(default=None, metadata={"help": "top_p value for nucleus (top_p) sampling."})
    clip_sample: bool = field(
        default=False,
        metadata={
            "help": "Whether to clip predicted sample between -1 and 1 for numerical stability in the noise scheduler."
        },
    )
    self_condition: Optional[str] = field(
        default=None,
        metadata={
            "help": (
                "If set, adds self-conditioning."
                "we consider the following options: `logits`: predicted logits, or `logits_with_projection`: to"
                "consider logits and apply the projection. After concatenating the inputs, we project inputs back"
                "with a projection layer to the half dimension. We also consider the cases of `logits_addition`"
                " and `logits_with_projection_addition` where we adds up the previous prediction to the logits,"
                "possibly with a projection operation. `logits_mean`: gets the average of logits and `logits_max`"
                "computes the maximum."
            )
        },
    )
    self_condition_mix_before_weights: bool = field(
        default=False, metadata={"help": "If set, mixes the softmax of simplexes and then apply the weights."}
    )
    self_condition_mix_logits_before_weights: bool = field(
        default=False, metadata={"help": "If set, mixes simplexes and then apply the weights."}
    )
    self_condition_mlp_projection: bool = field(default=False, metadata={"help": "If not set, uses a linear layer."})
    self_condition_zeros_after_softmax: bool = field(
        default=False,
        metadata={
            "help": "If set, makes the softmax of previous_logits,"
            "in case previous_logits are zero, zero. This avoid extra bias introduced with using Linear[softmax(previous_logits), logits]"
        },
    )
    deepmind_conditional: bool = field(
        default=False,
        metadata={
            "help": "This is the way conditional is explained in the DeepMind paper"
            "https://arxiv.org/abs/2211.15089, figure 3. In this setup, we mask the self-conditioned, noisy, and original emebeddings,"
            "then we concat mask to these, and project all of them, and then add timestep embeddings."
        },
    )
    guidance_scale: float = field(
        default=1.0, metadata={"help": "classifier-free guidance is applied if guidance_scale > 1.0."}
    )
    classifier_free_uncond_input: str = field(
        default="empty_token", metadata={"help": "This can be one of `empty_token` or `noisy_simplex`."}
    )
    empty_token_be_mask: bool = field(default=False, metadata={"help": "If set, makes the empty token a mask."})
    classifier_free_simplex_inputs: bool = field(
        default=False, metadata={"help": "If set to true, uses simplex representation for the unconditional input."}
    )
    temperature: float = field(default=1.0, metadata={"help": "Defines the softmax temperature before doing the sampling."})
    guidance_softmax_combination: bool = field(default=True, metadata={"help": "If set, first applies softmax, then combines logits."})
    generate_with_seed: bool = field(default=False, metadata={"help": "If set, generates with seed."})
sdlm/arguments.py
allenai/tess-diffusion
from dataclasses import dataclass
from typing import Optional, List, Dict, Any, Union
from transformers.tokenization_utils_base import PreTrainedTokenizerBase
from transformers.utils import PaddingStrategy
from sdlm.data.preprocessors import t5_random_spans_mask_batch, insert_extra_paddings, gpt_span_mask_batch
import numpy as np
import pdb
from random import choices
import torch
from enum import Enum
import random

class Objective(Enum):
    # Prefix language modeling like GPT style pretraining.
    prefix = 1
    # T5 objective with a range of 2 to 5 tokens as the span length, which masks about 15% of input tokens.
    t5 = 2
    # Aggressive denoising where approximately 50% of the input sequence is masked.
    aggressive_t5 = 3
    # Unconditional generation case.
    unconditional = 4


# TODO: automize this one.
# TODO: these are for sequence length of 100, adapt for 200.
OBJECTIVE_SETTINGS = {
    Objective.t5: [
        {"mask_ratio": 0.15, "mean_mask_span_length": 8},
        {"mask_ratio": 0.15, "mean_mask_span_length": 3},
    ],
    Objective.aggressive_t5: [
        {"mask_ratio": 0.5, "mean_mask_span_length": 8},
        {"mask_ratio": 0.5, "mean_mask_span_length": 3},
        {"mask_ratio": 0.5, "mean_mask_span_length": 48},
    ],
}


@dataclass
class SpanInfillingDataCollator:
    """
    Data collator that will dynamically pad the inputs received.
    Args:
        tokenizer ([`PreTrainedTokenizer`] or [`PreTrainedTokenizerFast`]):
            The tokenizer used for encoding the data.
        padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `True`):
            Select a strategy to pad the returned sequences (according to the model's padding side and padding index)
            among:
            - `True` or `'longest'` (default): Pad to the longest sequence in the batch (or no padding if only a single
              sequence is provided).
            - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum
              acceptable input length for the model if that argument is not provided.
            - `False` or `'do_not_pad'`: No padding (i.e., can output a batch with sequences of different lengths).
        max_length (`int`, *optional*):
            Maximum length of the returned list and optionally padding length (see above).
        pad_to_multiple_of (`int`, *optional*):
            If set will pad the sequence to a multiple of the provided value.
            This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability >=
            7.5 (Volta).
        return_tensors (`str`):
            The type of Tensor to return. Allowable values are "np", "pt" and "tf".
    """

    def __init__(
        self,
        mode,
        data_args,
        tokenizer: PreTrainedTokenizerBase,
        padding: Union[bool, str, PaddingStrategy] = True,
        max_length: Optional[int] = None,
        pad_to_multiple_of: Optional[int] = None,
        return_tensors: str = "pt",
        seed: int = 42,
        eval_context_size: int = None,
    ):
        self.tokenizer = tokenizer
        self.padding = padding
        self.max_length = max_length
        self.pad_to_multiple_of = pad_to_multiple_of
        self.return_tensors = return_tensors
        self.conditional_generation = data_args.conditional_generation
        self.extra_padding_ratio = data_args.extra_padding_ratio
        self.rng = np.random.default_rng(seed)
        self.eval_context_size = eval_context_size
        self.mode = mode
        if self.conditional_generation == "ul2_with_unconditional" and mode == "train":
            self.mask_generator = {}
            self.mask_generator[Objective.t5] = lambda batch, setting: t5_random_spans_mask_batch(
                batch, **setting, rng=self.rng
            )
            self.mask_generator[Objective.aggressive_t5] = lambda batch, setting: t5_random_spans_mask_batch(
                batch, **setting, rng=self.rng
            )
            self.mask_generator[Objective.prefix] = lambda batch: gpt_span_mask_batch(batch)
            self.mask_generator[Objective.unconditional] = lambda batch: None
        elif self.conditional_generation == "span_infilling":
            self.mask_generator = lambda batch: t5_random_spans_mask_batch(
                batch, data_args.mask_ratio, data_args.mean_mask_span_length, self.rng
            )
        elif self.conditional_generation == "prefix_lm":
            self.mask_generator = lambda batch: gpt_span_mask_batch(
                batch, use_half_length_as_prefix_size=(mode == "eval"), eval_context_size=eval_context_size
            )
        elif self.conditional_generation == "ul2" and mode == "train":
            self.mask_generator = {}
            self.mask_generator[Objective.t5] = lambda batch, setting: t5_random_spans_mask_batch(
                batch, **setting, rng=self.rng
            )
            self.mask_generator[Objective.aggressive_t5] = lambda batch, setting: t5_random_spans_mask_batch(
                batch, **setting, rng=self.rng
            )
            self.mask_generator[Objective.prefix] = lambda batch: gpt_span_mask_batch(batch)
        elif self.conditional_generation == "ul2_variable" and mode == "train":
            self.mask_generator = {}
            self.mask_generator[Objective.t5] = lambda batch, mask_ratio, mean_mask_span_length: t5_random_spans_mask_batch(
                batch, mask_ratio=mask_ratio, mean_mask_span_length=mean_mask_span_length, rng=self.rng
            )
            self.mask_generator[Objective.prefix] = lambda batch: gpt_span_mask_batch(batch)

    def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]:
        [f.pop("attention_mask") for f in features]

        if self.extra_padding_ratio:
            # Inserting random tokens uniformly, we do not modify start and end of
            # sequence tokens.
            for i in range(len(features)):
                features[i]["input_ids"] = insert_extra_paddings(
                    self.rng, features[i]["input_ids"], self.tokenizer.pad_token_id, self.extra_padding_ratio
                )

        masks = {}
        if self.conditional_generation in ["span_infilling", "prefix_lm"]:
            masks = {"span_mask": self.mask_generator(features)}
        elif self.conditional_generation == "ul2_with_unconditional" and self.mode == "train":
            objectives = [Objective.unconditional, Objective.t5, Objective.prefix, Objective.aggressive_t5]
            weights = [0.25, 0.25, 0.25, 0.25]
            objective = choices(objectives, weights)[0]
            if objective in [Objective.t5, Objective.aggressive_t5]:
                setting = choices(OBJECTIVE_SETTINGS[objective])[0]
                masks = {"span_mask": self.mask_generator[objective](features, setting)}
            else:
                masks = {"span_mask": self.mask_generator[objective](features)}
        elif self.conditional_generation == "ul2" and self.mode == "train":
            objectives = [Objective.t5, Objective.prefix, Objective.aggressive_t5]
            weights = [0.25, 0.25, 0.25]
            objective = choices(objectives, weights)[0]
            if objective in [Objective.t5, Objective.aggressive_t5]:
                setting = choices(OBJECTIVE_SETTINGS[objective])[0]
                masks = {"span_mask": self.mask_generator[objective](features, setting)}
            else:
                masks = {"span_mask": self.mask_generator[objective](features)}
        elif self.conditional_generation == "ul2_variable" and self.mode == "train":
            objectives = [Objective.t5, Objective.prefix]
            weights = [0.5, 0.5]
            objective = choices(objectives, weights)[0]
            if objective  == objective.t5:
                # Here we assume the length is the same for all data in a batch.
                length=len(features[0]["input_ids"])
                min_ratio = 1.0/length            
                mask_ratio = random.uniform(min_ratio, 0.5)
                mean_mask_span_length=int(random.uniform(1, mask_ratio*length))
                masks = {"span_mask": self.mask_generator[objective](features, mask_ratio, mean_mask_span_length)}
            else:
                masks = {"span_mask": self.mask_generator[objective](features)}
        elif self.mode == "eval" and self.conditional_generation in ["ul2", "ul2_with_unconditional", "ul2_variable"]:
            masks = {
                "span_mask": gpt_span_mask_batch(
                    features, use_half_length_as_prefix_size=True, eval_context_size=self.eval_context_size
                )
            }
        batch = self.tokenizer.pad(
            features,
            padding=self.padding,
            max_length=self.max_length,
            pad_to_multiple_of=self.pad_to_multiple_of,
            return_tensors=self.return_tensors,
        )
        if "attention_mask" in batch:
            del batch["attention_mask"]
        return {**batch, **masks}


@dataclass
class DataCollatorForSeq2Seq:
    """
    Data collator that will dynamically pad the inputs received, as well as the labels.
    Args:
        tokenizer ([`PreTrainedTokenizer`] or [`PreTrainedTokenizerFast`]):
            The tokenizer used for encoding the data.
        padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `True`):
            Select a strategy to pad the returned sequences (according to the model's padding side and padding index)
            among:
            - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single sequence
              is provided).
            - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum
              acceptable input length for the model if that argument is not provided.
            - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different
              lengths).
        max_length (`int`, *optional*):
            Maximum length of the returned list and optionally padding length (see above).
        pad_to_multiple_of (`int`, *optional*):
            If set will pad the sequence to a multiple of the provided value.
            This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability >=
            7.5 (Volta).
    """

    tokenizer: PreTrainedTokenizerBase
    padding: Union[bool, str, PaddingStrategy] = True
    max_length: Optional[int] = None
    pad_to_multiple_of: Optional[int] = None

    def __call__(self, features):
        input_ids = [feature["input_ids"] for feature in features]
        labels = [feature["labels"] for feature in features]
        input_target = [input + target for input, target in zip(input_ids, labels)]
        features = self.tokenizer.pad(
            {"input_ids": input_target},
            padding=self.padding,
            max_length=self.max_length,
            pad_to_multiple_of=self.pad_to_multiple_of,
            return_tensors="pt",
        )
        batch_length = features["input_ids"].shape[1]
        masks = [len(input) * [False] + (batch_length - len(input)) * [True] for input in input_ids]
        features["span_mask"] = torch.tensor(masks)
        return features

sdlm/data/data_collator.py
allenai/tess-diffusion
from itertools import chain
from datasets import load_dataset
import logging
import torch
import pdb
from datasets import DatasetDict

SMALL_GLUE_DATA = ["cola", "wnli", "rte", "mrpc", "stsb"]
LARGE_GLUE_DATA = ["qnli", "qqp", "sst2"]


logger = logging.getLogger(__name__)


def load_data(data_args, model_args):
    if data_args.dataset_name is not None:
        raw_datasets = load_dataset(
            data_args.dataset_name,
            data_args.dataset_config_name,
            cache_dir=model_args.cache_dir,
            use_auth_token=True if model_args.use_auth_token else None,
        )
    else:
        data_files = {}
        if data_args.train_file is not None:
            data_files["train"] = data_args.train_file
        if data_args.validation_file is not None:
            data_files["validation"] = data_args.validation_file
        extension = data_args.train_file.split(".")[-1]
        if extension == "txt":
            extension = "text"
        raw_datasets = load_dataset(
            extension,
            data_files=data_files,
            cache_dir=model_args.cache_dir,
            use_auth_token=True if model_args.use_auth_token else None,
        )
    return raw_datasets


def tokenize_data_new(data_args, tokenizer, raw_datasets, training_args):
    # Preprocessing the datasets.
    # First we tokenize all the texts.
    if training_args.do_train:
        column_names = raw_datasets["train"].column_names
    else:
        column_names = raw_datasets["validation"].column_names
    text_column_name = "text" if "text" in column_names else column_names[0]

    if data_args.max_seq_length is None:
        max_seq_length = tokenizer.model_max_length
        if max_seq_length > 1024:
            logger.warning(
                f"The tokenizer picked seems to have a very large `model_max_length` ({tokenizer.model_max_length}). "
                "Picking 1024 instead. You can change that default value by passing --max_seq_length xxx."
            )
            max_seq_length = 1024
    else:
        if data_args.max_seq_length > tokenizer.model_max_length:
            logger.warning(
                f"The max_seq_length passed ({data_args.max_seq_length}) is larger than the maximum length for the"
                f"model ({tokenizer.model_max_length}). Using max_seq_length={tokenizer.model_max_length}."
            )
        max_seq_length = min(data_args.max_seq_length, tokenizer.model_max_length)

    if data_args.line_by_line:
        # When using line_by_line, we just tokenize each nonempty line.
        padding = "max_length" if data_args.pad_to_max_length else False

        def tokenize_function(examples):
            # Remove empty lines
            examples[text_column_name] = [
                line for line in examples[text_column_name] if len(line) > 0 and not line.isspace()
            ]
            return tokenizer(
                examples[text_column_name],
                padding=padding,
                truncation=True,
                max_length=max_seq_length,
                # We use this option because DataCollatorForLanguageModeling (see below) is more efficient when it
                # receives the `special_tokens_mask`.
                return_special_tokens_mask=True,
            )

        with training_args.main_process_first(desc="dataset map tokenization"):
            tokenized_datasets = raw_datasets.map(
                tokenize_function,
                batched=True,
                num_proc=data_args.preprocessing_num_workers,
                remove_columns=[text_column_name],
                load_from_cache_file=not data_args.overwrite_cache,
                desc="Running tokenizer on dataset line_by_line",
            )
    else:
        # Otherwise, we tokenize every text, then concatenate them together before splitting them in smaller parts.
        # We use `return_special_tokens_mask=True` because DataCollatorForLanguageModeling (see below) is more
        # efficient when it receives the `special_tokens_mask`.
        def tokenize_function(examples):
            return tokenizer(examples[text_column_name], return_special_tokens_mask=True)

        with training_args.main_process_first(desc="dataset map tokenization"):
            tokenized_datasets = raw_datasets.map(
                tokenize_function,
                batched=True,
                num_proc=data_args.preprocessing_num_workers,
                remove_columns=column_names,
                load_from_cache_file=not data_args.overwrite_cache,
                desc="Running tokenizer on every text in dataset",
            )

        # Main data processing function that will concatenate all texts from our dataset and generate chunks of
        # max_seq_length.
        def group_texts(examples):
            # Concatenate all texts.
            concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()}
            total_length = len(concatenated_examples[list(examples.keys())[0]])
            # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
            # customize this part to your needs.
            if total_length >= max_seq_length:
                total_length = (total_length // max_seq_length) * max_seq_length
            # Split by chunks of max_len.
            result = {
                k: [t[i : i + max_seq_length] for i in range(0, total_length, max_seq_length)]
                for k, t in concatenated_examples.items()
            }
            return result

        # Note that with `batched=True`, this map processes 1,000 texts together, so group_texts throws away a
        # remainder for each of those groups of 1,000 texts. You can adjust that batch_size here but a higher value
        # might be slower to preprocess.
        #
        # To speed up this part, we use multiprocessing. See the documentation of the map method for more information:
        # https://huggingface.co/docs/datasets/package_reference/main_classes.html#datasets.Dataset.map

        with training_args.main_process_first(desc="grouping texts together"):
            tokenized_datasets = tokenized_datasets.map(
                group_texts,
                batched=True,
                num_proc=data_args.preprocessing_num_workers,
                load_from_cache_file=not data_args.overwrite_cache,
                desc=f"Grouping texts in chunks of {max_seq_length}",
            )
    return tokenized_datasets


# TODO: we need to remove this one and update process_data.py.
def tokenize_data(data_args, tokenizer, raw_datasets, accelerator):
    if data_args.max_seq_length is None:
        max_seq_length = tokenizer.model_max_length
        if max_seq_length > 1024:
            logger.warning(
                f"The tokenizer picked seems to have a very large `model_max_length` ({tokenizer.model_max_length}). "
                "Picking 1024 instead. You can change that default value by passing --max_seq_length xxx."
            )
            max_seq_length = 1024
    else:
        if data_args.max_seq_length > tokenizer.model_max_length:
            logger.warning(
                f"The max_seq_length passed ({data_args.max_seq_length}) is larger than the maximum length for the"
                f"model ({tokenizer.model_max_length}). Using max_seq_length={tokenizer.model_max_length}."
            )
        max_seq_length = min(data_args.max_seq_length, tokenizer.model_max_length)

    # Preprocessing the datasets.
    # First we tokenize all the texts.
    column_names = raw_datasets["train"].column_names
    text_column_name = "text" if "text" in column_names else column_names[0]

    if data_args.line_by_line:
        # When using line_by_line, we just tokenize each nonempty line.
        padding = "max_length" if data_args.pad_to_max_length else False

        def tokenize_function(examples):
            # Remove empty lines
            examples[text_column_name] = [
                line for line in examples[text_column_name] if len(line) > 0 and not line.isspace()
            ]
            return tokenizer(examples[text_column_name], padding=padding, truncation=True, max_length=max_seq_length)

        with accelerator.main_process_first():
            tokenized_datasets = raw_datasets.map(
                tokenize_function,
                batched=True,
                num_proc=data_args.preprocessing_num_workers,
                remove_columns=[text_column_name],
                load_from_cache_file=not data_args.overwrite_cache,
                desc="Running tokenizer on dataset line_by_line",
            )
    else:
        # Otherwise, we tokenize every text, then concatenate them together before splitting them in smaller parts.
        def tokenize_function(examples):
            return tokenizer(examples[text_column_name])

        with accelerator.main_process_first():
            tokenized_datasets = raw_datasets.map(
                tokenize_function,
                batched=True,
                num_proc=data_args.preprocessing_num_workers,
                remove_columns=column_names,
                load_from_cache_file=not data_args.overwrite_cache,
                desc="Running tokenizer on every text in dataset",
            )

        # Main data processing function that will concatenate all texts from our dataset and generate chunks of
        # max_seq_length.
        def group_texts(examples):
            # Concatenate all texts.
            concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()}
            total_length = len(concatenated_examples[list(examples.keys())[0]])
            # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
            # customize this part to your needs.
            if total_length >= max_seq_length:
                total_length = (total_length // max_seq_length) * max_seq_length
            # Split by chunks of max_len.
            result = {
                k: [t[i : i + max_seq_length] for i in range(0, total_length, max_seq_length)]
                for k, t in concatenated_examples.items()
            }
            return result

        # Note that with `batched=True`, this map processes 1,000 texts together, so group_texts throws away a
        # remainder for each of those groups of 1,000 texts. You can adjust that batch_size here but a higher value
        # might be slower to preprocess.
        #
        # To speed up this part, we use multiprocessing. See the documentation of the map method for more information:
        # https://huggingface.co/docs/datasets/package_reference/main_classes.html#datasets.Dataset.map

        with accelerator.main_process_first():
            tokenized_datasets = tokenized_datasets.map(
                group_texts,
                batched=True,
                num_proc=data_args.preprocessing_num_workers,
                load_from_cache_file=not data_args.overwrite_cache,
                desc=f"Grouping texts in chunks of {max_seq_length}",
            )
    return tokenized_datasets


def split_data_to_train_validation(data_args, data, seed):
    total_size = len(data["train"])
    validation_size = int(total_size * data_args.validation_split_ratio)
    train_size = total_size - validation_size

    # TODO(rabeeh): we need to do this for the other ones as well and think how to do it cleanly.
    if data_args.max_train_samples is not None:
        train_size = min(train_size, data_args.max_train_samples)
    if data_args.max_eval_samples is not None:
        validation_size = min(validation_size, data_args.max_eval_samples)

    remaining_size = total_size - train_size - validation_size
    train, validation, _ = torch.utils.data.random_split(
        data["train"], [train_size, validation_size, remaining_size], generator=torch.Generator().manual_seed(seed)
    )
    data["train"], data["validation"] = train, validation
    assert len(data["train"]) == train_size and len(data["validation"]) == validation_size
    return data


def split_glue(raw_datasets, dataset_name, seed):
    """Since glue test sets are not public, splits the data splits to form test sets.

    For large datasets (#samples > 10K), divides training set into 1K as validation and
    rest as train, using original validation as test. Otherwise, divides validation set
    to half (half for validation and half for test)."""
    if dataset_name == "mnli":
        raw_datasets = DatasetDict(
            {
                "test": raw_datasets["validation_matched"],
                "validation": raw_datasets["validation_mismatched"],
                "train": raw_datasets["train"],
            }
        )
    elif dataset_name in SMALL_GLUE_DATA:
        # Splits the validation set into half for validation and half for test.
        splits = raw_datasets["validation"].train_test_split(test_size=0.5, shuffle=True, seed=seed)
        raw_datasets = DatasetDict({"validation": splits["train"], "test": splits["test"], "train": raw_datasets["train"]})
    elif dataset_name in LARGE_GLUE_DATA:
        # Splits the training set into 1K as validation, rest as train.
        test_size = 1000 / len(raw_datasets["train"])
        splits = raw_datasets["train"].train_test_split(test_size=test_size, shuffle=True, seed=seed)
        raw_datasets = DatasetDict(
            {"train": splits["train"], "validation": splits["test"], "test": raw_datasets["validation"]}
        )
    else:
        raise NotImplementedError
    return raw_datasets

sdlm/data/data_utils.py
allenai/tess-diffusion
import pdb

from datasets import DatasetDict, load_from_disk

tokenized_data_path = "${LOCAL_DIR}/simplex-diffusion/processed_data/openwebtext_256_split"
output_dir = "${LOCAL_DIR}/simplex-diffusion/processed_data/openwebtext_256_split_gpt_eval"
seed = 42
tokenized_datasets = load_from_disk(tokenized_data_path)
validation_split_ratio = 0.1414827391058291
train_testvalid = tokenized_datasets["validation"].train_test_split(
    test_size=validation_split_ratio, shuffle=True, seed=seed
)
tokenized_datasets = DatasetDict({"train": tokenized_datasets["train"], "validation": train_testvalid["test"]})
tokenized_datasets.save_to_disk(output_dir)

sdlm/data/gpt_eval_split.py
allenai/tess-diffusion
import nltk  # Here to have a nice missing dependency error message early on
from transformers.utils import is_offline_mode
from filelock import FileLock

try:
    nltk.data.find("tokenizers/punkt")
except (LookupError, OSError):
    if is_offline_mode():
        raise LookupError("Offline mode: run this script without TRANSFORMERS_OFFLINE first to download nltk data files")
    with FileLock(".lock") as lock:
        nltk.download("punkt", quiet=True)

def string_to_float(string, default=-1.):
    """Converts string to float, using default when conversion not possible."""
    try:
        return float(string)
    except ValueError:
        return default


def string_to_int(string, default=-1):
    """Converts string to int, using default when conversion not possible."""
    try:
        return int(string)
    except ValueError:
        return default


def get_post_processor(task):
    """Returns post processor required to apply on the predictions/targets
    before computing metrics for each task."""
    if task == "stsb":
        return string_to_float
    elif task in ["qqp", "cola", "mrpc"]:
        return string_to_int
    else:
        return None


def postprocess_text_for_metric(metric, preds, labels=None, sources=None):
    if metric == "sari":
        assert sources is not None
        preds = [pred.strip() for pred in preds]
        labels = [label.strip() for label in labels]
        sources = [source.strip() for source in sources]
        return preds, labels, sources
    elif metric == "rouge":
        preds = [pred.strip() for pred in preds]
        labels = [label.strip() for label in labels]
        # rougeLSum expects newline after each sentence
        preds = ["\n".join(nltk.sent_tokenize(pred)) for pred in preds]
        labels = ["\n".join(nltk.sent_tokenize(label)) for label in labels]
        return preds, labels
    elif metric == "bleu":
        preds = [pred.strip() for pred in preds]
        labels = [[label.strip()] for label in labels]
        return preds, labels
    elif metric in ["bertscore", "bertscore_them"]:
        preds = [pred.strip() for pred in preds]
        labels = [label.strip() for label in labels]
        return preds, labels
    elif metric in ["dist"]:
        preds = [pred.strip() for pred in preds]
        return preds
    else:
        raise NotImplementedError
sdlm/data/postprocessors.py
allenai/tess-diffusion
"""Implements data preprocessings including the T5 preprocessing."""
import numpy as np
import itertools
import pdb
import torch


# TODO: here the max perhaps needs to be also the half-length.
def gpt_span_mask(length, pad_length, use_half_length_as_prefix_size, eval_context_size):
    """Given the length and pad_length for an input generates a prefix (GPT-style) mask."""
    # Start of the sequence is not masked, so we consider length-1.
    # TODO: we need an assert for length not be smaller than a value.
    if not use_half_length_as_prefix_size:
        # high should be higher than low, otherwise we set prefix_size=1.
        prefix_size = np.random.randint(low=1, high=int((length - 1) / 2)) if length >= 5 else 1
    else:
        # If eval_context_size is set, we consider it, otherwise we use half of the given length as
        # context. Note that since the start token is also masked, we deduct one from the given
        # context size.
        prefix_size = eval_context_size - 1 if eval_context_size is not None else int((length - 1) / 2)
    # The start token is not masked.
    return [False] + [False] * prefix_size + [True] * (length - prefix_size - 1) + [False] * pad_length


def gpt_span_mask_batch(batch, use_half_length_as_prefix_size=False, eval_context_size=None):
    lengths = [len(feature["input_ids"]) for feature in batch]
    max_length = max(lengths)
    masks = [
        gpt_span_mask(length, max_length - length, use_half_length_as_prefix_size, eval_context_size) for length in lengths
    ]
    return torch.tensor(masks)


def t5_random_spans_mask(length, mask_ratio, mean_mask_span_length=3.0, rng=None, pad_length=None):
    """Noise mask consisting of random spans of mask tokens.

    The number of mask tokens and the number of mask spans and non-mask spans
    are determined deterministically as follows:
      num_mask_tokens = round(length * mask_ratio)
      num_nonmask_spans = num_mask_spans = round(
         num_mask_tokens / mean_mask_span_length)
    Spans alternate between non-mask and mask, beginning with non-mask.
    Subject to the above restrictions, all masks are equally likely.
    Note that this function do not mask start/end of sequence.
    Args:
      length: an int32 scalar (length of the incoming token sequence)
      mask_ratio: a float - approximate ratio of output mask (between 0 and 1).
      mean_mask_span_length: Average mask length.
      rng = a np.random.default_rng() instance or None
    Returns:
      a boolean list of shape [length]
    adapted from https://github.com/google-research/text-to-text-transfer-transformer/blob/master/t5/data/preprocessors.py#L2704
    and https://github.com/allenai/contrastive_pretraining/blob/95fe35d3257402c7df362c3e0f746a40d9fba8f0/cpt/data.py#L288
    """
    # By default, we do not maks start and end of sequence.
    # TODO: we need to put assert for this!
    # NOTE: this only works if we use line_by_line which we do not. So I had to remove it.
    # length -= 2
    orig_length = length
    # Increase length to avoid degeneracy.
    length = max(length, 2)

    # Compute number of mask tokens and mask spans.
    num_mask_tokens = int(length * mask_ratio)
    # Avoid degeneracy by ensuring positive numbers of mask and nonmask tokens.
    num_mask_tokens = min(max(num_mask_tokens, 1), length - 1)
    num_mask_spans = int(num_mask_tokens / mean_mask_span_length)
    # Avoid degeneracy by ensuring positive number of mask spans.
    num_mask_spans = max(num_mask_spans, 1)
    num_nonmask_tokens = length - num_mask_tokens
    mask_span_lengths = _random_segmentation(num_mask_tokens, num_mask_spans, rng=rng)
    nonmask_span_lengths = _random_segmentation(num_nonmask_tokens, num_mask_spans, rng=rng)
    mask = list(
        itertools.chain.from_iterable(
            [[False] * nonmask_span_lengths[k] + [True] * mask_span_lengths[k] for k in range(num_mask_spans)]
        )
    )[:orig_length]
    # Start and end of the sequence mask are set to False. Again since this is not line_by_line, we
    # remove this.
    # mask = [False] + mask + [False]
    if pad_length is not None:
        mask += [False for _ in range(pad_length)]
    return mask


def t5_random_spans_mask_batch(batch, mask_ratio, mean_mask_span_length=3.0, rng=None):
    """Given not padded inputs, generates the T5 mask for each input."""
    lengths = [len(feature["input_ids"]) for feature in batch]
    max_length = max(lengths)
    masks = [t5_random_spans_mask(length, mask_ratio, mean_mask_span_length, rng, max_length - length) for length in lengths]
    return torch.tensor(masks)


def _random_segmentation(num_items, num_segments, rng=None):
    """Partition a sequence of items randomly into non-empty segments.
    Args:
      num_items: an integer scalar > 0
      num_segments: an integer scalar in [1, num_items]
      rng = a np.random.default_rng() instance or None
    Returns:
      a list with shape [num_segments] containing positive integers that add up to num_items.
    forked from: https://github.com/allenai/contrastive_pretraining/blob/95fe35d3257402c7df362c3e0f746a40d9fba8f0/cpt/data.py#L265
    """
    first_in_segment = np.arange(num_items - 1) < num_segments - 1
    rng = rng or np.random.default_rng()
    rng.shuffle(first_in_segment)
    # The first position always starts a segment.
    # first_in_segment is boolean array for every position after the first that signals whether this location is the start of a new segment.
    segment_id = np.cumsum(first_in_segment)
    segment_length = [0] * num_segments
    segment_length[0] = 1  # first the first missing first in segment
    for k in range(num_items - 1):
        segment_length[segment_id[k]] += 1
    return segment_length


def insert_extra_paddings(rng, token_ids, pad_token_id, padding_ratio):
    """Inserts padding tokens with the ratio of `padding_ratio` into the token_ids."""
    # TODO: we need to assert to have start/end of sequence tokens.
    # We do not add the padding in the start and end of sequence.
    length = len(token_ids) - 2
    num_padding_tokens = int(length * padding_ratio)
    if num_padding_tokens == 0:
        # In this case, the rate of padding tokens was not enough to add extra tokens.
        return token_ids
    length = length + num_padding_tokens
    # We do not modify the start token.
    all_ids = np.arange(1, length + 1)
    # This is without shuffling.
    # original_ids = np.arange(1, length+1)
    rng = rng or np.random.default_rng()
    rng.shuffle(all_ids)
    # padding tokens positions.
    padding_ids = np.array(all_ids)[:num_padding_tokens] + 1
    token_ids_extended = []
    current_id = 0
    for i in range(length + 2):
        if i not in padding_ids:
            token_ids_extended.append(pad_token_id)
        else:
            token_ids_extended.append(token_ids[current_id])
            current_id += 1
    return token_ids_extended
    """
    # Other tokens positions, we do not change the start and end of sequence tokens.
    other_tokens_ids = [0]+[x for x in original_ids if x not in padding_ids]+[length+1]
    # Considers the start and end of sequence tokens in the final length.
    token_ids_extended = np.full((length+2), pad_token_id, dtype=int)
    token_ids_extended[other_tokens_ids] = token_ids
    return token_ids_extended.tolist()
    """

sdlm/data/preprocessors.py
allenai/tess-diffusion
"""Tokenize the dataset and saves the output."""
import logging
import os
import sys

import datasets
from datasets import load_dataset
from sdlm.arguments import DataTrainingArguments, ModelArguments, TrainingArguments
from transformers import AutoTokenizer, HfArgumentParser, set_seed
from transformers.utils.versions import require_version
from accelerate import Accelerator
from sdlm.data.data_utils import tokenize_data

require_version("datasets>=1.8.0")
logger = logging.getLogger(__name__)


def main():
    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        # If we pass only one argument to the script and it's the path to a json file,
        # let's parse it to get our arguments.
        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    accelerator = Accelerator()

    # Setup logging
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        handlers=[logging.StreamHandler(sys.stdout)],
    )

    log_level = training_args.get_process_log_level()
    logger.setLevel(log_level)
    datasets.utils.logging.set_verbosity(log_level)

    # Log on each process the small summary:
    logger.warning(
        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
        + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
    )
    # Set the verbosity to info of the Transformers logger (on main process only):
    logger.info(f"Training/evaluation parameters {training_args}")

    # Set seed before initializing model.
    set_seed(training_args.seed)

    assert data_args.dataset_name is not None
    # Downloading and loading a dataset from the hub.
    raw_datasets = datasets.DatasetDict()
    raw_datasets["train"] = load_dataset(
        data_args.dataset_name,
        data_args.dataset_config_name,
        cache_dir=model_args.cache_dir,
        use_auth_token=True if model_args.use_auth_token else None,
        split="train",
    )

    tokenizer_kwargs = {
        "cache_dir": model_args.cache_dir,
        "use_fast": model_args.use_fast_tokenizer,
        "revision": model_args.model_revision,
        "use_auth_token": True if model_args.use_auth_token else None,
    }
    if model_args.model_name_or_path:
        tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path, **tokenizer_kwargs)
    else:
        raise ValueError(
            "You are instantiating a new tokenizer from scratch. This is not supported by this script."
            "You can do it from another script, save it, and load it from here, using --tokenizer_name."
        )

    tokenized_datasets = tokenize_data(data_args, tokenizer, raw_datasets, accelerator)

    with training_args.main_process_first():
        tokenized_datasets.save_to_disk(training_args.output_dir)
        logger.info(f"The processed data are written in {training_args.output_dir}")


if __name__ == "__main__":
    main()

sdlm/data/process_data.py
allenai/tess-diffusion
import pdb

from datasets import DatasetDict, load_from_disk

tokenized_data_path = "${LOCAL_DIR}/simplex-diffusion/processed_data/openwebtext_50"
output_dir = "${LOCAL_DIR}/simplex-diffusion/processed_data/openwebtext_50_split"
seed = 42
validation_split_ratio = 0.001

tokenized_datasets = load_from_disk(tokenized_data_path)
train_testvalid = tokenized_datasets["train"].train_test_split(test_size=validation_split_ratio, shuffle=True, seed=seed)
tokenized_datasets = DatasetDict({"train": train_testvalid["train"], "validation": train_testvalid["test"]})
tokenized_datasets.save_to_disk(output_dir)

sdlm/data/split_data.py
allenai/tess-diffusion
import pdb

import numpy as np
import torch
import torch.nn.functional as F

from sdlm.metrics.metrics import distinct_n_grams, zipf
from sdlm.metrics.perplexity import conditional_perplexity, perplexity
from sdlm.metrics.repetition import repetition
from sdlm.utils import convert_to_simplex, join_texts


def sample_logits(sampling_type, logits, top_p, temperature):
    # top-p (nucleus) sampling.
    if sampling_type == "top_p":
        logits = logits / temperature
        probs = F.softmax(logits, dim=-1)
        if top_p is not None:
            sorted_probs, sorted_indices = torch.sort(probs, dim=-1, descending=True)
            cumsum_probs = torch.cumsum(sorted_probs, dim=-1)

            # Remove tokens with cumulative probability above the threshold.
            sorted_indices_to_keep = cumsum_probs < top_p

            # Shift the indices to the right to keep also the first token below the threshold.
            sorted_indices_to_keep[..., 1:] = sorted_indices_to_keep[..., :-1].clone()
            sorted_indices_to_keep[..., 0] = 1

            indices_to_keep = sorted_indices_to_keep.scatter(dim=2, index=sorted_indices, src=sorted_indices_to_keep)
            filtered_logits = logits.masked_fill(indices_to_keep == 0, -float("Inf"))

            # sample from the filtered distribution.
            token_ids = torch.distributions.categorical.Categorical(logits=filtered_logits).sample()
        else:
            token_ids = torch.argmax(probs, dim=-1)
    else:
        assert NotImplementedError
    return token_ids


def remove_first_occurrence(string, char):
    # We do not strip as we need the spaces as well.
    if char in string:
        idx = string.index(char)
        string = string[idx + len(char) :]
    return string


def keep_till_first_occurrence(string, chars):
    """Given a list of characters, trim the text after the first occurance between them."""
    idxs = [string.index(char) for char in chars if char in string]
    if len(idxs):
        min_idx = np.min(idxs)
        string = string[:min_idx]
    return string


def process_text(texts):
    # TODO(rabeeh): for now we only cover roberta case.
    texts = [keep_till_first_occurrence(text, ["</s>"]) for text in texts]
    texts = [remove_first_occurrence(text, "<s>") for text in texts]
    return texts


def split_into_masked_and_unmasked(token_ids, span_mask, return_masked=None):
    """Given an span_mask, splits the given token_ids into masked and unmasked parts.

    If return_masked is set, only returns the masked parts, if this is set to False,
    only returns the unmasked parts, and If set to None, returns both parts.
    """

    def update_spans(span, masked, unmasked, mask):
        # TODO: this needs to be here for previous version of the codes.
        # span = torch.stack(span)
        masked.append(span) if mask else unmasked.append(span)

    masked = []
    unmasked = []
    prev_mask = span_mask[0]
    span = []
    for _, (token_id, mask) in enumerate(zip(token_ids, span_mask)):
        if mask == prev_mask:
            span.append(token_id)
        else:
            # Adds the previous span.
            update_spans(span, masked, unmasked, prev_mask)
            prev_mask = mask
            span = [token_id]
    # Adds the last span.
    update_spans(span, masked, unmasked, prev_mask)

    if return_masked is None:
        return masked, unmasked

    return masked if return_masked else unmasked


def concatenate_alternatively(longer, shorter, mark=""):
    """Given two lists of strings, concatenates them alternatively.

    We assume that the concatenated string should starts from elements in the longer
    list (which has one extra element). The shorter text can optionally be embraced with
    a `mark` text on both sides.
    """
    concatenated_str = ""
    for l, s in zip(longer, shorter):
        concatenated_str += l + " " + mark + s + mark + " "
    if len(longer) == len(shorter) + 1:
        return concatenated_str + longer[-1]
    elif len(longer) == len(shorter):
        return concatenated_str[:-1]
    else:
        raise ValueError


def aggregate_list(x):
    str = ""
    if len(x) == 0:
        return str
    for l in x:
        str += l + " "
    return str[:-1]


def logits_projection(logits, sampling_type, top_p, simplex_value, temperature):
    # TODO(rabeeh): huggingface has different sampling, like constrastive one.
    # also there are more variant in diffusion-lm.
    token_ids = sample_logits(sampling_type, logits, top_p, temperature)
    return convert_to_simplex(token_ids, simplex_value, vocab_size=logits.shape[2])


def filter_empty(texts):
    """Filters empty texts and return the remained texts and the their indices."""
    list_of_tuples = [(text, i) for i, text in enumerate(texts) if text != ""]
    if len(list_of_tuples) == 0:
        return [], []
    non_empty_texts, remained_inds = list(zip(*list_of_tuples))
    return list(non_empty_texts), list(remained_inds)


def predict_conditional_generated(span_masks, input_ids, tokenizer, predicted_token_ids, prefix_name, skip_special_tokens):
    masked = list(
        map(lambda x, y: split_into_masked_and_unmasked(x, y, return_masked=True), predicted_token_ids, span_masks)
    )
    unmasked = list(map(lambda x, y: split_into_masked_and_unmasked(x, y, return_masked=False), input_ids, span_masks))
    pred_masked_texts = [tokenizer.batch_decode(x, skip_special_tokens=skip_special_tokens) for x in masked]
    pred_unmasked_texts = [tokenizer.batch_decode(x, skip_special_tokens=skip_special_tokens) for x in unmasked]
    pred_texts = list(map(lambda x, y: concatenate_alternatively(x, y), pred_unmasked_texts, pred_masked_texts))
    pred_texts_marked = list(
        map(lambda x, y: concatenate_alternatively(x, y, mark="***"), pred_unmasked_texts, pred_masked_texts)
    )
    aggregated_masked_texts = list(map(lambda x: aggregate_list(x), pred_masked_texts))
    predicted_tokens = [np.array(item).tolist() for submasked in masked for item in submasked]
    return {
        # prefix_name: pred_texts,
        prefix_name + "_marked": pred_texts_marked,
        prefix_name + "_masked": aggregated_masked_texts,
        prefix_name + "_masked_tokens": predicted_tokens,
    }


def evaluate_generation(
    results,
    data_args,
    causal_model,
    causal_tokenizer,
    is_conditional_generation,
    prefix_lm_eval=False,
    skip_special_tokens=True,
    eval_for_all_metrics=False,
):
    metrics = {}
    # In case of prefix_lm since the generated text is unified, we can evaluate only the masked parts.
    if prefix_lm_eval:
        gold_text_key = "gold_texts_masked"
        # In case of gpt2, we only have the key of `generated_texts_masked`.
        keys = (
            ["generated_texts_masked"]
            if "generated_texts_masked" in results
            else ["pred_texts_from_simplex_masked", "pred_texts_from_logits_masked"]
        )
    else:
        keys = ["pred_texts_from_simplex", "pred_texts_from_logits"]
        gold_text_key = "gold_texts"

    if is_conditional_generation:
        gold_texts = results[gold_text_key]
        if not skip_special_tokens:
            gold_texts = process_text(gold_texts)
    if "prefixes" in results:
        prefixes = results["prefixes"]
    else:
        prefixes = None

    for key in keys:
        key_metrics = {}
        texts = results[key]
        if not skip_special_tokens:
            texts = process_text(texts)

        non_empty_texts, remained_indices = filter_empty(texts)
        if len(non_empty_texts) == 0:
            continue

        # Perplexity measured by a causal model.
        if prefixes is None:
            key_metrics.update(
                {"perplexity": perplexity(non_empty_texts, causal_model, causal_tokenizer)["mean_perplexity"]}
            )
        else:
            non_empty_prefixes = [prefix for i, prefix in enumerate(prefixes) if i in remained_indices]
            perplexity_results = conditional_perplexity(non_empty_texts, non_empty_prefixes, causal_model, causal_tokenizer)
            key_metrics.update(
                {
                    "perplexity": perplexity_results["mean_perplexity"],
                    "total_perplexity": perplexity_results["mean_perplexity_total"],
                }
            )

        # Dist-1,2,3 measurements.
        key_metrics.update(distinct_n_grams(texts))

        # Metrics requiring the gold text.
        if is_conditional_generation and eval_for_all_metrics:
            # Note that we need to pass both context and predicted texts to this metric.
            # remained_gold_texts = [text for i, text in enumerate(gold_texts) if i in remained_indices]
            # remained_prefixes = [text for i, text in enumerate(prefixes) if i in remained_indices]
            texts_with_context = join_texts(prefixes, texts)
            gold_with_context = join_texts(prefixes, gold_texts)
            length = data_args.max_seq_length - data_args.truncation_length

        if key + "_tokens" in results and eval_for_all_metrics:
            key_metrics.update(repetition(results[key + "_tokens"], causal_tokenizer))
            key_metrics.update(zipf(results[key + "_tokens"]))

        # Adds the metrics.
        key_metrics = {f"{key}_{k}": v for k, v in key_metrics.items()}
        metrics.update(key_metrics)

    return metrics

sdlm/inference/inference_utils.py
allenai/tess-diffusion
"""Implements the metrics for evaluation of the diffusion models."""
import math
import operator
import pdb
from collections import Counter

import numpy as np
import scipy
import sklearn
from nltk.util import ngrams
from scipy import stats

MAX_TEXT_LENGTH = 256


def distinct_n_grams(texts):
    """Computes the average distinct n-grams of the generated texts.
    Args:
        texts (list of str): representing the generated texts.
    """
    dist_1, dist_2, dist_3, dist_4 = [], [], [], []
    for text in texts:
        total_words = len(text.split())
        unigrams = set(ngrams(text.split(), 1))
        bigrams = set(ngrams(text.split(), 2))
        trigrams = set(ngrams(text.split(), 3))
        fourgrams = set(ngrams(text.split(), 4))
        if total_words == 0:
            dist_1.append(0)
            dist_2.append(0)
            dist_3.append(0)
            dist_4.append(0)
        else:
            dist_1.append(len(unigrams) / total_words)
            dist_2.append(len(bigrams) / total_words)
            dist_3.append(len(trigrams) / total_words)
            dist_4.append(len(fourgrams) / total_words)
    return {
        "dist-1": np.nanmean(dist_1),
        "dist-2": np.nanmean(dist_2),
        "dist-3": np.nanmean(dist_3),
        "dist-4": np.nanmean(dist_4),
    }


def zipf(tokenized_texts, N=5000):
    """Computes the Zipf coefficient.

    Args:
        tokenized_texts (List[List[int]]) tokenized texts.
    Adapted from https://github.com/ari-holtzman/degen/blob/master/metrics/zipf.py
    """
    cnt = Counter()
    for tokenized_text in tokenized_texts:
        cnt.update(tokenized_text)

    xs = np.arange(1, min(len(cnt), N) + 1)
    ys = np.array(sorted(cnt.values(), key=operator.neg)[:N])
    a, b, r, p, std = stats.linregress(np.log(xs), np.log(ys))
    # Note that zipf_minus_a is the reported number.
    return {"zipf_minus_a": -a, "zipf_minus_r": -r, "zipf_p": p}


def accuracy(predictions, targets) -> dict:
    """Computes the average accuracy."""
    return {"accuracy": 100 * ((np.array(predictions) == np.array(targets)).mean())}


def pearson_corrcoef(predictions, targets) -> dict:
    """Computes Pearson correlation coefficient."""
    pearson_corrcoef = 100 * scipy.stats.pearsonr(targets, predictions)[0]

    # Note that if all the predictions will be the same, spearman
    # correlation is nan, to gaurad against this, we check the output
    # and return 0 in this case.
    if math.isnan(pearson_corrcoef):
        pearson_corrcoef = 0
    return {"pearson": pearson_corrcoef}


def spearman_corrcoef(predictions, targets) -> dict:
    """Computes Spearman correlation coefficient."""
    spearman_corrcoef = 100 * scipy.stats.spearmanr(targets, predictions)[0]

    # Note that if all the predictions will be the same, spearman
    # correlation is nan, to gaurad against this, we check the output
    # and return 0 in this case.
    if math.isnan(spearman_corrcoef):
        spearman_corrcoef = 0
    return {"spearmanr": spearman_corrcoef}


def f1_score_with_invalid(predictions, targets) -> dict:
    """Computes F1 score,  with any prediction != 0 or 1 is counted as incorrect.
    Args:
      targets: list of targets, either 0 or 1
      predictions: list of predictions, any integer value
    Returns:
      F1 score, where any prediction != 0 or 1 is counted as wrong.
    """
    targets, predictions = np.asarray(targets), np.asarray(predictions)
    # Get indices of invalid predictions.
    invalid_idx_mask = np.logical_and(predictions != 0, predictions != 1)
    # For any prediction != 0 or 1, we set the prediction to the opposite of its corresponding target.
    predictions[invalid_idx_mask] = 1 - targets[invalid_idx_mask]
    return {"f1": 100 * sklearn.metrics.f1_score(targets, predictions)}


# TODO: maybe gaurd against invalid values https://stackoverflow.com/questions/56865344/how-do-i-calculate-the-matthews-correlation-coefficient-in-tensorflow
def matthews_corrcoef(predictions, targets) -> dict:
    """Computes the Matthews correlation coefficient."""
    return {"matthews_correlation": 100 * sklearn.metrics.matthews_corrcoef(targets, predictions)}


def get_glue_metrics(task):
    GLUE_TASKS_TO_METRICS = {
        "mrpc": [f1_score_with_invalid, accuracy],
        "cola": [matthews_corrcoef],
        "sst2": [accuracy],
        "stsb": [pearson_corrcoef, spearman_corrcoef],
        "qqp": [f1_score_with_invalid, accuracy],
        "mnli": [accuracy],
        "qnli": [accuracy],
        "rte": [accuracy],
        "wnli": [accuracy],
    }
    return GLUE_TASKS_TO_METRICS[task]

sdlm/metrics/metrics.py
allenai/tess-diffusion
"""Perplexity Metric. This file is adapted from: https://huggingface.co/spaces/evaluate-measurement/perplexity/blob/main/perplexity.py"""

import numpy as np
import torch
from torch.nn import CrossEntropyLoss
from evaluate import logging
import pdb


def perplexity(
    texts, model, tokenizer, batch_size: int = 16, add_start_token: bool = True, max_length=None, only_return_loss=False
):
    """Perplexity (PPL) can be used for evaluating to what extent a dataset is similar to the distribution of text that
    a given model was trained on. It is defined as the exponentiated average negative log-likelihood of a sequence,
    calculated with exponent base `e`.

    For more information, see https://huggingface.co/docs/transformers/perplexity

    Args:
        texts (list of str): List of text strings.
        model: model used for calculating Perplexity
            NOTE: Perplexity can only be calculated for causal language models.
                    This includes models such as gpt2, causal variations of bert,
                    causal versions of t5, and more (the full list can be found
                    in the AutoModelForCausalLM documentation here:
                    https://huggingface.co/docs/transformers/master/en/model_doc/auto#transformers.AutoModelForCausalLM )
        tokenizer: the corresponding tokenizer for the given model.
        batch_size (int): the batch size to run texts through the model. Defaults to 16.
        add_start_token (bool): whether to add the start token to the texts,
            so the perplexity can include the probability of the first word. Defaults to True.
    Returns:
        perplexity: dictionary containing the perplexity scores for the texts
            in the input list, as well as the mean perplexity. If one of the input texts is
            longer than the max input length of the model, then it is truncated to the
            max length for the perplexity computation.
    """
    device = model.device
    # if batch_size > 1 (which generally leads to padding being required), and
    # if there is not an already assigned pad_token, assign an existing
    # special token to also be the padding token
    if tokenizer.pad_token is None and batch_size > 1:
        existing_special_tokens = list(tokenizer.special_tokens_map_extended.values())
        # check that the model already has at least one special token defined
        assert (
            len(existing_special_tokens) > 0
        ), "If batch_size > 1, model must have at least one special token to use for padding. Please use a different model or set batch_size=1."
        # assign one of the special tokens to also be the pad token
        tokenizer.add_special_tokens({"pad_token": existing_special_tokens[0]})

    if add_start_token and max_length:
        # leave room for <BOS> token to be added:
        assert (
            tokenizer.bos_token is not None
        ), "Input model must already have a BOS token if using add_start_token=True. Please use a different model, or set add_start_token=False"
        max_tokenized_len = max_length - 1
    else:
        max_tokenized_len = max_length

    encodings = tokenizer(
        texts,
        add_special_tokens=False,
        padding=True,
        truncation=True if max_tokenized_len else False,
        max_length=max_tokenized_len,
        return_tensors="pt",
        return_attention_mask=True,
    ).to(device)
    encoded_texts = encodings["input_ids"]
    attn_masks = encodings["attention_mask"]

    # check that each input is long enough:
    if add_start_token:
        assert torch.all(torch.ge(attn_masks.sum(1), 1)), "Each input text must be at least one token long."
    else:
        assert torch.all(
            torch.ge(attn_masks.sum(1), 2)
        ), "When add_start_token=False, each input text must be at least two tokens long. Run with add_start_token=True if inputting strings of only one token, and remove all empty input strings."

    ppls = []
    loss_fct = CrossEntropyLoss(reduction="none")
    if only_return_loss:
        all_losses, all_lengths = [], []
    for start_index in logging.tqdm(range(0, len(encoded_texts), batch_size)):
        end_index = min(start_index + batch_size, len(encoded_texts))
        encoded_batch = encoded_texts[start_index:end_index]
        attn_mask = attn_masks[start_index:end_index]

        if add_start_token:
            bos_tokens_tensor = torch.tensor([[tokenizer.bos_token_id]] * encoded_batch.size(dim=0)).to(device)
            encoded_batch = torch.cat([bos_tokens_tensor, encoded_batch], dim=1)
            attn_mask = torch.cat([torch.ones(bos_tokens_tensor.size(), dtype=torch.int64).to(device), attn_mask], dim=1)

        labels = encoded_batch

        with torch.no_grad():
            out_logits = model(encoded_batch, attention_mask=attn_mask).logits

        shift_logits = out_logits[..., :-1, :].contiguous()
        shift_labels = labels[..., 1:].contiguous()
        shift_attention_mask_batch = attn_mask[..., 1:].contiguous()

        loss = (loss_fct(shift_logits.transpose(1, 2), shift_labels) * shift_attention_mask_batch).sum(1)
        lengths = shift_attention_mask_batch.sum(1)
        if only_return_loss:
            all_losses.append(loss)
            all_lengths.append(lengths)
        else:
            perplexity_batch = torch.exp(loss / lengths)
            ppls += perplexity_batch.tolist()

    if only_return_loss:
        return all_losses, all_lengths
    else:
        return {"perplexities": ppls, "mean_perplexity": np.mean(ppls)}


def conditional_perplexity(
    texts, prefixes, model, tokenizer, batch_size: int = 16, add_start_token: bool = True, max_length=None
):
    """Computes the conditional perplexity for the case of prefix language modeling."""
    full_texts = [f"{prefix}{text}" for prefix,text in zip(prefixes, texts)]
    loss, lengths = perplexity(full_texts, model, tokenizer, batch_size, add_start_token, max_length, only_return_loss=True)
    prefix_loss, prefix_lengths = perplexity(
        prefixes, model, tokenizer, batch_size, add_start_token, max_length, only_return_loss=True
    )
    # Computing the perplexity over the whole examples.
    ppls = []
    total_nlls = 0
    total_tokens = 0
    for i in range(len(loss)):
        perplexity_batch = torch.exp((loss[i] - prefix_loss[i]) / (lengths[i] - prefix_lengths[i]))
        ppls.extend(perplexity_batch.tolist())
        total_nlls += torch.sum(loss[i] - prefix_loss[i]).item()
        total_tokens += torch.sum(lengths[i] - prefix_lengths[i]).item()
    return {"perplexities": ppls, "mean_perplexity": np.nanmean(ppls), "mean_perplexity_total": np.exp(total_nlls/total_tokens)}

sdlm/metrics/perplexity.py
allenai/tess-diffusion
"""Computes the repetition metric. Adapted from: https://raw.githubusercontent.com/ari-holtzman/degen/master/metrics/repetition.py"""
import pdb


def repetition(tokenized_texts, tokenizer):
    """
    Args:
        tokenized_texts: (List[List[int]]) generated input tokenized texts.

    Computes the repetition metric https://arxiv.org/pdf/1904.09751.pdf showing how each
    example is repeating itself, specifically the phrase the generation is repeating
    and how many times it is repeated.
    """
    SEP = tokenizer.encode(tokenizer.bos_token)[0]
    repetition_stats = []
    max_n = 90
    num_examples = len(tokenized_texts)
    n_repeated_examples = 0
    for tokenized_text in tokenized_texts:
        if tokenized_text[-1] == SEP:
            tokenized_text.pop(-1)
        rev_gen = list(reversed(tokenized_text))
        last_n_repeats = [0] * max_n
        for n in range(1, max_n + 1):
            n_repeat = 1
            while (
                len(rev_gen[n * n_repeat : n * (n_repeat + 1)]) == n
                and rev_gen[n * n_repeat : n * (n_repeat + 1)] == rev_gen[:n]
            ):
                n_repeat += 1
            last_n_repeats[n - 1] = n_repeat
        max_repeated_n = max(range(max_n), key=lambda x: last_n_repeats[x])
        if last_n_repeats[max_repeated_n] > 1 and (max_repeated_n + 1 >= 3 or last_n_repeats[max_repeated_n] > 50):
            repetition_stats.append(
                {
                    "repeated_phrase": list(reversed(rev_gen[: max_repeated_n + 1])),
                    "repeated_times": last_n_repeats[max_repeated_n],
                    "repeated_phrase_length": max_repeated_n + 1,
                }
            )
            n_repeated_examples += 1
        else:
            repetition_stats.append({})

    return {"repetition": n_repeated_examples * 1.0 / num_examples}  # , "repetition_stats": repetition_stats}

sdlm/metrics/repetition.py
allenai/tess-diffusion
from .configuration_roberta import RobertaDiffusionConfig
from .modeling_roberta import RobertaForDiffusionLM

from .configuration_xlm_roberta import XLMRobertaDiffusionConfig
from .modeling_xlm_roberta import XLMRobertaForDiffusionLM

sdlm/models/__init__.py
allenai/tess-diffusion
"""Adapted Roberta configuration for diffusion models."""

from transformers.models.roberta.configuration_roberta import RobertaConfig
from typing import Optional


class RobertaDiffusionConfig(RobertaConfig):
    def __init__(
        self,
        self_condition: Optional[str] = None,
        self_condition_zeros_after_softmax: bool = False,
        deepmind_conditional: bool = False,
        classifier_free_simplex_inputs: bool = False,
        classifier_free_uncond_input: str = "empty_token",
        self_condition_mlp_projection=False,
        self_condition_mix_before_weights=False,
        self_condition_mix_logits_before_weights=False,
        empty_token_be_mask=False,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.self_condition = self_condition
        self.self_condition_zeros_after_softmax = self_condition_zeros_after_softmax
        self.deepmind_conditional = deepmind_conditional
        self.classifier_free_simplex_inputs = classifier_free_simplex_inputs
        self.classifier_free_uncond_input = classifier_free_uncond_input
        self.self_condition_mlp_projection = self_condition_mlp_projection
        self.self_condition_mix_before_weights = self_condition_mix_before_weights
        self.self_condition_mix_logits_before_weights = self_condition_mix_logits_before_weights
        self.empty_token_be_mask=empty_token_be_mask
sdlm/models/configuration_roberta.py
allenai/tess-diffusion
"""Adapted XLM Roberta configuration for diffusion models."""

from transformers.models.xlm_roberta.configuration_xlm_roberta import XLMRobertaConfig
from typing import Optional


class XLMRobertaDiffusionConfig(XLMRobertaConfig):
    def __init__(
        self,
        self_condition: Optional[str] = None,
        self_condition_zeros_after_softmax: bool = False,
        deepmind_conditional: bool = False,
        classifier_free_simplex_inputs: bool = False,
        self_condition_mlp_projection=False,
        **kwargs
    ):
        super().__init__(**kwargs)
        self.self_condition = self_condition
        self.self_condition_zeros_after_softmax = self_condition_zeros_after_softmax
        self.deepmind_conditional = deepmind_conditional
        self.classifier_free_simplex_inputs = classifier_free_simplex_inputs
        self.self_condition_mlp_projection = self_condition_mlp_projection

sdlm/models/configuration_xlm_roberta.py
allenai/tess-diffusion
import torch
from typing import Optional, Tuple, Union
from transformers.models.roberta.modeling_roberta import RobertaPreTrainedModel, RobertaModel, RobertaLMHead
from transformers.utils import logging
from transformers.modeling_outputs import MaskedLMOutput
from torch.nn import CrossEntropyLoss
import torch.nn as nn
import torch.nn.functional as F
import pdb
import random
from sdlm.utils import convert_to_simplex
from transformers.activations import ACT2FN
from sdlm.utils import mix_values_based_on_self_condition

logger = logging.get_logger(__name__)


class RobertaForDiffusionLM(RobertaPreTrainedModel):
    _keys_to_ignore_on_save = [r"lm_head.decoder.weight", r"lm_head.decoder.bias"]
    _keys_to_ignore_on_load_missing = [r"position_ids", r"lm_head.decoder.weight", r"lm_head.decoder.bias"]
    _keys_to_ignore_on_load_unexpected = [r"pooler"]

    def __init__(self, config):
        super().__init__(config)

        if config.is_decoder:
            logger.warning(
                "If you want to use `RobertaForMaskedLM` make sure `config.is_decoder=False` for "
                "bi-directional self-attention."
            )

        self.roberta = RobertaModel(config, add_pooling_layer=False)
        self.lm_head = RobertaLMHead(config)

        # The LM head weights require special treatment only when they are tied with the word embeddings
        self.update_keys_to_ignore(config, ["lm_head.decoder.weight"])

        self.vocab_to_hidden_dim_embed = nn.Linear(config.vocab_size, config.hidden_size, bias=False)
        self.timestep_embed = nn.Linear(1, config.hidden_size, bias=True)

        if self.config.self_condition is not None and self.config.deepmind_conditional:
            # In this case, this is self-conditioning with conditional generation as done in DeepMind paper.
            # See Figure 3 in https://arxiv.org/pdf/2211.15089.pdf.
            # Here we concat masked word embeddings, noisy embeddings, mask, and self-conditioning inputs
            # and project them to the hidden_size.
            self.project_to_hidden_size = nn.Linear(config.hidden_size * 4, config.hidden_size, bias=False)
        elif self.config.self_condition is not None and not self.config.self_condition in [
            "logits_addition",
            "logits_with_projection_addition",
            "logits_max",
            "logits_mean",
        ]:
            if config.self_condition_mlp_projection:
                self.project_to_hidden_size = nn.Sequential(
                    nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False),
                    ACT2FN[config.hidden_act],
                    nn.Linear(config.hidden_size, config.hidden_size, bias=False),
                )
            else:
                self.project_to_hidden_size = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False)

        # Initialize weights and apply final processing
        self.post_init()

    def post_init(self):
        super().post_init()
        self.vocab_to_hidden_dim_embed.weight.data = self.get_input_embeddings().weight.data.T

    def get_output_embeddings(self):
        return self.lm_head.decoder

    def set_output_embeddings(self, new_embeddings):
        self.lm_head.decoder = new_embeddings

    def get_roberta_empty_tokens(self, shape, device):
        if self.config.empty_token_be_mask:
            empty_token_ids = torch.ones(shape, dtype=torch.int64, device=device)*50264
        else:
            # Padding token in roberta-large is 1.
            empty_token_ids = torch.ones(shape, dtype=torch.int64, device=device)
        empty_token_ids[:, 0] = 0
        empty_token_ids[:, -1] = 2
        return empty_token_ids

    def forward(
        self,
        timesteps: torch.FloatTensor,
        input_ids: torch.LongTensor,
        simplex: torch.FloatTensor,
        span_mask: Optional[torch.FloatTensor] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        token_type_ids: Optional[torch.LongTensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        encoder_hidden_states: Optional[torch.FloatTensor] = None,
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        previous_pred: Optional[torch.FloatTensor] = None,
        classifier_free_guidance: bool = False,
        classifier_free_guidance_in_train: bool = False,
        # unconditional_simplex: torch.FloatTensor = None,
    ) -> Union[Tuple[torch.Tensor], MaskedLMOutput]:
        r"""
        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
            config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
            loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
        kwargs (`Dict[str, any]`, optional, defaults to *{}*):
            Used to hide legacy arguments that have been deprecated.
        """
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        # If we have a mask, we need to mask the simplex values before softmax.
        """
        if span_mask is not None:
            mask_value = torch.finfo(simplex.dtype).min
            mask_value = torch.tensor(mask_value, dtype=simplex.dtype, device=simplex.device)
            simplex = torch.where(span_mask[:, :, None], simplex, mask_value)
        """
        inputs_probs = F.softmax(simplex, dim=-1)
        seq_length = inputs_probs.shape[1]
        inputs_embeds = self.vocab_to_hidden_dim_embed(inputs_probs)

        if classifier_free_guidance or classifier_free_guidance_in_train:
            if self.config.classifier_free_simplex_inputs:
                if self.config.classifier_free_uncond_input == "empty_token":
                    empty_token_ids = self.get_roberta_empty_tokens(shape=input_ids.shape, device=input_ids.device)
                    # TODO: fix the simplex_value later.
                    unconditional_simplex = convert_to_simplex(empty_token_ids, 5.0, self.config.vocab_size)
                elif self.config.classifier_free_uncond_input == "noisy_simplex":
                    simplex_shape = (input_ids.shape[0], input_ids.shape[1], self.config.vocab_size)
                    unconditional_simplex = 5.0 * torch.randn(simplex_shape, device=input_ids.device)
                else:
                    raise NotImplementedError
                unconditional_probs = F.softmax(unconditional_simplex, dim=-1)
                uncond_inputs_embeds = self.vocab_to_hidden_dim_embed(unconditional_probs)
            else:
                empty_token_ids = self.get_roberta_empty_tokens(shape=input_ids.shape, device=input_ids.device)
                uncond_inputs_embeds = self.get_input_embeddings()(empty_token_ids)

        if self.config.self_condition is not None:
            if self.config.self_condition_zeros_after_softmax and previous_pred is None:
                previous_pred_probs = torch.zeros_like(simplex, device=simplex.device)
            else:
                if previous_pred is None:
                    previous_pred = torch.zeros_like(simplex, device=simplex.device)
                """
                if span_mask is not None:
                    mask_value = torch.finfo(previous_pred.dtype).min
                    mask_value = torch.tensor(mask_value, dtype=previous_pred.dtype, device=previous_pred.device)
                    previous_pred = torch.where(span_mask[:, :, None], previous_pred, mask_value)
                """
                previous_pred_probs = F.softmax(previous_pred, dim=-1)
            if not self.config.self_condition_mix_logits_before_weights:
                previous_pred = self.vocab_to_hidden_dim_embed(previous_pred_probs)
            if not self.config.deepmind_conditional:
                # In this setting, we mix the probabilities then apply the weight.
                if self.config.self_condition_mix_logits_before_weights:
                    mixed_logits = mix_values_based_on_self_condition(self.config.self_condition, simplex, previous_pred)
                    mixed_probs = F.softmax(mixed_logits, dim=-1)
                    inputs_embeds = self.vocab_to_hidden_dim_embed(mixed_probs)
                elif self.config.self_condition_mix_before_weights:
                    mixed_probs = mix_values_based_on_self_condition(
                        self.config.self_condition, inputs_probs, previous_pred_probs
                    )
                    inputs_embeds = self.vocab_to_hidden_dim_embed(mixed_probs)
                else:
                    if self.config.self_condition in ["logits", "logits_with_projection"]:
                        inputs_embeds = self.project_to_hidden_size(torch.cat([inputs_embeds, previous_pred], axis=-1))
                    else:
                        inputs_embeds = mix_values_based_on_self_condition(
                            self.config.self_condition, inputs_embeds, previous_pred
                        )

        if span_mask is not None:
            # Original word embeddings without noise.
            if classifier_free_guidance_in_train and random.uniform(0, 1) < 0.1:
                inputs_word_embeds = uncond_inputs_embeds
            else:
                inputs_word_embeds = self.get_input_embeddings()(input_ids)

        if self.config.self_condition is not None and self.config.deepmind_conditional:
            inputs_embeds = torch.where(span_mask.unsqueeze(-1), inputs_embeds, torch.zeros_like(previous_pred))
            previous_pred = torch.where(span_mask.unsqueeze(-1), previous_pred, torch.zeros_like(previous_pred))
            inputs_word_embeds = torch.where(
                span_mask.unsqueeze(-1), torch.zeros_like(inputs_word_embeds), inputs_word_embeds
            )
            tiled_mask = span_mask.unsqueeze(-1).repeat(1, 1, self.config.hidden_size)
            inputs_embeds = self.project_to_hidden_size(
                torch.cat([inputs_embeds, inputs_word_embeds, previous_pred, tiled_mask], axis=-1)
            )

        # TODO: remove conversion.
        timesteps_embed = self.timestep_embed(timesteps.view(-1, 1).float())
        inputs_embeds = inputs_embeds + timesteps_embed.unsqueeze(1).repeat(1, seq_length, 1)

        if span_mask is not None and not self.config.deepmind_conditional:
            # For the unmasked tokens, we only compute their original word embeddings.
            # Note that this also sets the self-conditioned inputs wich we are conditioning on
            # to their original word embeddings values.
            inputs_embeds = torch.where(span_mask.unsqueeze(-1), inputs_embeds, inputs_word_embeds)
            # TODO: we need to fix classifier-free guidance for the case of deepmind_conditional.
            if classifier_free_guidance:
                inputs_embeds = torch.cat([uncond_inputs_embeds, inputs_embeds])
        outputs = self.roberta(
            input_ids=None,  # TODO(rabeeh): we can remove this hack when we moved loss to outside.
            attention_mask=None,  # attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        sequence_output = outputs[0]
        prediction_scores = self.lm_head(sequence_output)

        masked_lm_loss = None
        # In case of classifier-free guidance, since the number of output logits and input token ids do not match
        # we do not compute the loss.
        if input_ids is not None:
            # In case of classifier_free guidance we need to get rid of the unconditional part.
            prediction_scores_for_loss = prediction_scores.chunk(2)[1] if classifier_free_guidance else prediction_scores
            loss_fct = CrossEntropyLoss()
            labels = torch.where(span_mask, input_ids, -100) if span_mask is not None else input_ids
            masked_lm_loss = loss_fct(prediction_scores_for_loss.view(-1, self.config.vocab_size), labels.view(-1))

        if not return_dict:
            output = (prediction_scores,) + outputs[2:]
            return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output

        return MaskedLMOutput(
            loss=masked_lm_loss,
            logits=prediction_scores,
            hidden_states=outputs.last_hidden_state,
            attentions=outputs.attentions,
        )

    def resize_position_embeddings(self, new_num_position_embeddings: int, with_alternatation=False):
        """
        Resizes position embeddings of the model if `new_num_position_embeddings != config.max_position_embeddings`.
        Arguments:
            new_num_position_embeddings (`int`):
                The number of new position embedding matrix. If position embeddings are learned, increasing the size
                will add newly initialized vectors at the end, whereas reducing the size will remove vectors from the
                end. If position embeddings are not learned (*e.g.* sinusoidal position embeddings), increasing the
                size will add correct vectors at the end following the position encoding algorithm, whereas reducing
                the size will remove vectors from the end.
        """
        num_position_embeds_diff = new_num_position_embeddings - self.config.max_position_embeddings

        # no resizing needs to be done if the length stays the same
        if num_position_embeds_diff == 0:
            return

        logger.info(f"Setting `config.max_position_embeddings={new_num_position_embeddings}`...")
        self.config.max_position_embeddings = new_num_position_embeddings
        old_position_embeddings_weight = self.roberta.embeddings.position_embeddings.weight.clone()

        padding_idx = self.config.pad_token_id
        self.roberta.embeddings.position_embeddings = nn.Embedding(
            self.config.max_position_embeddings, self.config.hidden_size, padding_idx=padding_idx
        )
        with torch.no_grad():
            if num_position_embeds_diff > 0:
                self.roberta.embeddings.position_embeddings.weight[:-num_position_embeds_diff] = nn.Parameter(
                    old_position_embeddings_weight
                )
                if with_alternatation:
                    self.roberta.embeddings.position_embeddings.weight[-num_position_embeds_diff:] = nn.Parameter(
                        old_position_embeddings_weight[:num_position_embeds_diff]
                    )
            else:
                self.roberta.embeddings.position_embeddings.weight = nn.Parameter(
                    old_position_embeddings_weight[:num_position_embeds_diff]
                )
        # move position_embeddings to correct device
        self.roberta.embeddings.position_embeddings.to(self.device)
        # Update other needed parameters.
        self.roberta.embeddings.position_ids = (
            torch.arange(self.config.max_position_embeddings).expand((1, -1)).type_as(self.roberta.embeddings.position_ids)
        )
        self.roberta.embeddings.token_type_ids = torch.zeros(
            self.roberta.embeddings.position_ids.size(), dtype=torch.long
        ).type_as(self.roberta.embeddings.token_type_ids)

        # resize the distance embeddings.
        for i in range(self.config.num_hidden_layers):
            if (
                self.config.position_embedding_type == "relative_key"
                or self.config.position_embedding_type == "relative_key_query"
            ):
                self.roberta.encoder.layer[i].attention.self.distance_embedding = nn.Embedding(
                    2 * self.config.max_position_embeddings - 1, self.attention_head_size
                )
                old_distance_embedding_weight = self.layer[i].attention.self.distance_embedding.weight.clone()
                with torch.no_grad():
                    if num_position_embeds_diff > 0:
                        self.roberta.encoder.layer[i].attention.self.distance_embedding.weight[
                            : -2 * num_position_embeds_diff
                        ] = nn.Parameter(old_distance_embedding_weight)
                    else:
                        self.roberta.encoder.layer[i].attention.self.distance_embedding.weight = nn.Parameter(
                            old_distance_embedding_weight[: 2 * num_position_embeds_diff]
                        )

sdlm/models/modeling_roberta.py
allenai/tess-diffusion
import torch
from typing import Optional, Union, Tuple
from transformers.utils import logging
from transformers.models.xlm_roberta.modeling_xlm_roberta import XLMRobertaPreTrainedModel, XLMRobertaModel, XLMRobertaLMHead
from transformers.modeling_outputs import MaskedLMOutput
from torch.nn import CrossEntropyLoss
import torch.nn as nn
import torch.nn.functional as F
import pdb

logger = logging.get_logger(__name__)


class XLMRobertaForDiffusionLM(XLMRobertaPreTrainedModel):
    _keys_to_ignore_on_save = [r"lm_head.decoder.weight", r"lm_head.decoder.bias"]
    _keys_to_ignore_on_load_missing = [r"position_ids", r"lm_head.decoder.weight", r"lm_head.decoder.bias"]
    _keys_to_ignore_on_load_unexpected = [r"pooler"]

    def __init__(self, config):
        super().__init__(config)

        if config.is_decoder:
            logger.warning(
                "If you want to use `XLMRobertaForMaskedLM` make sure `config.is_decoder=False` for "
                "bi-directional self-attention."
            )

        self.roberta = XLMRobertaModel(config, add_pooling_layer=False)
        self.lm_head = XLMRobertaLMHead(config)

        # The LM head weights require special treatment only when they are tied with the word embeddings
        self.update_keys_to_ignore(config, ["lm_head.decoder.weight"])

        self.vocab_to_hidden_dim_embed = nn.Linear(config.vocab_size, config.hidden_size, bias=False)
        self.timestep_embed = nn.Linear(1, config.hidden_size, bias=True)

        if self.config.self_condition is not None and self.config.deepmind_conditional:
            # In this case, this is self-conditioning with conditional generation as done in DeepMind paper.
            # See Figure 3 in https://arxiv.org/pdf/2211.15089.pdf.
            # Here we concat masked word embeddings, noisy embeddings, mask, and self-conditioning inputs
            # and project them to the hidden_size.
            self.project_to_hidden_size = nn.Linear(config.hidden_size * 4, config.hidden_size, bias=False)
        elif self.config.self_condition is not None and not self.config.self_condition in [
            "logits_addition",
            "logits_with_projection_addition",
        ]:
            self.project_to_hidden_size = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False)

        # Initialize weights and apply final processing
        self.post_init()

    def post_init(self):
        super().post_init()
        self.vocab_to_hidden_dim_embed.weight.data = self.get_input_embeddings().weight.data.T

    def get_output_embeddings(self):
        return self.lm_head.decoder

    def set_output_embeddings(self, new_embeddings):
        self.lm_head.decoder = new_embeddings

    def forward(
        self,
        timesteps: torch.FloatTensor,
        simplex: torch.FloatTensor,
        input_ids: Optional[torch.LongTensor] = None,
        span_mask: Optional[torch.FloatTensor] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        token_type_ids: Optional[torch.LongTensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        encoder_hidden_states: Optional[torch.FloatTensor] = None,
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        previous_pred: Optional[torch.FloatTensor] = None,
        classifier_free_guidance: bool = False,
        unconditional_simplex: torch.FloatTensor = None,
    ) -> Union[Tuple[torch.Tensor], MaskedLMOutput]:
        r"""
        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
            config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
            loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
        kwargs (`Dict[str, any]`, optional, defaults to *{}*):
            Used to hide legacy arguments that have been deprecated.
        """
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        inputs_probs = F.softmax(simplex, dim=-1)
        seq_length = inputs_probs.shape[1]
        inputs_embeds = self.vocab_to_hidden_dim_embed(inputs_probs)

        if classifier_free_guidance:
            unconditional_probs = F.softmax(unconditional_simplex, dim=-1)
            uncond_inputs_embeds = self.vocab_to_hidden_dim_embed(unconditional_probs)

        if self.config.self_condition is not None:
            if self.config.self_condition_zeros_after_softmax and previous_pred is None:
                previous_pred_probs = torch.zeros_like(simplex, device=simplex.device)
            else:
                if previous_pred is None:
                    previous_pred = torch.zeros_like(simplex, device=simplex.device)
                previous_pred_probs = F.softmax(previous_pred, dim=-1)
            previous_pred = self.vocab_to_hidden_dim_embed(previous_pred_probs)
            if not self.config.deepmind_conditional:
                if self.config.self_condition in ["logits_with_projection_addition", "logits_addition"]:
                    inputs_embeds = inputs_embeds + previous_pred
                elif self.config.self_condition in ["logits", "logits_with_projection"]:
                    inputs_embeds = self.project_to_hidden_size(torch.cat([inputs_embeds, previous_pred], axis=-1))
                else:
                    raise NotImplementedError

        if span_mask is not None:
            # Original word embeddings without noise.
            inputs_word_embeds = self.get_input_embeddings()(input_ids)

        if self.config.self_condition is not None and self.config.deepmind_conditional:
            inputs_embeds = torch.where(span_mask.unsqueeze(-1), inputs_embeds, torch.zeros_like(previous_pred))
            previous_pred = torch.where(span_mask.unsqueeze(-1), previous_pred, torch.zeros_like(previous_pred))
            inputs_word_embeds = torch.where(
                span_mask.unsqueeze(-1), torch.zeros_like(inputs_word_embeds), inputs_word_embeds
            )
            tiled_mask = span_mask.unsqueeze(-1).repeat(1, 1, self.config.hidden_size)
            inputs_embeds = self.project_to_hidden_size(
                torch.cat([inputs_embeds, inputs_word_embeds, previous_pred, tiled_mask], axis=-1)
            )

        # TODO: remove conversion.
        timesteps_embed = self.timestep_embed(timesteps.view(-1, 1).float())
        inputs_embeds = inputs_embeds + timesteps_embed.unsqueeze(1).repeat(1, seq_length, 1)

        if span_mask is not None and not self.config.deepmind_conditional:
            # For the unmasked tokens, we only compute their original word embeddings.
            # Note that this also sets the self-conditioned inputs wich we are conditioning on
            # to their original word embeddings values.
            inputs_embeds = torch.where(span_mask.unsqueeze(-1), inputs_embeds, inputs_word_embeds)
            # TODO: we need to fix classifier-free guidance for the case of deepmind_conditional.
            if classifier_free_guidance:
                inputs_embeds = torch.cat([uncond_inputs_embeds, inputs_embeds])

        outputs = self.roberta(
            input_ids=None,
            attention_mask=None,  # attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        sequence_output = outputs[0]
        prediction_scores = self.lm_head(sequence_output)

        masked_lm_loss = None
        # In case of classifier-free guidance, since the number of output logits and input token ids do not match
        # we do not compute the loss.
        if input_ids is not None and not classifier_free_guidance:
            loss_fct = CrossEntropyLoss()
            labels = torch.where(span_mask, input_ids, -100) if span_mask is not None else input_ids
            masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))

        if not return_dict:
            output = (prediction_scores,) + outputs[2:]
            return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output

        return MaskedLMOutput(
            loss=masked_lm_loss,
            logits=prediction_scores,
            hidden_states=outputs.last_hidden_state,
            attentions=outputs.attentions,
        )

sdlm/models/modeling_xlm_roberta.py
allenai/tess-diffusion
from typing import Optional, Tuple, Union

import torch
import pdb
from diffusers.pipeline_utils import DiffusionPipeline
from sdlm.inference.inference_utils import logits_projection
from sdlm.utils import scale, self_condition_preds
from dataclasses import dataclass
import numpy as np
from diffusers.utils import BaseOutput
from sdlm.utils import convert_to_simplex
import torch.nn.functional as F

@dataclass
class SimplexDiffusionPipelineOutput(BaseOutput):
    """
    Output class for simplex diffusion pipelines.
    Args:
        simplex (`np.ndarray`)
            numpy array showing the denoised simplex representation.
        logits (`np.ndarray`) final generated logits before applying the projection.
    """

    simplex: np.ndarray
    logits: np.ndarray
    loss: np.ndarray


class SimplexDDPMPipeline(DiffusionPipeline):
    r"""
    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
    Parameters:
        model: Model architecture to denoise the latents (encoded token ids).
        scheduler ([`SchedulerMixin`]): A scheduler to denoise the encoded latent.
    """

    def __init__(
        self,
        model,
        scheduler,
        simplex_value,
        top_p,
        sampling_type,
        is_conditional_generation,
        tokenizer,
        classifier_free_uncond_input,
        temperature,
        guidance_softmax_combination
    ):
        super().__init__()
        self.register_modules(model=model, scheduler=scheduler)
        self.simplex_value = simplex_value
        self.top_p = top_p
        self.sampling_type = sampling_type
        self.is_conditional_generation = is_conditional_generation
        self.tokenizer = tokenizer
        self.classifier_free_uncond_input = classifier_free_uncond_input
        self.temperature = temperature
        self.guidance_softmax_combination=guidance_softmax_combination

    @torch.no_grad()
    def __call__(
        self,
        batch_size: int = 1,
        seq_length: int = 512,
        generator: Optional[torch.Generator] = None,
        batch: Optional[torch.FloatTensor] = None,
        guidance_scale: float = 1.0,
    ) -> Union[SimplexDiffusionPipelineOutput, Tuple]:
        r"""
        Args:
            batch_size (`int`, *optional*, defaults to 1):
                The number of images to generate.
            seq_length: (`int`), sequence length for the generated samples.
            generator (`torch.Generator`, *optional*):
                A [torch generator](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation
                deterministic.
            batch (`torch.FloatTensor`): batch of input data, mostly used in the conditional generation setting.
        Returns:
            [`~pipeline_utils.SimplexDiffusionPipelineOutput`]: returns the generated simplex.
        """
        # Classifier_free guidance works only in the conditional generation case.
        classifier_free_guidance = guidance_scale > 1.0 and self.is_conditional_generation
        """
        if classifier_free_guidance:
            # Makes unconditional input for max sequence length, later we truncate it.
            uncond_input = self.tokenizer(
                [""] * batch_size, padding="max_length", max_length=seq_length, return_tensors="pt"
            ).to(self.device)
            # Converts this to a simplex (batch_size, max_seq, vocab_size)
            uncond_simplex = convert_to_simplex(uncond_input["input_ids"], self.simplex_value, self.model.config.vocab_size)
        """
        # Sample gaussian noise to begin loop
        vocab_size = self.model.config.vocab_size
        if batch is not None:
            # TODO(rabeeh): is giving the length cheating for this setting?
            # Adapts the sequence length to the given `span_mask`'s length.
            seq_length = batch["input_ids"].shape[1]
        simplex_shape = (batch_size, seq_length, vocab_size)
        simplex = self.simplex_value * torch.randn(simplex_shape, generator=generator, device=self.device)
        if self.model.config.self_condition is not None:
            previous_pred = torch.zeros((batch_size, seq_length, vocab_size), device=self.device)
        logits_projection_fct = lambda x: logits_projection(
            x, self.sampling_type, self.top_p, self.simplex_value, self.temperature
        )

        for t in self.progress_bar(self.scheduler.timesteps):
            # TODO(rabeeh): also check without the scale.
            t_scaled = scale(t, len(self.scheduler))
            """
            if classifier_free_guidance:
                if self.classifier_free_uncond_input == "empty_token":
                    uncond_input = uncond_simplex[:, : batch["input_ids"].shape[1], :]
                elif self.classifier_free_uncond_input == "noisy_simplex":
                    uncond_input = self.simplex_value * torch.randn(simplex.shape, generator=generator, device=self.device)
                else:
                    raise NotImplementedError
            """
            # 1. predict noise model_output. Note we need not to pass the input_ids in case of
            # unconditional generation since the loss would be computed and it should not.
            model_output = self.model(
                input_ids=batch["input_ids"] if self.is_conditional_generation else None,
                span_mask=batch["span_mask"] if self.is_conditional_generation else None,
                simplex=simplex,
                timesteps=t_scaled,
                previous_pred=previous_pred if self.model.config.self_condition else None,
                classifier_free_guidance=classifier_free_guidance,
                # unconditional_simplex=uncond_input if classifier_free_guidance else None,
            )
            model_output_logits = model_output.logits

            # Performs classifier-free guidance.
            if classifier_free_guidance:
                logits_uncond, logits_pred = model_output_logits.chunk(2)
                if self.guidance_softmax_combination:
                    model_output_logits = F.softmax(logits_uncond, dim=-1) +  guidance_scale * (F.softmax(logits_pred, dim=-1) - F.softmax(logits_uncond, dim=-1))
                else:
                    model_output_logits = logits_uncond + guidance_scale * (logits_pred - logits_uncond)

            if self.model.config.self_condition is not None:
                if classifier_free_guidance:
                    prev_output_logits = model_output.logits.chunk(2)[1]
                else:
                    prev_output_logits = model_output_logits

                previous_pred = self_condition_preds(
                    self.model.config.self_condition, prev_output_logits, logits_projection_fct
                )

            # Projection.
            projected_logits = logits_projection_fct(model_output_logits)

            # 2. compute previous logits: x_t -> x_t-1
            noise = self.simplex_value * torch.randn(simplex_shape, generator=generator, device=self.device)
            simplex = self.scheduler.step(projected_logits, t, noise, generator=generator).prev_sample

        return SimplexDiffusionPipelineOutput(simplex=simplex, logits=model_output_logits, loss=model_output.loss)

sdlm/pipelines/simplex_ddpm.py
allenai/tess-diffusion
from .scheduling_simplex_ddpm import SimplexDDPMScheduler

sdlm/schedulers/__init__.py
allenai/tess-diffusion
"""DDPM scheduler for the simplex diffusion model."""

from diffusers import DDPMScheduler
from diffusers.schedulers.scheduling_ddpm import DDPMSchedulerOutput
from dataclasses import dataclass
from typing import Union, Tuple, Optional
import torch
import numpy as np
from diffusers.configuration_utils import register_to_config
from diffusers.utils import BaseOutput
import math
import pdb


@dataclass
class SimplexDDPMSchedulerOutput(BaseOutput):
    """
    Output class for the scheduler's step function output.
    Args:
        prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
            Computed sample (x_{t-1}) of previous timestep. `prev_sample` should be used as next model input in the
            denoising loop.
        projected_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, vocab_size)`):
            The projected logits sample (x_{0}) based on the model output from the current timestep.
    """

    prev_sample: torch.FloatTensor
    projected_logits: Optional[torch.FloatTensor] = None


def betas_for_alpha_bar(num_diffusion_timesteps, device, max_beta=0.999, improved_ddpm=False):
    """
    Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of
    (1-beta) over time from t = [0,1].
    Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up
    to that part of the diffusion process.
    Args:
        num_diffusion_timesteps (`int`): the number of betas to produce.
        max_beta (`float`): the maximum beta to use; use values lower than 1 to
                     prevent singularities.
    Returns:
        betas (`np.ndarray`): the betas used by the scheduler to step the model outputs
    """

    def default_alpha_bar(time_step):
        return math.cos((time_step + 1e-4) / (1 + 1e-4) * math.pi / 2) ** 2

    if improved_ddpm:
        # Implements eqn. 17 in https://arxiv.org/pdf/2102.09672.pdf.
        alpha_bar = lambda x: (default_alpha_bar(x) / default_alpha_bar(0.0))
        alphas_cumprod = []
    else:
        alpha_bar = default_alpha_bar
    betas = []
    for i in range(num_diffusion_timesteps):
        t1 = i / num_diffusion_timesteps
        t2 = (i + 1) / num_diffusion_timesteps
        alpha_bar_t1 = alpha_bar(t1)
        betas.append(min(1 - alpha_bar(t2) / alpha_bar_t1, max_beta))
        if improved_ddpm:
            alphas_cumprod.append(alpha_bar_t1)
    # TODO(rabeeh): maybe this cause memory issue.
    betas = torch.tensor(betas, dtype=torch.float32, device=device)
    if improved_ddpm:
        return betas, torch.tensor(alphas_cumprod, dtype=torch.torch.float32, device=device)
    return betas


class SimplexDDPMScheduler(DDPMScheduler):
    @register_to_config
    def __init__(
        self,
        device,
        simplex_value: float,
        num_train_timesteps: int = 1000,
        num_inference_timesteps: int = 1000,
        beta_start: float = 0.0001,
        beta_end: float = 0.02,
        beta_schedule: str = "linear",
        trained_betas: Optional[np.ndarray] = None,
        variance_type: str = "fixed_small",
        clip_sample: bool = False,
    ):
        if trained_betas is not None:
            self.betas = torch.from_numpy(trained_betas)
        elif beta_schedule == "linear":
            self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32, device=device)
        elif beta_schedule == "scaled_linear":
            # this schedule is very specific to the latent diffusion model.
            self.betas = (
                torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32, device=device)
                ** 2
            )
        elif beta_schedule == "squaredcos_cap_v2":
            # Glide cosine schedule
            self.betas = betas_for_alpha_bar(num_train_timesteps, device=device)
        elif beta_schedule == "squaredcos_improved_ddpm":
            self.betas, self.alphas_cumprod = betas_for_alpha_bar(num_train_timesteps, device=device, improved_ddpm=True)
        elif beta_schedule == "sigmoid":
            # GeoDiff sigmoid schedule
            betas = torch.linspace(-6, 6, num_train_timesteps, device=device)
            self.betas = torch.sigmoid(betas) * (beta_end - beta_start) + beta_start
        else:
            raise NotImplementedError(f"{beta_schedule} does is not implemented for {self.__class__}")

        if beta_schedule == "squaredcos_improved_ddpm":
            self.alphas = None
        else:
            self.alphas = 1.0 - self.betas
            self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)

        self.one = torch.tensor(1.0, device=device)

        # standard deviation of the initial noise distribution
        self.init_noise_sigma = 1.0

        # setable values
        self.num_inference_steps = None
        # TODO(rabeeh): if memory issue, we can not add this to GPU and convert them iteratively.
        self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy()).to(device=device)

        self.variance_type = variance_type

    def step(
        self,
        projected_logits: torch.FloatTensor,
        timestep: int,
        noise: torch.FloatTensor,
        generator=None,
    ) -> Union[DDPMSchedulerOutput, Tuple]:
        """
        Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion
        process from the learned model outputs (most often the predicted noise).
        Args:
            projected_logits (`torch.FloatTensor`): projected logits from the diffusion model.
            timestep (`int`): current discrete timestep in the diffusion chain.
            noise (`torch.FloatTensor`): a random noise with simplex_value standard deviation.
            generator: random number generator.
        Returns:
            [`~schedulers.scheduling_utils.DDPMSchedulerOutput`] resulted values.
        """
        t = timestep

        # 1. compute alphas, betas
        alpha_prod_t_prev = self.alphas_cumprod[t - 1] if t > 0 else self.one

        # 3. Clip "predicted x_0"
        if self.config.clip_sample:
            projected_logits = torch.clamp(projected_logits, -1, 1)

        # See algorithm 2 in Figure 3 in https://arxiv.org/pdf/2210.17432.pdf.
        predicted_logits_coeff = alpha_prod_t_prev ** (0.5)
        noise_coeff = (1 - alpha_prod_t_prev) ** (0.5)
        pred_prev_sample = predicted_logits_coeff * projected_logits + noise_coeff * noise

        return SimplexDDPMSchedulerOutput(prev_sample=pred_prev_sample, projected_logits=projected_logits)

    def add_noise(
        self,
        original_samples: torch.FloatTensor,
        noise: torch.FloatTensor,
        timesteps: torch.IntTensor,
    ) -> torch.FloatTensor:
        # Make sure alphas_cumprod and timestep have same device and dtype as original_samples
        # self.alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device, dtype=original_samples.dtype)
        # timesteps = timesteps.to(original_samples.device)

        alphas_cumprod_timesteps = self.alphas_cumprod[timesteps].view(-1, 1, 1)
        sqrt_alpha_prod = alphas_cumprod_timesteps**0.5
        sqrt_one_minus_alpha_prod = (1 - alphas_cumprod_timesteps) ** 0.5
        noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise
        return noisy_samples

sdlm/schedulers/scheduling_simplex_ddpm.py
allenai/tess-diffusion
from transformers import Trainer
from transformers.utils import is_apex_available
from typing import Dict, Union, Any, Optional, List, Tuple, NamedTuple
from torch.utils.data import Dataset
import torch
import pdb
import math
import time
import os
import numpy as np
from torch import nn
import torch.nn.functional as F
from sdlm.utils import convert_to_simplex, scale
from transformers.trainer_pt_utils import (
    nested_detach,
    nested_numpify,
    find_batch_size,
    nested_concat,
    nested_truncate,
    IterableDatasetShard,
)
import torch.distributed as dist
from torch.utils.data.distributed import DistributedSampler
from tqdm.auto import tqdm
from torch.utils.data import RandomSampler
from transformers.debug_utils import DebugOption, DebugUnderflowOverflow
import sys
from transformers.trainer_callback import TrainerState
from transformers.integrations import TensorBoardCallback, hp_params
from transformers.trainer_utils import (
    has_length,
    denumpify_detensorize,
    speed_metrics,
    seed_worker,
    HPSearchBackend,
    TrainOutput,
    ShardedDDPOption,
)
from transformers.utils import logging, is_datasets_available, is_torch_tpu_available
from transformers.pytorch_utils import is_torch_less_than_1_11
from torch.utils.data import DataLoader
from transformers.deepspeed import deepspeed_init
from sdlm.pipelines.simplex_ddpm import SimplexDDPMPipeline
from sdlm.inference.inference_utils import predict_conditional_generated, logits_projection
from sdlm.utils import self_condition_preds, pad_data
from torch.nn import CrossEntropyLoss
from transformers.utils import is_sagemaker_mp_enabled
from transformers import AdamW
import shutil
from transformers.trainer import TRAINER_STATE_NAME

if is_apex_available():
    from apex import amp

if is_datasets_available():
    import datasets

if is_torch_tpu_available(check_device=False):
    import torch_xla.core.xla_model as xm
    import torch_xla.debug.metrics as met
    import torch_xla.distributed.parallel_loader as pl

if is_sagemaker_mp_enabled():
    import smdistributed.modelparallel.torch as smp

IS_SAGEMAKER_MP_POST_1_10 = False
GENERATION_RESULTS = "generated"


logger = logging.get_logger(__name__)


class EvalLoopOutput(NamedTuple):
    logits: Union[np.ndarray, Tuple[np.ndarray]]
    simplex: Union[np.ndarray, Tuple[np.ndarray]]
    input_ids: Optional[Union[np.ndarray, Tuple[np.ndarray]]]
    metrics: Optional[Dict[str, float]]
    results: Optional[Dict[str, List[str]]]
    num_samples: Optional[int]


class DiffusionTrainer(Trainer):
    def __init__(
        self,
        noise_scheduler,
        inference_noise_scheduler,
        diffusion_args,
        data_args,
        *args,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        self.noise_scheduler = noise_scheduler
        self.diffusion_args = diffusion_args
        self.data_args = data_args
        self.vocab_size = self.model.config.vocab_size
        self.inference_noise_scheduler = inference_noise_scheduler
        self.tb_writer = self.get_tb_writer()
        self.eos_token_id = self.tokenizer.eos_token_id
        self.classifier_free_guidance = diffusion_args.guidance_scale > 1.0 and data_args.conditional_generation is not None

    def annotated_split(self, split):
        return f"{split}_top_p_{self.diffusion_args.top_p}_temperature_{self.diffusion_args.temperature}_seed_{self.args.seed}_guidance_scale_{self.diffusion_args.guidance_scale}"

    def save_metrics(self, split, metrics, combined=True):
        super().save_metrics(self.annotated_split(split), metrics, combined)

    def log_metrics(self, split, metrics):
        super().log_metrics(self.annotated_split(split), metrics)


    def get_tb_writer(self):
        for cb in self.callback_handler.callbacks:
            if isinstance(cb, TensorBoardCallback):
                return cb
        return None

    def training_step(self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]]) -> torch.Tensor:
        """
        Perform a training step on a batch of inputs.
        Subclass and override to inject custom behavior.
        Args:
            model (`nn.Module`):
                The model to train.
            inputs (`Dict[str, Union[torch.Tensor, Any]]`):
                The inputs and targets of the model.
                The dictionary will be unpacked before being fed to the model. Most models expect the targets under the
                argument `labels`. Check your model's documentation for all accepted arguments.
        Return:
            `torch.Tensor`: The tensor with training loss on this batch.
        """
        model.train()
        inputs = self._prepare_inputs(inputs)

        # Truncate the length if needed.
        if self.data_args.truncation_length > 0:
            inputs["input_ids"] = inputs["input_ids"][:, : -self.data_args.truncation_length]
            inputs["span_mask"] = inputs["span_mask"][:, : -self.data_args.truncation_length]

        # Creates the noisy simplex and timesteps.
        simplex = convert_to_simplex(inputs["input_ids"], self.diffusion_args.simplex_value, self.vocab_size)
        noise = self.diffusion_args.simplex_value * torch.randn(simplex.shape, device=simplex.device, dtype=simplex.dtype)
        bsz = simplex.shape[0]
        # Sample a random timestep for each simplex token representation.
        timesteps = torch.randint(0, len(self.noise_scheduler), (bsz,), device=simplex.device, dtype=torch.int64)
        # Adds noise to each simplex representation (Forward diffusion process).
        noisy_simplex = self.noise_scheduler.add_noise(simplex, noise, timesteps)
        timesteps = scale(timesteps, len(self.noise_scheduler))

        inputs.update({"timesteps": timesteps, "simplex": noisy_simplex})
        if self.diffusion_args.self_condition is not None:
            previous_pred = None
            if np.random.rand(1) > 0.5:
                outputs = model(**inputs, previous_pred=previous_pred)
                logits_projection_fct = lambda x: logits_projection(
                    x,
                    self.diffusion_args.sampling_type,
                    self.diffusion_args.top_p,
                    self.diffusion_args.simplex_value,
                    self.diffusion_args.temperature,
                )
                previous_pred = self_condition_preds(
                    self.diffusion_args.self_condition, outputs.logits, logits_projection_fct
                )
            inputs.update({"previous_pred": previous_pred})
        # NOTE: we do this after computation of self-conditioning to not affect that one.
        inputs.update({"classifier_free_guidance_in_train": self.classifier_free_guidance})
        with self.autocast_smart_context_manager():
            loss = self.compute_loss(model, inputs)

        if self.args.n_gpu > 1:
            loss = loss.mean()  # mean() to average on multi-gpu parallel training

        if self.args.gradient_accumulation_steps > 1 and not self.deepspeed:
            # deepspeed handles loss scaling by gradient_accumulation_steps in its `backward`
            loss = loss / self.args.gradient_accumulation_steps

        if self.do_grad_scaling:
            self.scaler.scale(loss).backward()
        elif self.use_apex:
            with amp.scale_loss(loss, self.optimizer) as scaled_loss:
                scaled_loss.backward()
        elif self.deepspeed:
            # loss gets scaled under gradient_accumulation_steps in deepspeed
            loss = self.deepspeed.backward(loss)
        else:
            loss.backward()

        return loss.detach()

    def prediction_step(
        self, inputs: Dict[str, Union[torch.Tensor, Any]], pipeline
    ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:

        inputs = self._prepare_inputs(inputs)
        is_conditional_generation = True if "span_mask" in inputs else False
        with torch.no_grad():
            with self.compute_loss_context_manager():
                outputs = pipeline(
                    batch_size=inputs["input_ids"].shape[0]
                    if is_conditional_generation
                    else self.args.per_device_eval_batch_size,
                    seq_length=self.data_args.max_seq_length - self.data_args.truncation_length,
                    batch=inputs,
                    guidance_scale=self.diffusion_args.guidance_scale,
                    generator=torch.Generator(device=self.args.device).manual_seed(self.args.seed) if self.diffusion_args.generate_with_seed else None,
                )
                if is_conditional_generation:
                    loss = outputs.loss.mean().detach()
                else:
                    loss = None
        logits = nested_detach(outputs.logits)
        simplex = nested_detach(outputs.simplex)

        # Computes the evaluation loss from the simplex values.
        if self.args.compute_eval_loss_with_simplex and is_conditional_generation:
            loss_fct = CrossEntropyLoss()
            labels = torch.where(inputs["span_mask"], inputs["input_ids"], -100)
            loss = loss_fct(simplex.view(-1, self.model.config.vocab_size), labels.view(-1)).detach()

        return (simplex, logits, loss)

    def evaluation_loop(
        self,
        dataloader: DataLoader,
        description: str,
        prediction_loss_only: Optional[bool] = None,
        ignore_keys: Optional[List[str]] = None,
        metric_key_prefix: str = "eval",
    ) -> EvalLoopOutput:
        """
        Prediction/evaluation loop, shared by `Trainer.evaluate()` and `Trainer.predict()`.
        Works both with or without labels.
        """
        args = self.args
        is_conditional_generation = self.data_args.conditional_generation is not None
        save_prefixes = is_conditional_generation  # and self.data_args.conditional_generation != "seq2seq"

        prediction_loss_only = prediction_loss_only if prediction_loss_only is not None else args.prediction_loss_only
        # if eval is called w/o train init deepspeed here
        if args.deepspeed and not self.deepspeed:

            # XXX: eval doesn't have `resume_from_checkpoint` arg but we should be able to do eval
            # from the checkpoint eventually
            deepspeed_engine, _, _ = deepspeed_init(self, num_training_steps=0, resume_from_checkpoint=None, inference=True)
            self.model = deepspeed_engine.module
            self.model_wrapped = deepspeed_engine
            self.deepspeed = deepspeed_engine

        model = self._wrap_model(self.model, training=False, dataloader=dataloader)

        # if full fp16 or bf16 eval is wanted and this ``evaluation`` or ``predict`` isn't called
        # while ``train`` is running, cast it to the right dtype first and then put on device
        if not self.is_in_train:
            if args.fp16_full_eval:
                model = model.to(dtype=torch.float16, device=args.device)
            elif args.bf16_full_eval:
                model = model.to(dtype=torch.bfloat16, device=args.device)

        batch_size = self.args.eval_batch_size

        logger.info(f"***** Running {description} *****")
        if has_length(dataloader):
            logger.info(f"  Num examples = {self.num_examples(dataloader)}")
        else:
            logger.info("  Num examples: Unknown")
        logger.info(f"  Batch size = {batch_size}")

        model.eval()

        pipeline = SimplexDDPMPipeline(
            model=model,
            scheduler=self.inference_noise_scheduler,
            simplex_value=self.diffusion_args.simplex_value,
            top_p=self.diffusion_args.top_p,
            sampling_type=self.diffusion_args.sampling_type,
            is_conditional_generation=is_conditional_generation,
            tokenizer=self.tokenizer,
            classifier_free_uncond_input=self.diffusion_args.classifier_free_uncond_input,
            temperature=self.diffusion_args.temperature,
            guidance_softmax_combination=self.diffusion_args.guidance_softmax_combination
        )

        self.callback_handler.eval_dataloader = dataloader
        # Do this before wrapping.
        eval_dataset = getattr(dataloader, "dataset", None)

        # Initialize containers
        # logits/simplex/labels on GPU/TPU (accumulated for eval_accumulation_steps)
        losses_host = None
        logits_host = None
        simplex_host = None
        inputs_host = None
        masks_host = None
        prefixes_host = None

        # logits/simplex/labels on CPU (final containers)
        all_losses = None
        all_logits = None
        all_simplex = None
        all_inputs = None
        all_masks = None
        all_prefixes = None
        observed_num_examples = 0

        # Main evaluation loop
        for step, inputs in enumerate(dataloader):
            has_mask = True if "span_mask" in inputs else False

            # Truncate the length if needed.
            if self.data_args.truncation_length > 0:
                inputs["input_ids"] = inputs["input_ids"][:, : -self.data_args.truncation_length]
                inputs["span_mask"] = inputs["span_mask"][:, : -self.data_args.truncation_length]
                max_seq_length = self.data_args.max_seq_length - self.data_args.truncation_length
                assert self.data_args.eval_context_size < max_seq_length

            # Update the observed num examples
            observed_batch_size = find_batch_size(inputs)
            if observed_batch_size is not None:
                observed_num_examples += observed_batch_size
                # For batch samplers, batch_size is not known by the dataloader in advance.
                if batch_size is None:
                    batch_size = observed_batch_size

            # Prediction step
            simplex, logits, loss = self.prediction_step(inputs, pipeline=pipeline)
            inputs_decode = self._prepare_input(inputs["input_ids"])
            masks = self._prepare_input(inputs["span_mask"]) if has_mask else None
            if save_prefixes:
                prefixes = (
                    pad_data([input[~mask] for input, mask in zip(inputs_decode, masks)], self.tokenizer)
                    if has_mask
                    else None
                )
                prefixes = self._prepare_input(prefixes)
            else:
                prefixes = None
            # Update containers on host
            if prefixes is not None:
                prefixes = self._pad_across_processes(prefixes, pad_index=self.eos_token_id)
                prefixes = self._nested_gather(prefixes)
                prefixes_host = (
                    prefixes
                    if prefixes_host is None
                    else nested_concat(prefixes_host, prefixes, padding_index=self.eos_token_id)
                )
            if inputs_decode is not None:
                inputs_decode = self._pad_across_processes(inputs_decode, pad_index=self.eos_token_id)
                inputs_decode = self._nested_gather(inputs_decode)
                inputs_host = (
                    inputs_decode
                    if inputs_host is None
                    else nested_concat(inputs_host, inputs_decode, padding_index=self.eos_token_id)
                )
            if loss is not None:
                losses = self._nested_gather(loss.repeat(batch_size))
                losses_host = losses if losses_host is None else torch.cat((losses_host, losses), dim=0)
            # Note that this block should be before masks block, since we need masks here.
            if simplex is not None:
                # In case of having a mask softmax is applied over the simplex non-masked values.
                if has_mask:
                    mask_value = torch.finfo(simplex.dtype).min
                    mask_value = torch.tensor(mask_value, dtype=simplex.dtype, device=simplex.device)
                    simplex = torch.where(masks[:, :, None], simplex, mask_value)
                simplex = F.softmax(simplex, dim=-1)
                if self.preprocess_logits_for_metrics is not None:
                    simplex = self.preprocess_logits_for_metrics(simplex)
                simplex = self._pad_across_processes(simplex, pad_index=self.eos_token_id)
                simplex = self._nested_gather(simplex)
                # TODO: note that this is no more a simplex, but the processed one.
                simplex_host = (
                    simplex
                    if simplex_host is None
                    else nested_concat(simplex_host, simplex, padding_index=self.eos_token_id)
                )
            if masks is not None:
                masks = self._pad_across_processes(masks, pad_index=0)
                masks = self._nested_gather(masks)
                # We pad masks with False tokens.
                masks_host = masks if masks_host is None else nested_concat(masks_host, masks, padding_index=0)
            if logits is not None:
                if self.preprocess_logits_for_metrics is not None:
                    logits = self.preprocess_logits_for_metrics(logits)
                logits = self._pad_across_processes(logits, pad_index=self.eos_token_id)
                logits = self._nested_gather(logits)
                logits_host = (
                    logits if logits_host is None else nested_concat(logits_host, logits, padding_index=self.eos_token_id)
                )

            self.control = self.callback_handler.on_prediction_step(args, self.state, self.control)

        # Gather all remaining tensors and put them back on the CPU
        if losses_host is not None:
            losses = nested_numpify(losses_host)
            all_losses = losses if all_losses is None else np.concatenate((all_losses, losses), axis=0)
        if logits_host is not None:
            logits = nested_numpify(logits_host)
            all_logits = logits if all_logits is None else nested_concat(all_logits, logits, padding_index=self.eos_token_id)
        if simplex_host is not None:
            simplex = nested_numpify(simplex_host)
            all_simplex = (
                simplex if all_simplex is None else nested_concat(all_simplex, simplex, padding_index=self.eos_token_id)
            )
        if inputs_host is not None:
            inputs_decode = nested_numpify(inputs_host)
            all_inputs = (
                inputs_decode
                if all_inputs is None
                else nested_concat(all_inputs, inputs_decode, padding_index=self.eos_token_id)
            )
        if masks_host is not None:
            masks = nested_numpify(masks_host)
            all_masks = masks if all_masks is None else nested_concat(all_masks, masks, padding_index=0)
        if prefixes_host is not None:
            prefixes = nested_numpify(prefixes_host)
            all_prefixes = (
                prefixes if all_prefixes is None else nested_concat(all_prefixes, prefixes, padding_index=self.eos_token_id)
            )

        # Number of samples
        num_samples = len(eval_dataset)
        if num_samples == 0 and observed_num_examples > 0:
            num_samples = observed_num_examples

        # Number of losses has been rounded to a multiple of batch_size and in a distributed training, the number of
        # samplers has been rounded to a multiple of batch_size, so we truncate.
        if all_losses is not None:
            all_losses = all_losses[:num_samples]
        if all_masks is not None:
            all_masks = nested_truncate(all_masks, num_samples)
        if all_simplex is not None:
            all_simplex = nested_truncate(all_simplex, num_samples)
        if all_logits is not None:
            all_logits = nested_truncate(all_logits, num_samples)
        if all_inputs is not None:
            all_inputs = nested_truncate(all_inputs, num_samples)
        if all_prefixes is not None:
            all_prefixes = nested_truncate(all_prefixes, num_samples)

        # Generates the texts.
        results = {}
        if is_conditional_generation:

            # We predict the masked tokens only. Here, we compute the masked tokens.
            results.update(
                predict_conditional_generated(
                    all_masks,
                    all_inputs,
                    self.tokenizer,
                    all_simplex,
                    "pred_texts_from_simplex",
                    self.data_args.skip_special_tokens,
                )
            )
            results.update(
                predict_conditional_generated(
                    all_masks,
                    all_inputs,
                    self.tokenizer,
                    all_logits,
                    "pred_texts_from_logits",
                    self.data_args.skip_special_tokens,
                )
            )
        else:
            results.update(
                {
                    "pred_texts_from_simplex": self.tokenizer.batch_decode(
                        all_simplex, skip_special_tokens=self.data_args.skip_special_tokens
                    )
                }
            )
            results.update(
                {
                    "pred_texts_from_logits": self.tokenizer.batch_decode(
                        all_logits, skip_special_tokens=self.data_args.skip_special_tokens
                    )
                }
            )
        if is_conditional_generation:

            results.update(
                {
                    "gold_texts_masked": [
                        self.tokenizer.decode(input[mask], skip_special_tokens=self.data_args.skip_special_tokens)
                        for mask, input in zip(all_masks, all_inputs)
                    ]
                }
            )
            if save_prefixes:
                results.update(
                    {
                        "prefixes": [
                            self.tokenizer.decode(x, skip_special_tokens=True)  # self.data_args.skip_special_tokens)
                            for x in all_prefixes
                        ]
                    }
                )

        # Metrics.
        if self.compute_metrics is not None:
            metrics = self.compute_metrics(results)
        else:
            metrics = {}

        # To be JSON-serializable, we need to remove numpy types or zero-d tensors
        metrics = denumpify_detensorize(metrics)

        if all_losses is not None:
            metrics[f"{metric_key_prefix}_loss"] = all_losses.mean().item()

        # Prefix all keys with metric_key_prefix + '_'
        for key in list(metrics.keys()):
            if not key.startswith(f"{metric_key_prefix}_"):
                metrics[f"{metric_key_prefix}_{key}"] = metrics.pop(key)
        return EvalLoopOutput(
            logits=all_logits,
            simplex=all_simplex,
            input_ids=all_inputs,
            metrics=metrics,
            num_samples=num_samples,
            results=results,
        )

    def evaluate(
        self,
        eval_dataset: Optional[Dataset] = None,
        ignore_keys: Optional[List[str]] = None,
        metric_key_prefix: str = "eval",
    ) -> Dict[str, float]:
        """
        Run evaluation and returns metrics.
        The calling script will be responsible for providing a method to compute metrics, as they are task-dependent
        (pass it to the init `compute_metrics` argument).
        You can also subclass and override this method to inject custom behavior.
        Args:
            eval_dataset (`Dataset`, *optional*):
                Pass a dataset if you wish to override `self.eval_dataset`. If it is a [`~datasets.Dataset`], columns
                not accepted by the `model.forward()` method are automatically removed. It must implement the `__len__`
                method.
            ignore_keys (`Lst[str]`, *optional*):
                A list of keys in the output of your model (if it is a dictionary) that should be ignored when
                gathering predictions.
            metric_key_prefix (`str`, *optional*, defaults to `"eval"`):
                An optional prefix to be used as the metrics key prefix. For example the metrics "bleu" will be named
                "eval_bleu" if the prefix is "eval" (default)
        Returns:
            A dictionary containing the evaluation loss and the potential metrics computed from the predictions. The
            dictionary also contains the epoch number which comes from the training state.
        """
        # memory metrics - must set up as early as possible
        self._memory_tracker.start()
        eval_dataloader = self.get_eval_dataloader(eval_dataset)

        start_time = time.time()

        output = self.evaluation_loop(
            eval_dataloader,
            description="Evaluation",
            # No point gathering the predictions if there are no metrics, otherwise we defer to
            # self.args.prediction_loss_only
            prediction_loss_only=True if self.compute_metrics is None else None,
            ignore_keys=ignore_keys,
            metric_key_prefix=metric_key_prefix,
        )

        # Adds the generated texts to tensorboard.
        if self.args.log_generated_texts:
            self.log_results_to_tensorboard(self.state, output)

        total_batch_size = self.args.eval_batch_size * self.args.world_size
        output.metrics.update(
            speed_metrics(
                metric_key_prefix,
                start_time,
                num_samples=output.num_samples,
                num_steps=math.ceil(output.num_samples / total_batch_size),
            )
        )
        self.log(output.metrics)
        self.control = self.callback_handler.on_evaluate(self.args, self.state, self.control, output.metrics)
        self._memory_tracker.stop_and_update_metrics(output.metrics)

        # Save the results
        self.save_metrics(GENERATION_RESULTS + "_" + metric_key_prefix, output.results)
        logger.info("Results are saved now")

        return output.metrics

    def log_results_to_tensorboard(self, state, output):
        # TODO: we need to fix this which happens during the only eval option.
        if self.tb_writer.tb_writer is None:
            return
        for i in range(len(output.logits)):
            total_text = ""
            for k, v in output.results.items():
                total_text += f"*** {k} ***: {v[i]}" + "  \n"
            self.tb_writer.tb_writer.add_text(f"sample_{i}", total_text, state.global_step)

    def get_train_dataloader(self) -> DataLoader:
        """
        Returns the training [`~torch.utils.data.DataLoader`].
        Will use no sampler if `train_dataset` does not implement `__len__`, a random sampler (adapted to distributed
        training if necessary) otherwise.
        Subclass and override this method if you want to inject some custom behavior.
        """
        if self.train_dataset is None:
            raise ValueError("Trainer: training requires a train_dataset.")

        train_dataset = self.train_dataset
        data_collator = self.data_collator("train")
        if is_datasets_available() and isinstance(train_dataset, datasets.Dataset):
            train_dataset = self._remove_unused_columns(train_dataset, description="training")
        else:
            data_collator = self._get_collator_with_removed_columns(data_collator, description="training")

        train_sampler = self._get_train_sampler()

        return DataLoader(
            train_dataset,
            batch_size=self._train_batch_size,
            sampler=train_sampler,
            collate_fn=data_collator,
            drop_last=self.args.dataloader_drop_last,
            num_workers=self.args.dataloader_num_workers,
            pin_memory=self.args.dataloader_pin_memory,
            worker_init_fn=seed_worker,
        )

    def get_eval_dataloader(self, eval_dataset: Optional[Dataset] = None) -> DataLoader:
        """
        Returns the evaluation [`~torch.utils.data.DataLoader`].
        Subclass and override this method if you want to inject some custom behavior.
        Args:
            eval_dataset (`torch.utils.data.Dataset`, *optional*):
                If provided, will override `self.eval_dataset`. If it is a [`~datasets.Dataset`], columns not accepted
                by the `model.forward()` method are automatically removed. It must implement `__len__`.
        """
        if eval_dataset is None and self.eval_dataset is None:
            raise ValueError("Trainer: evaluation requires an eval_dataset.")
        eval_dataset = eval_dataset if eval_dataset is not None else self.eval_dataset
        data_collator = self.data_collator("eval")

        if is_datasets_available() and isinstance(eval_dataset, datasets.Dataset):
            eval_dataset = self._remove_unused_columns(eval_dataset, description="evaluation")
        else:
            data_collator = self._get_collator_with_removed_columns(data_collator, description="evaluation")

        eval_sampler = self._get_eval_sampler(eval_dataset)

        return DataLoader(
            eval_dataset,
            sampler=eval_sampler,
            batch_size=self.args.eval_batch_size,
            collate_fn=data_collator,
            drop_last=self.args.dataloader_drop_last,
            num_workers=self.args.dataloader_num_workers,
            pin_memory=self.args.dataloader_pin_memory,
        )

    def create_optimizer(self):
        """
        Setup the optimizer.
        We provide a reasonable default that works well. If you want to use something else, you can pass a tuple in the
        Trainer's init through `optimizers`, or subclass and override this method in a subclass.
        """
        if not self.args.ssdlm_optimizer:
            return super().create_optimizer()

        # SDDLM optimizer.
        opt_model = self.model

        if self.optimizer is None:
            no_decay = ["bias", "LayerNorm.weight", "timestep_embed.weight", "timestep_embed.bias"]
            optimizer_grouped_parameters = [
                {
                    "params": [p for n, p in opt_model.named_parameters() if not any(nd in n for nd in no_decay)],
                    "weight_decay": self.args.weight_decay,
                },
                {
                    "params": [p for n, p in opt_model.named_parameters() if any(nd in n for nd in no_decay)],
                    "weight_decay": 0.0,
                },
            ]
            self.optimizer = AdamW(optimizer_grouped_parameters, lr=self.args.learning_rate)
        return self.optimizer

    def _rotate_checkpoints(self, use_mtime=False, output_dir=None) -> None:
        if self.args.save_total_limit is None or self.args.save_total_limit <= 0:
            return

        # Check if we should delete older checkpoint(s)
        checkpoints_sorted = self._sorted_checkpoints(use_mtime=use_mtime, output_dir=output_dir)
        if len(checkpoints_sorted) <= self.args.save_total_limit:
            return

        # If save_total_limit=1 with load_best_model_at_end=True, we could end up deleting the last checkpoint, which
        # we don't do to allow resuming.
        save_total_limit = self.args.save_total_limit
        if (
            self.state.best_model_checkpoint is not None
            and self.args.save_total_limit == 1
            and checkpoints_sorted[-1] != self.state.best_model_checkpoint
        ):
            save_total_limit = 2

        number_of_checkpoints_to_delete = max(0, len(checkpoints_sorted) - save_total_limit)
        checkpoints_to_be_deleted = checkpoints_sorted[:number_of_checkpoints_to_delete]
        for checkpoint in checkpoints_to_be_deleted:
            logger.info(f"Deleting older checkpoint [{checkpoint}] due to args.save_total_limit")
            shutil.rmtree(checkpoint)

    def _inner_training_loop(
        self, batch_size=None, args=None, resume_from_checkpoint=None, trial=None, ignore_keys_for_eval=None
    ):
        self._train_batch_size = batch_size
        # Data loader and number of training steps
        train_dataloader = self.get_train_dataloader()

        # Setting up training control variables:
        # number of training epochs: num_train_epochs
        # number of training steps per epoch: num_update_steps_per_epoch
        # total number of training steps to execute: max_steps
        total_train_batch_size = args.train_batch_size * args.gradient_accumulation_steps * args.world_size

        len_dataloader = None
        if has_length(train_dataloader):
            len_dataloader = len(train_dataloader)
            num_update_steps_per_epoch = len_dataloader // args.gradient_accumulation_steps
            num_update_steps_per_epoch = max(num_update_steps_per_epoch, 1)
            num_examples = self.num_examples(train_dataloader)
            if args.max_steps > 0:
                max_steps = args.max_steps
                num_train_epochs = args.max_steps // num_update_steps_per_epoch + int(
                    args.max_steps % num_update_steps_per_epoch > 0
                )
                # May be slightly incorrect if the last batch in the training dataloader has a smaller size but it's
                # the best we can do.
                num_train_samples = args.max_steps * total_train_batch_size
            else:
                max_steps = math.ceil(args.num_train_epochs * num_update_steps_per_epoch)
                num_train_epochs = math.ceil(args.num_train_epochs)
                num_train_samples = self.num_examples(train_dataloader) * args.num_train_epochs
        elif args.max_steps > 0:  # Rely on max_steps when dataloader does not have a working size
            max_steps = args.max_steps
            # Setting a very large number of epochs so we go as many times as necessary over the iterator.
            num_train_epochs = sys.maxsize
            num_update_steps_per_epoch = max_steps
            num_examples = total_train_batch_size * args.max_steps
            num_train_samples = args.max_steps * total_train_batch_size
        else:
            raise ValueError(
                "args.max_steps must be set to a positive value if dataloader does not have a length, was"
                f" {args.max_steps}"
            )

        if DebugOption.UNDERFLOW_OVERFLOW in self.args.debug:
            if self.args.n_gpu > 1:
                # nn.DataParallel(model) replicates the model, creating new variables and module
                # references registered here no longer work on other gpus, breaking the module
                raise ValueError(
                    "Currently --debug underflow_overflow is not supported under DP. Please use DDP"
                    " (torch.distributed.launch)."
                )
            else:
                debug_overflow = DebugUnderflowOverflow(self.model)  # noqa

        delay_optimizer_creation = (
            self.sharded_ddp is not None
            and self.sharded_ddp != ShardedDDPOption.SIMPLE
            or is_sagemaker_mp_enabled()
            or self.fsdp is not None
        )
        if args.deepspeed:
            deepspeed_engine, optimizer, lr_scheduler = deepspeed_init(
                self, num_training_steps=max_steps, resume_from_checkpoint=resume_from_checkpoint
            )
            self.model = deepspeed_engine.module
            self.model_wrapped = deepspeed_engine
            self.deepspeed = deepspeed_engine
            self.optimizer = optimizer
            self.lr_scheduler = lr_scheduler
        elif not delay_optimizer_creation:
            self.create_optimizer_and_scheduler(num_training_steps=max_steps)

        self.state = TrainerState()
        self.state.is_hyper_param_search = trial is not None

        # Activate gradient checkpointing if needed
        if args.gradient_checkpointing:
            self.model.gradient_checkpointing_enable()

        model = self._wrap_model(self.model_wrapped)

        if is_sagemaker_mp_enabled() and resume_from_checkpoint is not None:
            self._load_from_checkpoint(resume_from_checkpoint, model)

        # for the rest of this function `model` is the outside model, whether it was wrapped or not
        if model is not self.model:
            self.model_wrapped = model

        if delay_optimizer_creation:
            self.create_optimizer_and_scheduler(num_training_steps=max_steps)

        # Check if saved optimizer or scheduler states exist
        self._load_optimizer_and_scheduler(resume_from_checkpoint)

        # important: at this point:
        # self.model         is the Transformers Model
        # self.model_wrapped is DDP(Transformers Model), Deepspeed(Transformers Model), etc.

        # Train!
        logger.info("***** Running training *****")
        logger.info(f"  Num examples = {num_examples}")
        logger.info(f"  Num Epochs = {num_train_epochs}")
        logger.info(f"  Instantaneous batch size per device = {args.per_device_train_batch_size}")
        logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_train_batch_size}")
        logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
        logger.info(f"  Total optimization steps = {max_steps}")
        logger.info(f"  Number of trainable parameters = {sum(p.numel() for p in model.parameters() if p.requires_grad)}")

        self.state.epoch = 0
        start_time = time.time()
        epochs_trained = 0
        steps_trained_in_current_epoch = 0
        steps_trained_progress_bar = None

        # Check if continuing training from a checkpoint
        if resume_from_checkpoint is not None and os.path.isfile(os.path.join(resume_from_checkpoint, TRAINER_STATE_NAME)):
            self.state = TrainerState.load_from_json(os.path.join(resume_from_checkpoint, TRAINER_STATE_NAME))
            epochs_trained = self.state.global_step // num_update_steps_per_epoch
            if not args.ignore_data_skip:
                steps_trained_in_current_epoch = self.state.global_step % (num_update_steps_per_epoch)
                steps_trained_in_current_epoch *= args.gradient_accumulation_steps
            else:
                steps_trained_in_current_epoch = 0

            logger.info("  Continuing training from checkpoint, will skip to saved global_step")
            logger.info(f"  Continuing training from epoch {epochs_trained}")
            logger.info(f"  Continuing training from global step {self.state.global_step}")
            if not args.ignore_data_skip:
                logger.info(
                    f"  Will skip the first {epochs_trained} epochs then the first {steps_trained_in_current_epoch} "
                    "batches in the first epoch. If this takes a lot of time, you can add the `--ignore_data_skip` "
                    "flag to your launch command, but you will resume the training on data already seen by your model."
                )
                if self.is_local_process_zero() and not args.disable_tqdm:
                    steps_trained_progress_bar = tqdm(total=steps_trained_in_current_epoch)
                    steps_trained_progress_bar.set_description("Skipping the first batches")

        # Update the references
        self.callback_handler.model = self.model
        self.callback_handler.optimizer = self.optimizer
        self.callback_handler.lr_scheduler = self.lr_scheduler
        self.callback_handler.train_dataloader = train_dataloader
        if self.hp_name is not None and self._trial is not None:
            # use self._trial because the SigOpt/Optuna hpo only call `_hp_search_setup(trial)` instead of passing trial
            # parameter to Train when using DDP.
            self.state.trial_name = self.hp_name(self._trial)
        if trial is not None:
            assignments = trial.assignments if self.hp_search_backend == HPSearchBackend.SIGOPT else trial
            self.state.trial_params = hp_params(assignments)
        else:
            self.state.trial_params = None
        # This should be the same if the state has been saved but in case the training arguments changed, it's safer
        # to set this after the load.
        self.state.max_steps = max_steps
        self.state.num_train_epochs = num_train_epochs
        self.state.is_local_process_zero = self.is_local_process_zero()
        self.state.is_world_process_zero = self.is_world_process_zero()

        # tr_loss is a tensor to avoid synchronization of TPUs through .item()
        tr_loss = torch.tensor(0.0).to(args.device)
        # _total_loss_scalar is updated everytime .item() has to be called on tr_loss and stores the sum of all losses
        self._total_loss_scalar = 0.0
        self._globalstep_last_logged = self.state.global_step
        model.zero_grad()

        self.control = self.callback_handler.on_train_begin(args, self.state, self.control)

        # Skip the first epochs_trained epochs to get the random state of the dataloader at the right point.
        if not args.ignore_data_skip:
            for epoch in range(epochs_trained):
                is_random_sampler = hasattr(train_dataloader, "sampler") and isinstance(
                    train_dataloader.sampler, RandomSampler
                )
                if is_torch_less_than_1_11 or not is_random_sampler:
                    # We just need to begin an iteration to create the randomization of the sampler.
                    # That was before PyTorch 1.11 however...
                    for _ in train_dataloader:
                        break
                else:
                    # Otherwise we need to call the whooooole sampler cause there is some random operation added
                    # AT THE VERY END!
                    _ = list(train_dataloader.sampler)

        for epoch in range(epochs_trained, num_train_epochs):
            if isinstance(train_dataloader, DataLoader) and isinstance(train_dataloader.sampler, DistributedSampler):
                train_dataloader.sampler.set_epoch(epoch)
            elif hasattr(train_dataloader, "dataset") and isinstance(train_dataloader.dataset, IterableDatasetShard):
                train_dataloader.dataset.set_epoch(epoch)

            if is_torch_tpu_available():
                parallel_loader = pl.ParallelLoader(train_dataloader, [args.device]).per_device_loader(args.device)
                epoch_iterator = parallel_loader
            else:
                epoch_iterator = train_dataloader

            # Reset the past mems state at the beginning of each epoch if necessary.
            if args.past_index >= 0:
                self._past = None

            steps_in_epoch = (
                len(epoch_iterator) if len_dataloader is not None else args.max_steps * args.gradient_accumulation_steps
            )
            self.control = self.callback_handler.on_epoch_begin(args, self.state, self.control)

            if epoch == epochs_trained and resume_from_checkpoint is not None and steps_trained_in_current_epoch == 0:
                self._load_rng_state(resume_from_checkpoint)

            step = -1
            for step, inputs in enumerate(epoch_iterator):

                # Skip past any already trained steps if resuming training
                if steps_trained_in_current_epoch > 0:
                    steps_trained_in_current_epoch -= 1
                    if steps_trained_progress_bar is not None:
                        steps_trained_progress_bar.update(1)
                    if steps_trained_in_current_epoch == 0:
                        self._load_rng_state(resume_from_checkpoint)
                    continue
                elif steps_trained_progress_bar is not None:
                    steps_trained_progress_bar.close()
                    steps_trained_progress_bar = None

                if step % args.gradient_accumulation_steps == 0:
                    self.control = self.callback_handler.on_step_begin(args, self.state, self.control)

                if (
                    ((step + 1) % args.gradient_accumulation_steps != 0)
                    and args.local_rank != -1
                    and args._no_sync_in_gradient_accumulation
                ):
                    # Avoid unnecessary DDP synchronization since there will be no backward pass on this example.
                    with model.no_sync():
                        tr_loss_step = self.training_step(model, inputs)
                else:
                    tr_loss_step = self.training_step(model, inputs)

                if (
                    args.logging_nan_inf_filter
                    and not is_torch_tpu_available()
                    and (torch.isnan(tr_loss_step) or torch.isinf(tr_loss_step))
                ):
                    # if loss is nan or inf simply add the average of previous logged losses
                    tr_loss += tr_loss / (1 + self.state.global_step - self._globalstep_last_logged)
                else:
                    tr_loss += tr_loss_step

                self.current_flos += float(self.floating_point_ops(inputs))

                # Optimizer step for deepspeed must be called on every step regardless of the value of gradient_accumulation_steps
                if self.deepspeed:
                    self.deepspeed.step()

                if (step + 1) % args.gradient_accumulation_steps == 0 or (
                    # last step in epoch but step is always smaller than gradient_accumulation_steps
                    steps_in_epoch <= args.gradient_accumulation_steps
                    and (step + 1) == steps_in_epoch
                ):
                    # Gradient clipping
                    if args.max_grad_norm is not None and args.max_grad_norm > 0 and not self.deepspeed:
                        # deepspeed does its own clipping

                        if self.do_grad_scaling:
                            # Reduce gradients first for XLA
                            if is_torch_tpu_available():
                                gradients = xm._fetch_gradients(self.optimizer)
                                xm.all_reduce("sum", gradients, scale=1.0 / xm.xrt_world_size())
                            # AMP: gradients need unscaling
                            self.scaler.unscale_(self.optimizer)

                        if is_sagemaker_mp_enabled() and args.fp16:
                            self.optimizer.clip_master_grads(args.max_grad_norm)
                        elif hasattr(self.optimizer, "clip_grad_norm"):
                            # Some optimizers (like the sharded optimizer) have a specific way to do gradient clipping
                            self.optimizer.clip_grad_norm(args.max_grad_norm)
                        elif hasattr(model, "clip_grad_norm_"):
                            # Some models (like FullyShardedDDP) have a specific way to do gradient clipping
                            model.clip_grad_norm_(args.max_grad_norm)
                        else:
                            # Revert to normal clipping otherwise, handling Apex or full precision
                            nn.utils.clip_grad_norm_(
                                amp.master_params(self.optimizer) if self.use_apex else model.parameters(),
                                args.max_grad_norm,
                            )

                    # Optimizer step
                    optimizer_was_run = True
                    if self.deepspeed:
                        pass  # called outside the loop
                    elif is_torch_tpu_available():
                        if self.do_grad_scaling:
                            self.scaler.step(self.optimizer)
                            self.scaler.update()
                        else:
                            xm.optimizer_step(self.optimizer)
                    elif self.do_grad_scaling:
                        scale_before = self.scaler.get_scale()
                        self.scaler.step(self.optimizer)
                        self.scaler.update()
                        scale_after = self.scaler.get_scale()
                        optimizer_was_run = scale_before <= scale_after
                    else:
                        self.optimizer.step()

                    if optimizer_was_run and not self.deepspeed:
                        self.lr_scheduler.step()

                    model.zero_grad()
                    self.state.global_step += 1
                    self.state.epoch = epoch + (step + 1) / steps_in_epoch
                    self.control = self.callback_handler.on_step_end(args, self.state, self.control)

                    self._maybe_log_save_evaluate(tr_loss, model, trial, epoch, ignore_keys_for_eval)
                else:
                    self.control = self.callback_handler.on_substep_end(args, self.state, self.control)

                if self.control.should_epoch_stop or self.control.should_training_stop:
                    break
            if step < 0:
                logger.warning(
                    "There seems to be not a single sample in your epoch_iterator, stopping training at step"
                    f" {self.state.global_step}! This is expected if you're using an IterableDataset and set"
                    f" num_steps ({max_steps}) higher than the number of available samples."
                )
                self.control.should_training_stop = True

            self.control = self.callback_handler.on_epoch_end(args, self.state, self.control)
            self._maybe_log_save_evaluate(tr_loss, model, trial, epoch, ignore_keys_for_eval)

            if DebugOption.TPU_METRICS_DEBUG in self.args.debug:
                if is_torch_tpu_available():
                    # tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
                    xm.master_print(met.metrics_report())
                else:
                    logger.warning(
                        "You enabled PyTorch/XLA debug metrics but you don't have a TPU "
                        "configured. Check your training configuration if this is unexpected."
                    )
            if self.control.should_training_stop:
                break

        if args.past_index and hasattr(self, "_past"):
            # Clean the state at the end of training
            delattr(self, "_past")

        logger.info("\n\nTraining completed. Do not forget to share your model on huggingface.co/models =)\n\n")
        if args.load_best_model_at_end and self.state.best_model_checkpoint is not None:
            # Wait for everyone to get here so we are sur the model has been saved by process 0.
            if is_torch_tpu_available():
                xm.rendezvous("load_best_model_at_end")
            elif args.local_rank != -1:
                dist.barrier()
            elif is_sagemaker_mp_enabled():
                smp.barrier()

            self._load_best_model()

        # add remaining tr_loss
        self._total_loss_scalar += tr_loss.item()
        train_loss = self._total_loss_scalar / self.state.global_step

        metrics = speed_metrics("train", start_time, num_samples=num_train_samples, num_steps=self.state.max_steps)
        self.store_flos()
        metrics["total_flos"] = self.state.total_flos
        metrics["train_loss"] = train_loss

        self.is_in_train = False

        self._memory_tracker.stop_and_update_metrics(metrics)

        self.log(metrics)

        run_dir = self._get_output_dir(trial)
        checkpoints_sorted = self._sorted_checkpoints(use_mtime=False, output_dir=run_dir)

        # Delete the last checkpoint when save_total_limit=1 if it's different from the best checkpoint.
        if self.state.best_model_checkpoint is not None and self.args.save_total_limit == 1:
            if self.is_world_process_zero():
                for checkpoint in checkpoints_sorted:
                    if checkpoint != self.state.best_model_checkpoint:
                        logger.info(f"Deleting older checkpoint [{checkpoint}] due to args.save_total_limit")
                        shutil.rmtree(checkpoint)

        self.control = self.callback_handler.on_train_end(args, self.state, self.control)

        return TrainOutput(self.state.global_step, train_loss, metrics)

sdlm/trainer.py
allenai/tess-diffusion
"""Defines the utilities used during the training/infernece of diffusion language models."""
import torch.nn.functional as F
import os
import re
import pdb
from pathlib import Path
from transformers.utils import logging
import shutil
import numpy as np
from typing import Callable, Iterable, List
import torch

logger = logging.get_logger(__name__)


def join_texts(prefixes, sentences):
    """Joins prefixes to setences."""
    return [f"{prefix}{sentence}" for prefix, sentence in zip(prefixes, sentences)]


def convert_to_simplex(token_ids, simplex_value, vocab_size):
    return 2 * simplex_value * F.one_hot(token_ids, vocab_size) - simplex_value


def scale(inputs, scale_value):
    return inputs / scale_value


def get_last_checkpoint(folder, prefix_checkpoint_dir="step"):
    re_checkpoint = re.compile(r"^" + prefix_checkpoint_dir + r"\_(\d+)$")
    content = os.listdir(folder)
    checkpoints = [
        path for path in content if re_checkpoint.search(path) is not None and os.path.isdir(os.path.join(folder, path))
    ]
    if len(checkpoints) == 0:
        return
    return os.path.join(folder, max(checkpoints, key=lambda x: int(re_checkpoint.search(x).groups()[0])))


def remove_checkpoints(output_dir, checkpoint_prefix="step"):
    checkpoints = [str(x) for x in Path(output_dir).glob(f"{checkpoint_prefix}_*") if os.path.isdir(x)]
    for checkpoint in checkpoints:
        logger.info(f"Deleting older checkpoint [{checkpoint}] due to args.save_total_limit")
        shutil.rmtree(checkpoint)


def get_norm_stats(model):
    # Gradient norm of word embeddings and lm_head.
    input_embed_grad_norm = 0
    if model.roberta.embeddings.word_embeddings.weight.grad is not None:
        input_embed_grad_norm = model.roberta.embeddings.word_embeddings.weight.grad.detach().data.norm(2).item()

    output_embed_grad_norm = 0.0
    if model.lm_head.decoder.weight.grad is not None:
        output_embed_grad_norm = model.lm_head.decoder.weight.grad.detach().data.norm(2).item()

    """
    total_grad_norm = 0.0
    for p in model.parameters():
        grad_norm = 0.0
        if  p.grad is not None:
            grad_norm = p.grad.detach().data.norm(2).item()
        total_grad_norm += grad_norm ** 2
    total_grad_norm = total_grad_norm ** 0.5

    # Norms of word embeddings and lm_head.
    input_embed_norm = model.roberta.embeddings.word_embeddings.weight.detach().data.norm(2).item()
    output_embed_norm = model.lm_head.decoder.weight.detach().data.norm(2).item()
    total_param_norm = 0.0
    for p in model.parameters():
        param_norm = p.detach().data.norm(2)
        total_param_norm += param_norm.item() ** 2
    total_param_norm = total_param_norm ** 0.5
    """
    return {
        "input_embed_grad_norm": input_embed_grad_norm,
        "output_embed_grad_norm": output_embed_grad_norm,
        # "total_grad_norm": total_grad_norm,
        # "input_embed_norm": input_embed_norm,
        # "output_embed_norm": output_embed_norm,
        # "total_param_norm": total_param_norm
    }


def self_condition_preds(self_condition, logits, logits_projection=None):
    if self_condition in ["logits", "logits_addition", "logits_mean", "logits_max", "logits_multiply"]:
        previous_pred = logits.detach()
    elif self_condition in ["logits_with_projection", "logits_with_projection_addition"]:
        previous_pred = logits_projection(logits.detach())
    else:
        assert NotImplementedError(f"{self_condition} is not implemented.")
    return previous_pred

def mix_values_based_on_self_condition(self_condition_type, value_1, value_2):
    if self_condition_type in ["logits_with_projection_addition", "logits_addition"]:
        mixed_values = value_1 + value_2
    elif self_condition_type == "logits_mean":
        mixed_values = (value_1 + value_2) / 2.0
    elif self_condition_type == "logits_max":
        mixed_values = torch.max(value_1, value_2)
    elif self_condition_type == "logits_multiply":
        mixed_values = value_1 * value_2
    else:
        assert NotImplementedError
    return mixed_values

def round_stsb_target(label):
    """STSB maps two sentences to a floating point number between 1 and 5
    representing their semantic similarity. Since we are treating all tasks as
    text-to-text tasks we need to convert this floating point number to a string.
    The vast majority of the similarity score labels in STSB are in the set
    [0, 0.2, 0.4, ..., 4.8, 5.0]. So, we first round the number to the closest
    entry in this set, and then we convert the result to a string (literally e.g.
    "3.4"). This converts STSB roughly into a 26-class classification dataset.
    Args:
      label: original label.
    Returns:
      A preprocessed label.
    """
    return np.round((label * 5) / 5, decimals=1)


def lmap(f: Callable, x: Iterable) -> List:
    """list(map(f, x))"""
    return list(map(f, x))


def pad_data(data_list, tokenizer):
    return tokenizer.pad({"input_ids": data_list}, padding=True)["input_ids"]

sdlm/utils.py
allenai/tess-diffusion
"""Setups diffusion for text."""

import setuptools


def setup_package():
    setuptools.setup(
        name="sdlm",
        version="0.0.1",
        description="Simplex diffusion for language models.",
        long_description=open("README.md").read(),
        long_description_content_type="text/markdown",
        classifiers=[
            "Intended Audience :: Science/Research",
            "Development Status :: 3 - Alpha",
            "License :: OSI Approved :: Apache Software License",
            "Programming Language :: Python :: 3",
            "Programming Language :: Python :: 3.7",
            "Topic :: Scientific/Engineering :: Artificial Intelligence",
        ],
        keywords="",
        url="https://github.com/allenai/tess_diffusion",
        author="Allen Institute for Artificial Intelligence",
        author_email="rabeehk@allenai.org",
        license="Apache",
        packages=setuptools.find_packages(
            exclude=["tests", "docs", "scripts", "examples"],
        ),
        dependency_links=[
            "https://download.pytorch.org/whl/torch_stable.html",
        ],
        python_requires=">=3.7",
    )


if __name__ == "__main__":
    setup_package()

setup.py
allenai/tess-diffusion