WIP templating

examples/seq2seq/aggregate_distributed_results.py

@@ -0,0 +1,46 @@
+from pathlib import Path
+
+import fire
+
+
+try:
+    from .utils import calculate_bleu, calculate_rouge, load_json, save_json, write_txt_file
+except ImportError:
+    from utils import calculate_bleu, calculate_rouge, load_json, save_json, write_txt_file
+
+
+def combine_partial_results(
+    result_dir: str, save_dir: str = None, save_prefix=None, calc_bleu=False, just_metrics=False
+):
+    """Write first n lines of each file f in src_dir to dest_dir/f """
+    src_dir = Path(result_dir)
+    save_dir = Path(save_dir)
+    save_dir.mkdir(exist_ok=True)
+    paths_to_combine = list(src_dir.glob("rank*.json"))
+    records = []
+    for partial_result in paths_to_combine:
+        records.extend(load_json(partial_result))
+    preds = [x["pred"] for x in records]
+    labels = [x["label"] for x in records]
+    score_fn = calculate_bleu if calc_bleu else calculate_rouge
+    metrics = score_fn(preds, labels)
+    save_json(metrics, save_dir.joinpath("metrics.json"))  # better would be be {prefix}_{rouge|bleu}.json
+    print(metrics)
+    if just_metrics:
+        return
+
+    if save_prefix is None:
+        save_prefix = "generated"
+        print("using generated as prefix")
+
+    tgt_path = save_dir.joinpath(f"{save_prefix}.target")
+    write_txt_file(labels, tgt_path)
+    pred_path = save_dir.joinpath(f"{save_prefix}.pred_target")
+    write_txt_file(preds, pred_path)
+    if "source" in records[0]:
+        src_path = save_dir.joinpath(f"{save_prefix}.source")
+        write_txt_file([x["source"] for x in records], src_path)
+
+
+if __name__ == "__main__":
+    fire.Fire(combine_partial_results)

examples/seq2seq/distillation.py

@@ -472,6 +472,7 @@ LAYERS_TO_COPY = {
         6: [0, 3, 6, 9, 12, 15],
         8: [0, 2, 4, 6, 8, 10, 12, 15],
         9: [0, 1, 3, 5, 7, 9, 11, 13, 15],
+        12: [0, 1, 2, 3, 4, 5, 6, 7, 9, 11, 13, 15],
         16: list(range(16)),
     },
     6: {1: [0], 2: [0, 5], 3: [0, 2, 5], 4: [0, 1, 3, 5], 6: list(range(6))},

examples/seq2seq/run_distributed_eval.py

@@ -0,0 +1,142 @@
+import argparse
+from logging import getLogger
+from pathlib import Path
+from typing import Dict
+
+import torch
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+
+
+logger = getLogger(__name__)
+
+try:
+    from .utils import Seq2SeqDataset, parse_numeric_cl_kwargs, save_json, use_task_specific_params
+except ImportError:
+    from utils import Seq2SeqDataset, parse_numeric_cl_kwargs, save_json, use_task_specific_params
+
+
+DEFAULT_DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+def eval_data_dir(
+    data_dir,
+    save_dir: str,
+    model_name: str,
+    bs: int = 8,
+    max_source_length: int = 1024,
+    type_path="val",
+    n_obs=None,
+    fp16=False,
+    save_source=False,
+    num_beams: int = 4,
+    task="summarization",
+    local_rank=None,
+    **generate_kwargs,
+) -> Dict:
+    """Run evaluation on part of the data for one gpu and save to {save_dir}/rank_{rank}_output.json"""
+    model_name = str(model_name)
+    assert local_rank is not None
+    torch.distributed.init_process_group(backend="nccl", rank=local_rank)
+
+    save_dir = Path(save_dir)
+    save_path = save_dir.joinpath(f"rank_{local_rank}_output.json")
+    torch.cuda.set_device(local_rank)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name).cuda()
+    if fp16:
+        model = model.half()
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    logger.info(f"Inferred tokenizer type: {tokenizer.__class__}")  # if this is wrong, check config.model_type.
+    use_task_specific_params(model, task)  # update config with task specific params
+    if max_source_length is None:
+        max_source_length = tokenizer.model_max_length
+    ds = Seq2SeqDataset(
+        tokenizer,
+        data_dir,
+        max_source_length,
+        max_target_length=1024,
+        type_path=type_path,
+        n_obs=n_obs,
+        prefix=model.config.prefix,
+    )
+    sampler = ds.make_sortish_sampler(bs, distributed=True)
+    data_loader = DataLoader(ds, sampler=sampler, batch_size=bs, collate_fn=ds.collate_fn)
+    dec_kwargs = dict(skip_special_tokens=True, clean_up_tokenization_spaces=False)  # tokenizer.decode
+    results = []
+    for batch in tqdm(data_loader):
+        summaries = model.generate(
+            input_ids=batch["input_ids"].to(model.device),
+            attention_mask=batch["attention_mask"].to(model.device),
+            num_beams=num_beams,
+            **generate_kwargs,
+        )
+        preds = tokenizer.batch_decode(summaries, **dec_kwargs)
+        labels = tokenizer.batch_decode(batch["labels"], **dec_kwargs)
+        if save_source:
+            docs = tokenizer.batch_decode(batch["input_ids"], **dec_kwargs)
+        for i in range(len(labels)):
+            label, pred = labels[i], preds[i]
+            if save_source:
+                results.append(dict(pred=pred, label=label, source=docs[i]))
+            else:
+                results.append(dict(pred=pred, label=label))
+    save_json(results, save_path)
+    return results
+
+
+def run_generate():
+    parser = argparse.ArgumentParser(
+        epilog="Unspecified args like --num_beams=2 --decoder_start_token_id=4 are passed to model.generate"
+    )
+    parser.add_argument("--input_path", type=str, help="like cnn_dm/test.source")
+    parser.add_argument(
+        "--model_name",
+        type=str,
+        help="like facebook/bart-large-cnn,t5-base, etc.",
+        default="sshleifer/distilbart-xsum-12-3",
+    )
+    parser.add_argument("--save_dir", type=str, help="where to save", default="tmp_gen")
+    parser.add_argument("--max_source_length", type=int, default=None)
+    parser.add_argument(
+        "--type_path", type=str, default="test", help="which subset to evaluate typically train/val/test"
+    )
+    parser.add_argument("--reference_path", type=str, required=False, help="like cnn_dm/test.target")
+    parser.add_argument("--task", type=str, default="summarization", help="used for task_specific_params + metrics")
+    parser.add_argument("--bs", type=int, default=8, required=False, help="batch size")
+    parser.add_argument(
+        "--local_rank", type=int, default=-1, required=False, help="should be passed by distributed.launch"
+    )
+
+    parser.add_argument(
+        "--n_obs", type=int, default=None, required=False, help="How many observations. Defaults to all."
+    )
+    parser.add_argument("--fp16", action="store_true")
+    parser.add_argument("--save_source", action="store_true")
+
+    args, rest = parser.parse_known_args()
+    generate_kwargs = parse_numeric_cl_kwargs(rest)
+    if generate_kwargs:
+        print(f"parsed the following generate kwargs: {generate_kwargs}")
+    Path(args.save_dir).mkdir(exist_ok=True)
+    eval_data_dir(
+        args.input_path,
+        args.save_dir,
+        args.model_name,
+        type_path=args.type_path,
+        batch_size=args.bs,
+        fp16=args.fp16,
+        task=args.task,
+        local_rank=args.local_rank,
+        n_obs=args.n_obs,
+        save_source=args.save_source,
+        max_source_length=args.max_source_length,
+        **generate_kwargs,
+    )
+
+
+if __name__ == "__main__":
+    # Usage for MT:
+    run_generate()

examples/seq2seq/utils.py

@@ -98,7 +98,8 @@ class AbstractSeq2SeqDataset(Dataset):
         self.max_target_length = max_target_length
         assert min(self.src_lens) > 0, f"found empty line in {self.src_file}"
         self.tokenizer = tokenizer
-        self.prefix = prefix
+        self.prefix = prefix if prefix is not None else ""
+
         if n_obs is not None:
             self.src_lens = self.src_lens[:n_obs]
         self.pad_token_id = self.tokenizer.pad_token_id
@@ -387,3 +388,10 @@ def parse_numeric_cl_kwargs(unparsed_args: List[str]) -> Dict[str, Union[int, fl
 
         result[unparsed_args[i][2:]] = value
     return result
+
+
+def write_txt_file(ordered_tgt, path):
+    f = Path(path).open("w")
+    for ln in ordered_tgt:
+        f.write(ln + "\n")
+        f.flush()

src/transformers/configuration_pegasus.py

@@ -49,7 +49,7 @@ DEFAULTS = dict(
 # Config values that vary between checkpoints: for testing and conversion
 task_specific_params = {
     # These are task specific params for pegasus-large and normal params for finetuned checkpoints
-    "summarization_xsum": {"length_penalty": 0.8, "max_length": 64, "max_position_embeddings": 512},
+    "summarization_xsum": {"length_penalty": 0.6, "max_length": 64, "max_position_embeddings": 512},
     "summarization_cnn_dailymail": {"length_penalty": 0.8, "max_length": 128, "max_position_embeddings": 1024},
     "summarization_newsroom": {"length_penalty": 0.8, "max_length": 128, "max_position_embeddings": 512},
     "summarization_wikihow": {"length_penalty": 0.6, "max_length": 256, "max_position_embeddings": 512},

src/transformers/convert_pegasus_tf_to_pytorch.py

@@ -47,8 +47,8 @@ PATTERNS = [
 
 def rename_state_dict_key(k):
 
-    for pegasus_name, bart_name in PATTERNS:
-        k = k.replace(pegasus_name, bart_name)
+    for pegasus_name, hf_name in PATTERNS:
+        k = k.replace(pegasus_name, hf_name)
     return k
 
 
@@ -57,13 +57,12 @@ def rename_state_dict_key(k):
 # TODO(SS): one constant
 
 
-def convert_pegasus_to_bart(tf_weights: dict, cfg_updates: dict) -> PegasusForConditionalGeneration:
+def convert_pegasus(tf_weights: dict, cfg_updates: dict) -> PegasusForConditionalGeneration:
     cfg_kwargs = DEFAULTS.copy()
     cfg_kwargs.update(cfg_updates)
-
-    cfg = PegasusConfig(**cfg_updates)
-    bart = PegasusForConditionalGeneration(cfg)
-    sd = bart.model.state_dict()
+    cfg = PegasusConfig(**cfg_kwargs)
+    torch_model = PegasusForConditionalGeneration(cfg)
+    sd = torch_model.model.state_dict()
     mapping = {}
     for k, v in tf_weights.items():
         new_k = rename_state_dict_key(k)
@@ -80,13 +79,13 @@ def convert_pegasus_to_bart(tf_weights: dict, cfg_updates: dict) -> PegasusForCo
     mapping["decoder.embed_tokens.weight"] = mapping["shared.weight"]
     empty_biases = {k: torch.zeros_like(v) for k, v in sd.items() if k.endswith("bias") and k not in mapping}
     mapping.update(**empty_biases)
-    missing, extra = bart.model.load_state_dict(mapping, strict=False)
+    missing, extra = torch_model.model.load_state_dict(mapping, strict=False)
     unexpected_missing = [
         k for k in missing if k not in ["encoder.embed_positions.weight", "decoder.embed_positions.weight"]
     ]
     assert unexpected_missing == [], f"no matches found for the following torch keys {unexpected_missing}"
     assert extra == [], f"no matches found for the following tf keys {extra}"
-    return bart
+    return torch_model
 
 
 def get_tf_weights_as_numpy(path="./ckpt/aeslc/model.ckpt-32000") -> Dict:
@@ -115,7 +114,7 @@ def convert_pegasus_ckpt_to_pytorch(ckpt_path: str, save_dir: str):
     cfg_updates = task_specific_params[f"summarization_{dataset}"]
     if dataset == "large":
         cfg_updates["task_specific_params"] = task_specific_params
-    torch_model = convert_pegasus_to_bart(tf_weights, cfg_updates)
+    torch_model = convert_pegasus(tf_weights, cfg_updates)
     torch_model.save_pretrained(save_dir)
     sd = torch_model.state_dict()
     sd.pop("model.decoder.embed_positions.weight")

templates/examples/type/scripts/explicit_example.py

@@ -0,0 +1,63 @@
+import torch
+from torch.utils.data import DataLoader
+from datasets import load_dataset, load_metric
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, AdamW, logging
+
+logger = logging.get_logger(__name__)
+
+model_name = 'bert-base-cased'
+dataset_name, dataset_subset = 'glue', 'mrpc'
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
+model = AutoModelForSequenceClassification.from_pretrained(model_name)
+dataset = load_dataset(dataset_name, dataset_subset)
+metric = load_metric(dataset_name, dataset_subset)
+
+
+# Prepare the dataset: encode it as input to our model and truncate to max model length
+def encode(examples):
+    output = tokenizer(examples['sentence1'], examples['sentence2'], truncation=True)
+    output['labels'] = examples['label']
+    return output
+
+encoded_dataset = dataset.map(encode, batched=True)
+encoded_dataset.set_format(columns=['input_ids', 'token_type_ids', 'labels'])
+
+def train_epoch(model, optimizer, train_dataloader, eval_dataloader, device):
+    model.to(device)
+    model.train()
+    for batch in train_dataloader:
+        batch.to(device)
+        model(**batch)
+        loss = model.losss
+        loss.backward()
+        optimizer.step()
+
+    model.eval()
+    for batch in eval_dataloader:
+        batch.to(device)
+        outputs = model(**batch)
+        predictions = outputs.logits.argmax(dim=-1)
+        metric.add_batch(predictions=predictions, references=batch['labels'])
+    scores = metric.compute()
+    return scores
+
+learning_rate = 1e-5
+train_batch_size = 16
+seed = 42
+epochs = 2
+
+# Prepare pytorch dataloader (with dynamic batch i.e. pad the sequences to the longest in the batch)
+def collate_fn(examples):
+  return tokenizer.pad(examples, padding='longest', return_tensors='pt')
+
+train_dataloader = DataLoader(dataset['train'], shuffle=True, collate_fn=collate_fn, batch_size=32)
+valid_dataloader = DataLoader(dataset['validation'], shuffle=False, collate_fn=collate_fn, batch_size=16)
+
+optimizer = AdamW(params=model.parameters(), lr=learning_rate)
+
+for epoch in range(epochs):
+    eval_metrics = train_epoch(model, optimizer, train_dataloader, valid_dataloader)
+    print(f"eval_metrics for trial {trial.number} epoch {epoch}:", eval_metrics)
+    sum_metrics = sum(eval_metrics.values())

templates/examples/type/scripts/{{cookiecutter.application}}_{{cookiecutter.style}}.py

@@ -0,0 +1,127 @@
+from transformers import AutoTokenizer, AdamW, get_constant_schedule_with_warmup, get_linear_schedule_with_warmup, {{cookiecutter.model_class_name}}
+from datasets import load_dataset, load_metric, temp_seed
+from torch.utils.data import DataLoader
+import argparse
+import os
+
+MAX_GPU_BATCH_SIZE = 8
+EVAL_BATCH_SIZE = 16
+
+def prepare_dataset(args):
+    """ Load a dataset, a tokenizer and encode the dataset """
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
+
+    dataset = load_dataset(args.dataset_name, args.dataset_config_name)
+    # If you have a custom dataset in a CSV/JSON/TEXT file
+    # dataset = load_dataset("csv", data_files={'train': "train_data.csv", 'validation': "valid_data.csv"})
+
+    # Prepare the dataset (cached)
+    def encode(example):
+        output = tokenizer(example['sentence1'], example['sentence2'], truncation=True)
+        output['labels'] = example['label']
+        return output
+
+    dataset = dataset.map(encode)
+    dataset.set_format(columns=['attention_mask', 'input_ids', 'token_type_ids', 'labels'])
+
+    return dataset, tokenizer
+
+
+def train_and_evaluate(args):
+    # Load dataset, tokenizer and metric
+    dataset, tokenizer = prepare_dataset(args)
+    metric = load_metric(args.dataset_name, args.dataset_config_name)
+
+    # Instantiate train and evaluation dataloaders
+    def collate_fn(examples):
+        """ Collate function which pads the batches to the longest sequence in the batch and returns PyTorch tensors """
+        return tokenizer.pad(examples, padding='longest', return_tensors='pt')
+
+    train_dataloader = DataLoader(dataset['train'], shuffle=True, collate_fn=collate_fn, batch_size=args.train_batch_size)
+    eval_dataloader = DataLoader(dataset['validation'], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE)
+
+    # Instantiate model
+    model = {{cookiecutter.model_class_name}}.from_pretrained(args.model_name_or_path, return_dict=True)
+
+    # Instantiate optimizer and learning rate schedule
+    optimizer = AdamW(params=model.parameters(), lr=args.lr, correct_bias=args.correct_bias)
+    if args.lr_schedule == 'constant':
+        lr_scheduler = get_constant_schedule_with_warmup(optimizer=optimizer, num_warmup_steps=100)
+    else:
+        lr_scheduler = get_linear_schedule_with_warmup(optimizer=optimizer, num_warmup_steps=100, num_training_steps=len(train_dataloader)*args.num_epochs)
+
+    # Train the model
+    eval_metrics = {}
+    for epoch in range(args.num_epochs):
+        model.train()
+        model.to(args.device)
+        for step, batch in enumerate(train_dataloader):
+            batch.to(args.device)
+            outputs = model(**batch)
+            loss = outputs.loss
+            loss = loss / args.gradient_accumulation_steps
+            loss.backward()
+            if step % args.gradient_accumulation_steps == 0:
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
+
+        model.eval()
+        for step, batch in enumerate(eval_dataloader):
+            batch.to(args.device)
+            outputs = model(**batch)
+            predictions = outputs.logits.argmax(dim=-1)
+            metric.add_batch(predictions=predictions, references=batch['labels'])
+
+        eval_metrics = metric.compute()
+        print(f"eval_metric for epoch {epoch}:", eval_metrics)
+
+    return eval_metrics
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_name_or_path", default=None, type=str, required=True, help="Path to pre-trained model or shortcut name")
+    parser.add_argument("--config_name", default="", type=str, help="Pretrained config name or path if not the same as model_name")
+    parser.add_argument("--tokenizer_name", default="", type=str, help="Pretrained tokenizer name or path if not the same as model_name")
+    parser.add_argument("--task_name", default=None, type=str, required=True, help="The name of the task to train.")
+    parser.add_argument("--max_seq_length", default=128, type=int, help="The maximum total input sequence length after tokenization. Sequences longer than this will be truncated, sequences shorter will be padded.")
+    parser.add_argument("--output_dir", default=None, type=str, required=True, help="The output directory where the model predictions and checkpoints will be written.")
+    parser.add_argument("--overwrite_output_dir", default=False, type=bool, help="Overwrite the content of the output directory. Use this to continue training if output_dir points to a checkpoint directory.")
+    parser.add_argument("--do_train", action='store_true', help="Whether to run training.")
+    parser.add_argument("--do_eval", action='store_true', help="Whether to run eval on the dev set.")
+    parser.add_argument("--evaluate_during_training", action='store_true', help="Run evaluation during training at each logging step.")
+    parser.add_argument("--prediction_loss_only", action='store_true', help="When performing evaluation and predictions, only returns the loss.")
+    parser.add_argument("--per_device_train_batch_size", default=8, type=int, help="Batch size per GPU/TPU core/CPU for training.")
+    parser.add_argument("--per_device_eval_batch_size", default=8, type=int, help="Batch size per GPU/TPU core/CPU for evaluation.")
+    parser.add_argument('--gradient_accumulation_steps', type=int, default=1, help="Number of updates steps to accumulate before performing a backward/update pass.")     
+    parser.add_argument("--learning_rate", default=5e-5, type=float, help="The initial learning rate for Adam.")
+    parser.add_argument("--weight_decay", default=0.0, type=float, help="Weight decay if we apply some.")
+    parser.add_argument("--adam_beta1", default=0.9, type=float, help="Beta1 for Adam optimizer.")
+    parser.add_argument("--adam_beta2", default=0.999, type=float, help="Beta2 for Adam optimizer.")
+    parser.add_argument("--adam_epsilon", default=1e-8, type=float, help="Epsilon for Adam optimizer.")
+    parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
+    parser.add_argument("--num_train_epochs", default=3.0, type=float, help="Total number of training epochs to perform.")
+    parser.add_argument("--max_steps", default=-1, type=int, help="If > 0: set total number of training steps to perform. Override num_train_epochs.")
+    parser.add_argument("--warmup_steps", default=0, type=int, help="Linear warmup over warmup_steps.")
+    parser.add_argument("--logging_dir", type=str, default_factory=default_logdir, help="Tensorboard log dir.")
+    parser.add_argument("--logging_first_step", action='store_true', help="Log and eval the first global_step")
+    parser.add_argument('--logging_steps', type=int, default=500, help="Log every X updates steps.")
+    parser.add_argument('--save_steps', type=int, default=500, help="Save checkpoint every X updates steps.")
+    parser.add_argument('--save_total_limit', type=int, default=None, help="Limit the total amount of checkpoints. Deletes the older checkpoints in the output_dir. Default is unlimited checkpoints")
+    parser.add_argument('--seed', type=int, default=42, help="random seed for initialization")
+    parser.add_argument("--no_cuda", action='store_true', help="Avoid using CUDA when available")
+    parser.add_argument('--fp16', action='store_true', help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit")
+    parser.add_argument('--fp16_opt_level', type=str, default='O1', help="For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']. See details at https://nvidia.github.io/apex/amp.html")
+    parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank")
+    parser.add_argument("--dataloader_drop_last", action='store_true', help="Drop the last incomplete batch if it is not divisible by the batch size.")
+    parser.add_argument("--eval_steps", type=int, default=1000, help="Run an evaluation every X steps.")
+    parser.add_argument("--past_index", type=int, default=-1, help="If >=0, uses the corresponding part of the output as the past state for next step.")
+    parser.add_argument("--run_name", type=str, default=None, help="An optional descriptor for the run. Notably used for wandb logging.")
+    parser.add_argument("--disable_tqdm", action='store_true', help="Whether or not to disable the tqdm progress bars.")
+    args = parser.parse_args()
+
+    with temp_seed(args.seed, set_pytorch=True):
+        return train_and_evaluate(args)
+
+if __name__ == "__main__":
+    main()