Add files via upload

training/run_retriever_no_trainer.py

@@ -0,0 +1,381 @@
+import argparse
+import functools
+import logging
+import math
+from random import choice, randint
+
+import torch
+from accelerate import Accelerator
+from accelerate.utils import set_seed
+from datasets import load_dataset
+from torch.utils import checkpoint
+from torch.utils.data import Dataset, RandomSampler, DataLoader, SequentialSampler
+from tqdm.auto import tqdm
+from transformers import get_scheduler, AutoTokenizer, AdamW, SchedulerType, AutoModelForSequenceClassification
+
+logger = logging.getLogger(__name__)
+
+
+def get_parser():
+    parser = argparse.ArgumentParser(description="Train ELI5 retriever")
+    parser.add_argument(
+        "--dataset_name",
+        type=str,
+        default="vblagoje/lfqa",
+        help="The name of the dataset to use (via the datasets library).",
+    )
+
+    parser.add_argument(
+        "--per_device_train_batch_size",
+        type=int,
+        default=1024,
+    )
+
+    parser.add_argument(
+        "--per_device_eval_batch_size",
+        type=int,
+        default=1024,
+        help="Batch size (per device) for the evaluation dataloader.",
+    )
+
+    parser.add_argument(
+        "--max_length",
+        type=int,
+        default=128,
+    )
+
+    parser.add_argument(
+        "--checkpoint_batch_size",
+        type=int,
+        default=32,
+    )
+
+    parser.add_argument(
+        "--pretrained_model_name",
+        type=str,
+        default="google/bert_uncased_L-8_H-768_A-12",
+    )
+
+    parser.add_argument(
+        "--model_save_name",
+        type=str,
+        default="eli5_retriever_model_l-12_h-768_b-512-512",
+    )
+
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=2e-4,
+    )
+
+    parser.add_argument(
+        "--weight_decay",
+        type=float,
+        default=0.2,
+    )
+
+    parser.add_argument(
+        "--log_freq",
+        type=int,
+        default=500,
+        help="Log train/validation loss every log_freq update steps"
+    )
+
+    parser.add_argument(
+        "--num_train_epochs",
+        type=int,
+        default=4,
+    )
+
+    parser.add_argument(
+        "--max_train_steps",
+        type=int,
+        default=None,
+        help="Total number of training steps to perform. If provided, overrides num_train_epochs.",
+    )
+
+    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(
+        "--lr_scheduler_type",
+        type=SchedulerType,
+        default="linear",  # this is linear with warmup
+        help="The scheduler type to use.",
+        choices=["linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup"],
+    )
+
+    parser.add_argument(
+        "--num_warmup_steps",
+        type=int,
+        default=100,
+        help="Number of steps for the warmup in the lr scheduler."
+    )
+
+    parser.add_argument(
+        "--warmup_percentage",
+        type=float,
+        default=0.08,
+        help="Number of steps for the warmup in the lr scheduler."
+    )
+    return parser
+
+
+class RetrievalQAEmbedder(torch.nn.Module):
+    def __init__(self, sent_encoder):
+        super(RetrievalQAEmbedder, self).__init__()
+        dim = sent_encoder.config.hidden_size
+        self.bert_query = sent_encoder
+        self.output_dim = 128
+        self.project_query = torch.nn.Linear(dim, self.output_dim, bias=False)
+        self.project_doc = torch.nn.Linear(dim, self.output_dim, bias=False)
+        self.ce_loss = torch.nn.CrossEntropyLoss(reduction="mean")
+
+    def embed_sentences_checkpointed(self, input_ids, attention_mask, checkpoint_batch_size=-1):
+        # reproduces BERT forward pass with checkpointing
+        if checkpoint_batch_size < 0 or input_ids.shape[0] < checkpoint_batch_size:
+            return self.bert_query(input_ids, attention_mask=attention_mask)[1]
+        else:
+            # prepare implicit variables
+            device = input_ids.device
+            input_shape = input_ids.size()
+            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
+            head_mask = [None] * self.bert_query.config.num_hidden_layers
+            extended_attention_mask: torch.Tensor = self.bert_query.get_extended_attention_mask(
+                attention_mask, input_shape, device
+            )
+
+            # define function for checkpointing
+            def partial_encode(*inputs):
+                encoder_outputs = self.bert_query.encoder(inputs[0], attention_mask=inputs[1], head_mask=head_mask, )
+                sequence_output = encoder_outputs[0]
+                pooled_output = self.bert_query.pooler(sequence_output)
+                return pooled_output
+
+            # run embedding layer on everything at once
+            embedding_output = self.bert_query.embeddings(
+                input_ids=input_ids, position_ids=None, token_type_ids=token_type_ids, inputs_embeds=None
+            )
+            # run encoding and pooling on one mini-batch at a time
+            pooled_output_list = []
+            for b in range(math.ceil(input_ids.shape[0] / checkpoint_batch_size)):
+                b_embedding_output = embedding_output[b * checkpoint_batch_size: (b + 1) * checkpoint_batch_size]
+                b_attention_mask = extended_attention_mask[b * checkpoint_batch_size: (b + 1) * checkpoint_batch_size]
+                pooled_output = checkpoint.checkpoint(partial_encode, b_embedding_output, b_attention_mask)
+                pooled_output_list.append(pooled_output)
+            return torch.cat(pooled_output_list, dim=0)
+
+    def embed_questions(self, q_ids, q_mask, checkpoint_batch_size=-1):
+        q_reps = self.embed_sentences_checkpointed(q_ids, q_mask, checkpoint_batch_size)
+        return self.project_query(q_reps)
+
+    def embed_answers(self, a_ids, a_mask, checkpoint_batch_size=-1):
+        a_reps = self.embed_sentences_checkpointed(a_ids, a_mask, checkpoint_batch_size)
+        return self.project_doc(a_reps)
+
+    def forward(self, q_ids, q_mask, a_ids, a_mask, checkpoint_batch_size=-1):
+        device = q_ids.device
+        q_reps = self.embed_questions(q_ids, q_mask, checkpoint_batch_size)
+        a_reps = self.embed_answers(a_ids, a_mask, checkpoint_batch_size)
+        compare_scores = torch.mm(q_reps, a_reps.t())
+        loss_qa = self.ce_loss(compare_scores, torch.arange(compare_scores.shape[1]).to(device))
+        loss_aq = self.ce_loss(compare_scores.t(), torch.arange(compare_scores.shape[0]).to(device))
+        loss = (loss_qa + loss_aq) / 2
+        return loss
+
+
+class ELI5DatasetQARetriever(Dataset):
+    def __init__(self, examples_array, extra_answer_threshold=3, min_answer_length=64, training=True, n_samples=None):
+        self.data = examples_array
+        self.answer_thres = extra_answer_threshold
+        self.min_length = min_answer_length
+        self.training = training
+        self.n_samples = self.data.num_rows if n_samples is None else n_samples
+
+    def __len__(self):
+        return self.n_samples
+
+    def make_example(self, idx):
+        example = self.data[idx]
+        question = example["title"]
+        if self.training:
+            answers = [a for i, (a, sc) in enumerate(zip(example["answers"]["text"], example["answers"]["score"]))]
+            answer_tab = choice(answers).split(" ")
+            start_idx = randint(0, max(0, len(answer_tab) - self.min_length))
+            answer_span = " ".join(answer_tab[start_idx:])
+        else:
+            answer_span = example["answers"]["text"][0]
+        return question, answer_span
+
+    def __getitem__(self, idx):
+        return self.make_example(idx % self.data.num_rows)
+
+
+def make_qa_retriever_batch(qa_list, tokenizer, max_len=64):
+    q_ls = [q for q, a in qa_list]
+    a_ls = [a for q, a in qa_list]
+    q_toks = tokenizer(q_ls, padding="max_length", max_length=max_len, truncation=True)
+    q_ids, q_mask = (
+        torch.LongTensor(q_toks["input_ids"]),
+        torch.LongTensor(q_toks["attention_mask"])
+    )
+    a_toks = tokenizer(a_ls, padding="max_length", max_length=max_len, truncation=True)
+    a_ids, a_mask = (
+        torch.LongTensor(a_toks["input_ids"]),
+        torch.LongTensor(a_toks["attention_mask"]),
+    )
+    return q_ids, q_mask, a_ids, a_mask
+
+
+def evaluate_qa_retriever(model, data_loader):
+    # make iterator
+    epoch_iterator = tqdm(data_loader, desc="Iteration", disable=True)
+    tot_loss = 0.0
+    with torch.no_grad():
+        for step, batch in enumerate(epoch_iterator):
+            q_ids, q_mask, a_ids, a_mask = batch
+            loss = model(q_ids, q_mask, a_ids, a_mask)
+            tot_loss += loss.item()
+        return tot_loss / (step + 1)
+
+
+def train(config):
+    set_seed(42)
+    args = config["args"]
+    data_files = {"train": "train.json", "validation": "validation.json", "test": "test.json"}
+    eli5 = load_dataset(args.dataset_name, data_files=data_files)
+
+    # Initialize the accelerator. We will let the accelerator handle device placement for us in this example.
+    accelerator = Accelerator()
+    # Make one log on every process with the configuration for debugging.
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        level=logging.INFO,
+    )
+    logger.setLevel(logging.INFO if accelerator.is_local_main_process else logging.ERROR)
+    logger.info(accelerator.state)
+
+    # prepare torch Dataset objects
+    train_dataset = ELI5DatasetQARetriever(eli5['train'], training=True)
+    valid_dataset = ELI5DatasetQARetriever(eli5['validation'], training=False)
+
+    tokenizer = AutoTokenizer.from_pretrained(args.pretrained_model_name)
+    base_model = AutoModel.from_pretrained(args.pretrained_model_name)
+
+    model = RetrievalQAEmbedder(base_model)
+    no_decay = ['bias', 'LayerNorm.weight']
+    optimizer_grouped_parameters = [
+        {'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
+         'weight_decay': args.weight_decay},
+        {'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
+    ]
+    optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, weight_decay=args.weight_decay)
+
+    model_collate_fn = functools.partial(make_qa_retriever_batch, tokenizer=tokenizer, max_len=args.max_length)
+    train_dataloader = DataLoader(train_dataset, batch_size=args.per_device_train_batch_size,
+                                  sampler=RandomSampler(train_dataset), collate_fn=model_collate_fn)
+
+    model_collate_fn = functools.partial(make_qa_retriever_batch, tokenizer=tokenizer, max_len=args.max_length)
+    eval_dataloader = DataLoader(valid_dataset, batch_size=args.per_device_eval_batch_size,
+                                 sampler=SequentialSampler(valid_dataset), collate_fn=model_collate_fn)
+
+    # train the model
+    model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(model, optimizer,
+                                                                              train_dataloader, eval_dataloader)
+    # Scheduler and math around the number of training steps.
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+    if args.max_train_steps is None:
+        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
+    else:
+        args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+
+    num_warmup_steps = args.num_warmup_steps if args.num_warmup_steps else math.ceil(args.max_train_steps *
+                                                                                     args.warmup_percentage)
+    scheduler = get_scheduler(
+        name=args.lr_scheduler_type,
+        optimizer=optimizer,
+        num_warmup_steps=args.num_warmup_steps,
+        num_training_steps=args.max_train_steps,
+    )
+
+    # Train!
+    total_batch_size = args.per_device_train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+
+    logger.info("***** Running training *****")
+    logger.info(f"  Num examples = {len(train_dataset)}")
+    logger.info(f"  Num Epochs = {args.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_batch_size}")
+    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
+    logger.info(f"  Total optimization steps = {args.max_train_steps}")
+    logger.info(f"  Warmup steps = {num_warmup_steps}")
+    logger.info(f"  Logging training progress every {args.log_freq} optimization steps")
+
+    loc_loss = 0.0
+    current_loss = 0.0
+    checkpoint_step = 0
+
+    completed_steps = checkpoint_step
+    progress_bar = tqdm(range(args.max_train_steps), initial=checkpoint_step,
+                        disable=not accelerator.is_local_main_process)
+    for epoch in range(args.num_train_epochs):
+        model.train()
+        batch = next(iter(train_dataloader))
+        for step in range(1000):
+        #for step, batch in enumerate(train_dataloader, start=checkpoint_step):
+            # model inputs
+            q_ids, q_mask, a_ids, a_mask = batch
+            pre_loss = model(q_ids, q_mask, a_ids, a_mask, checkpoint_batch_size=args.checkpoint_batch_size)
+            loss = pre_loss.sum() / args.gradient_accumulation_steps
+            accelerator.backward(loss)
+            loc_loss += loss.item()
+            if ((step + 1) % args.gradient_accumulation_steps == 0) or (step + 1 == len(train_dataloader)):
+                current_loss = loc_loss
+                optimizer.step()
+                scheduler.step()
+                optimizer.zero_grad()
+                progress_bar.update(1)
+                progress_bar.set_postfix(loss=loc_loss)
+                loc_loss = 0
+                completed_steps += 1
+
+            if step % (args.log_freq * args.gradient_accumulation_steps) == 0:
+                accelerator.wait_for_everyone()
+                unwrapped_model = accelerator.unwrap_model(model)
+                eval_loss = evaluate_qa_retriever(unwrapped_model, eval_dataloader)
+                logger.info(f"Train loss {current_loss} , eval loss {eval_loss}")
+                if args.wandb and accelerator.is_local_main_process:
+                    import wandb
+                    wandb.log({"loss": current_loss, "eval_loss": eval_loss, "step": completed_steps})
+
+            if completed_steps >= args.max_train_steps:
+                break
+
+        logger.info("Saving model {}".format(args.model_save_name))
+        accelerator.wait_for_everyone()
+        unwrapped_model = accelerator.unwrap_model(model)
+        accelerator.save(unwrapped_model.state_dict(), "{}_{}.bin".format(args.model_save_name, epoch))
+        eval_loss = evaluate_qa_retriever(unwrapped_model, eval_dataloader)
+        logger.info("Evaluation loss epoch {:4d}: {:.3f}".format(epoch, eval_loss))
+
+
+if __name__ == "__main__":
+    parser = get_parser()
+    parser.add_argument(
+        "--wandb",
+        action="store_true",
+        help="Whether to use W&B logging",
+    )
+    main_args, _ = parser.parse_known_args()
+    config = {"args": main_args}
+    if main_args.wandb:
+        import wandb
+        wandb.init(project="Retriever")
+
+    train(config=config)
+

training/run_retriever_no_trainer_gpl.py

@@ -0,0 +1,403 @@
+import argparse
+import logging
+import math
+from dataclasses import dataclass
+from typing import List, Any, Union, Optional
+
+import torch
+import ujson
+from accelerate import Accelerator
+from accelerate.utils import set_seed
+from torch import nn, Tensor
+from torch.nn import functional as F
+from torch.utils.data import Dataset, RandomSampler, DataLoader, SequentialSampler
+from tqdm.auto import tqdm
+from transformers import get_scheduler, AutoTokenizer, AutoModel, AdamW, SchedulerType, PreTrainedTokenizerBase, AutoModelForSequenceClassification, BatchEncoding
+from transformers.file_utils import PaddingStrategy
+
+logger = logging.getLogger(__name__)
+
+
+def get_parser():
+    parser = argparse.ArgumentParser(description="Train LFQA retriever")
+    parser.add_argument(
+        "--dpr_input_file",
+        type=str,
+        help="DPR formatted input file with question/positive/negative pairs in a JSONL file",
+    )
+
+    parser.add_argument(
+        "--per_device_train_batch_size",
+        type=int,
+        default=32,
+    )
+
+    parser.add_argument(
+        "--per_device_eval_batch_size",
+        type=int,
+        default=32,
+        help="Batch size (per device) for the evaluation dataloader.",
+    )
+
+    parser.add_argument(
+        "--max_length",
+        type=int,
+        default=128,
+    )
+
+
+    parser.add_argument(
+        "--pretrained_model_name",
+        type=str,
+        default="sentence-transformers/all-MiniLM-L6-v2",
+    )
+
+    parser.add_argument(
+        "--ce_model_name",
+        type=str,
+        default="cross-encoder/ms-marco-MiniLM-L-6-v2",
+    )
+
+    parser.add_argument(
+        "--model_save_name",
+        type=str,
+        default="eli5_retriever_model_l-12_h-768_b-512-512",
+    )
+
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=2e-5,
+    )
+
+    parser.add_argument(
+        "--weight_decay",
+        type=float,
+        default=0.01,
+    )
+
+    parser.add_argument(
+        "--log_freq",
+        type=int,
+        default=500,
+        help="Log train/validation loss every log_freq update steps"
+    )
+
+    parser.add_argument(
+        "--num_train_epochs",
+        type=int,
+        default=4,
+    )
+
+    parser.add_argument(
+        "--max_train_steps",
+        type=int,
+        default=None,
+        help="Total number of training steps to perform. If provided, overrides num_train_epochs.",
+    )
+
+    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(
+        "--lr_scheduler_type",
+        type=SchedulerType,
+        default="linear",  # this is linear with warmup
+        help="The scheduler type to use.",
+        choices=["linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup"],
+    )
+
+    parser.add_argument(
+        "--num_warmup_steps",
+        type=int,
+        default=100,
+        help="Number of steps for the warmup in the lr scheduler."
+    )
+
+    parser.add_argument(
+        "--warmup_percentage",
+        type=float,
+        default=0.08,
+        help="Number of steps for the warmup in the lr scheduler."
+    )
+    return parser
+
+
+@dataclass
+class InputExample:
+    guid: str = ""
+    texts: List[str] = None
+    label: Union[int, float] = 0
+
+
+class DPRDataset(Dataset):
+    """
+    Dataset DPR format of question, answers, positive, negative, and hard negative passages
+    See https://github.com/facebookresearch/DPR#retriever-input-data-format for more details
+    """
+
+    def __init__(self, file_path: str, include_all_positive: bool = False) -> None:
+        super().__init__()
+        with open(file_path, "r") as fp:
+            self.data = []
+
+            def dpr_example_to_input_example(idx, dpr_item):
+                examples = []
+                for p_idx, p_item in enumerate(dpr_item["positive_ctxs"]):
+                    for n_idx, n_item in enumerate(dpr_item["negative_ctxs"]):
+                        examples.append(InputExample(guid=[idx, p_idx, n_idx], texts=[dpr_item["question"],
+                                                                                      p_item["text"],
+                                                                                      n_item["text"]]))
+                    if not include_all_positive:
+                        break
+                return examples
+
+            for idx, line in enumerate(fp):
+                self.data.extend(dpr_example_to_input_example(idx, ujson.loads(line)))
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, index):
+        return self.data[index]
+
+
+def dpr_collate_fn(batch):
+    query_id, pos_id, neg_id = zip(*[example.guid for example in batch])
+    query, pos, neg = zip(*[example.texts for example in batch])
+    return (query_id, pos_id, neg_id), (query, pos, neg)
+
+
+# Mean Pooling - Take attention mask into account for correct averaging
+def mean_pooling(model_output, attention_mask):
+    token_embeddings = model_output[0]  # First element of model_output contains all token embeddings
+    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
+    sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, 1)
+    sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)
+    return sum_embeddings / sum_mask
+
+
+@dataclass
+class CrossEncoderCollator:
+    tokenizer: PreTrainedTokenizerBase
+    model: Any
+    target_tokenizer: PreTrainedTokenizerBase
+    padding: Union[bool, str, PaddingStrategy] = True
+    max_length: Optional[int] = None
+    pad_to_multiple_of: Optional[int] = None
+    return_tensors: str = "pt"
+
+    def __call__(self, batch):
+        query_id, pos_id, neg_id = zip(*[example.guid for example in batch])
+        query, pos_passage, neg_passage = zip(*[example.texts for example in batch])
+        batch_input: List[List[str]] = list(zip(query, pos_passage)) + list(zip(query, neg_passage))
+        features = self.tokenizer(batch_input, padding=self.padding, truncation=True,
+                                  return_tensors=self.return_tensors)
+        with torch.no_grad():
+            scores = self.model(**features).logits
+
+        labels = scores[:len(query)] - scores[len(query):]
+        batch_input: List[str] = list(query) + list(pos_passage) + list(neg_passage)
+        #breakpoint()
+        encoded_input = self.target_tokenizer(batch_input, padding=True, truncation=True,
+                                              max_length=256, return_tensors='pt')
+
+        encoded_input["labels"] = labels
+
+        return encoded_input
+
+
+class RetrievalQAEmbedder(torch.nn.Module):
+    def __init__(self, sent_encoder, sent_tokenizer, batch_size:int = 32):
+        super(RetrievalQAEmbedder, self).__init__()
+        dim = sent_encoder.config.hidden_size
+        self.model = sent_encoder
+        self.tokenizer = sent_tokenizer
+        self.scale = 1
+        self.similarity_fct = 'dot'
+        self.batch_size = 32
+        self.loss_fct = nn.MSELoss()
+
+    def forward(self, examples: BatchEncoding):
+        # Tokenize sentences
+        labels = examples.pop("labels")
+        # Compute token embeddings
+        model_output = self.model(**examples)
+
+        examples["labels"] = labels
+
+        # Perform pooling. In this case, mean pooling
+        sentence_embeddings = mean_pooling(model_output, examples['attention_mask'])
+        target_shape = (3, self.batch_size, sentence_embeddings.shape[-1])
+        sentence_embeddings_reshaped = torch.reshape(sentence_embeddings, target_shape)
+        
+        #breakpoint()
+
+        embeddings_query = sentence_embeddings_reshaped[0]
+        embeddings_pos = sentence_embeddings_reshaped[1]
+        embeddings_neg = sentence_embeddings_reshaped[2]
+
+        if self.similarity_fct == 'cosine':
+            embeddings_query = F.normalize(embeddings_query, p=2, dim=1)
+            embeddings_pos = F.normalize(embeddings_pos, p=2, dim=1)
+            embeddings_neg = F.normalize(embeddings_neg, p=2, dim=1)
+
+        scores_pos = (embeddings_query * embeddings_pos).sum(dim=-1) * self.scale
+        scores_neg = (embeddings_query * embeddings_neg).sum(dim=-1) * self.scale
+        margin_pred = scores_pos - scores_neg
+        #breakpoint()
+        return self.loss_fct(margin_pred, labels.squeeze())
+
+
+def evaluate_qa_retriever(model, data_loader):
+    # make iterator
+    epoch_iterator = tqdm(data_loader, desc="Iteration", disable=True)
+    tot_loss = 0.0
+    with torch.no_grad():
+        for step, batch in enumerate(epoch_iterator):
+            q_ids, q_mask, a_ids, a_mask = batch
+            loss = model(q_ids, q_mask, a_ids, a_mask)
+            tot_loss += loss.item()
+        return tot_loss / (step + 1)
+
+
+def train(config):
+    set_seed(42)
+    args = config["args"]
+
+    # Initialize the accelerator. We will let the accelerator handle device placement for us in this example.
+    accelerator = Accelerator()
+    # Make one log on every process with the configuration for debugging.
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        level=logging.INFO,
+    )
+    logger.setLevel(logging.INFO if accelerator.is_local_main_process else logging.ERROR)
+    logger.info(accelerator.state)
+
+    # prepare torch Dataset objects
+    train_dataset = DPRDataset(file_path=args.dpr_input_file)
+    valid_dataset = Dataset()
+
+    base_tokenizer = AutoTokenizer.from_pretrained(args.pretrained_model_name)
+    base_model = AutoModel.from_pretrained(args.pretrained_model_name)
+
+    ce_tokenizer = AutoTokenizer.from_pretrained(args.ce_model_name)
+    ce_model = AutoModelForSequenceClassification.from_pretrained(args.ce_model_name)
+    _ = ce_model.eval()
+
+    model = RetrievalQAEmbedder(base_model, base_tokenizer)
+    no_decay = ['bias', 'LayerNorm.weight']
+    optimizer_grouped_parameters = [
+        {'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
+         'weight_decay': args.weight_decay},
+        {'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
+    ]
+    optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate)
+
+    cec = CrossEncoderCollator(model=ce_model, tokenizer=ce_tokenizer, target_tokenizer=base_tokenizer)
+    train_dataloader = DataLoader(train_dataset, batch_size=args.per_device_train_batch_size,
+                                  sampler=RandomSampler(train_dataset), collate_fn=cec)
+
+    eval_dataloader = DataLoader(valid_dataset, batch_size=args.per_device_eval_batch_size,
+                                 sampler=SequentialSampler(valid_dataset), collate_fn=cec)
+
+    # train the model
+    model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(model, optimizer,
+                                                                              train_dataloader, eval_dataloader)
+    # Scheduler and math around the number of training steps.
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+    if args.max_train_steps is None:
+        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
+    else:
+        args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+
+    num_warmup_steps = args.num_warmup_steps if args.num_warmup_steps else math.ceil(args.max_train_steps *
+                                                                                     args.warmup_percentage)
+    scheduler = get_scheduler(
+        name=args.lr_scheduler_type,
+        optimizer=optimizer,
+        num_warmup_steps=args.num_warmup_steps,
+        num_training_steps=args.max_train_steps,
+    )
+
+    # Train!
+    total_batch_size = args.per_device_train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+
+    logger.info("***** Running training *****")
+    logger.info(f"  Num examples = {len(train_dataset)}")
+    logger.info(f"  Num Epochs = {args.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_batch_size}")
+    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
+    logger.info(f"  Total optimization steps = {args.max_train_steps}")
+    logger.info(f"  Warmup steps = {num_warmup_steps}")
+    logger.info(f"  Logging training progress every {args.log_freq} optimization steps")
+
+    loc_loss = 0.0
+    current_loss = 0.0
+    checkpoint_step = 0
+
+    completed_steps = checkpoint_step
+    progress_bar = tqdm(range(args.max_train_steps), initial=checkpoint_step,
+                        disable=not accelerator.is_local_main_process)
+    for epoch in range(args.num_train_epochs):
+        model.train()
+        for step, batch in enumerate(train_dataloader, start=checkpoint_step):
+            # model inputs
+            pre_loss = model(batch)
+            loss = pre_loss / args.gradient_accumulation_steps
+            accelerator.backward(loss)
+            loc_loss += loss.item()
+            if ((step + 1) % args.gradient_accumulation_steps == 0) or (step + 1 == len(train_dataloader)):
+                current_loss = loc_loss
+                optimizer.step()
+                scheduler.step()
+                optimizer.zero_grad()
+                progress_bar.update(1)
+                progress_bar.set_postfix(loss=loc_loss)
+                loc_loss = 0
+                completed_steps += 1
+
+            if step % (args.log_freq * args.gradient_accumulation_steps) == 0:
+                # accelerator.wait_for_everyone()
+                # unwrapped_model = accelerator.unwrap_model(model)
+                # eval_loss = evaluate_qa_retriever(unwrapped_model, eval_dataloader)
+                eval_loss = 0
+                logger.info(f"Train loss {current_loss} , eval loss {eval_loss}")
+                if args.wandb and accelerator.is_local_main_process:
+                    import wandb
+                    wandb.log({"loss": current_loss, "eval_loss": eval_loss, "step": completed_steps})
+
+            if completed_steps >= args.max_train_steps:
+                break
+
+        logger.info("Saving model {}".format(args.model_save_name))
+        accelerator.wait_for_everyone()
+        unwrapped_model = accelerator.unwrap_model(model)
+        accelerator.save(unwrapped_model.state_dict(), "{}_{}.bin".format(args.model_save_name, epoch))
+        eval_loss = evaluate_qa_retriever(unwrapped_model, eval_dataloader)
+        logger.info("Evaluation loss epoch {:4d}: {:.3f}".format(epoch, eval_loss))
+
+
+if __name__ == "__main__":
+    parser = get_parser()
+    parser.add_argument(
+        "--wandb",
+        action="store_true",
+        help="Whether to use W&B logging",
+    )
+    main_args, _ = parser.parse_known_args()
+    config = {"args": main_args}
+    if main_args.wandb:
+        import wandb
+
+        wandb.init(project="Retriever")
+
+    train(config=config)
+

training/run_seq2seq_no_trainer.py

@@ -0,0 +1,446 @@
+import argparse
+import logging
+import math
+import re
+
+import numpy as np
+import torch
+from accelerate import Accelerator
+from accelerate.utils import set_seed
+from torch.utils.data import DataLoader
+from tqdm.auto import tqdm
+from transformers import get_scheduler, AutoTokenizer, AdamW, SchedulerType, AutoModelForSeq2SeqLM, \
+    DataCollatorWithPadding
+
+from datasets import load_dataset
+
+logger = logging.getLogger(__name__)
+
+
+def get_parser():
+    parser = argparse.ArgumentParser(description="Train ELI5 seq2seq answer generation model")
+    parser.add_argument(
+        "--dataset_name",
+        type=str,
+        default="vblagoje/lfqa",
+        help="The name of the dataset to use (via the datasets library).",
+    )
+
+    parser.add_argument(
+        "--per_device_train_batch_size",
+        type=int,
+        default=4,
+    )
+
+    parser.add_argument(
+        "--per_device_eval_batch_size",
+        type=int,
+        default=4,
+        help="Batch size (per device) for the evaluation dataloader.",
+    )
+
+    parser.add_argument(
+        "--pretrained_model_name",
+        type=str,
+        default="facebook/bart-large",
+    )
+
+    parser.add_argument(
+        "--model_save_name",
+        type=str,
+        default="eli5_bart_model",
+    )
+
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=2e-4,
+    )
+
+    parser.add_argument(
+        "--weight_decay",
+        type=float,
+        default=0.0,
+        help="Weight decay to use."
+    )
+
+    parser.add_argument(
+        "--log_freq",
+        type=int,
+        default=100,
+        help="Log train/validation loss every log_freq update steps"
+    )
+
+    parser.add_argument(
+        "--ignore_pad_token_for_loss",
+        type=bool,
+        default=True,
+        help="Whether to ignore the tokens corresponding to " "padded labels in the loss computation or not.",
+    )
+
+    parser.add_argument(
+        "--num_train_epochs",
+        type=int,
+        default=3,
+    )
+
+    parser.add_argument(
+        "--max_train_steps",
+        type=int,
+        default=None,
+        help="Total number of training steps to perform. If provided, overrides num_train_epochs.",
+    )
+
+    parser.add_argument(
+        "--gradient_accumulation_steps",
+        type=int,
+        default=16,
+        help="Number of updates steps to accumulate before performing a backward/update pass.",
+    )
+
+    parser.add_argument(
+        "--pad_to_max_length",
+        action="store_true",
+        help="If passed, pad all samples to `max_length`. Otherwise, dynamic padding is used.",
+    )
+
+    parser.add_argument(
+        "--overwrite_cache", type=bool, default=None, help="Overwrite the cached training and evaluation sets"
+    )
+
+    parser.add_argument(
+        "--max_source_length",
+        type=int,
+        default=1024,
+        help="The maximum total input sequence length after "
+             "tokenization.Sequences longer than this will be truncated, sequences shorter will be padded.",
+    )
+
+    parser.add_argument(
+        "--max_target_length",
+        type=int,
+        default=360,
+        help="The maximum total sequence length for target text after "
+             "tokenization. Sequences longer than this will be truncated, sequences shorter will be padded."
+    )
+
+    parser.add_argument(
+        "--lr_scheduler_type",
+        type=SchedulerType,
+        default="linear",  # this is linear with warmup
+        help="The scheduler type to use.",
+        choices=["linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup"],
+    )
+
+    parser.add_argument(
+        "--num_warmup_steps",
+        type=int,
+        default=None,
+        help="Number of steps for the warmup in the lr scheduler."
+    )
+
+    parser.add_argument(
+        "--warmup_percentage",
+        type=float,
+        default=0.08,
+        help="Number of steps for the warmup in the lr scheduler."
+    )
+    return parser
+
+
+def cleanup_references(text):
+    # URL reference where we need to remove both the link text and URL
+    # ...and this letter is used by most biographers as the cornerstone of Lee's personal
+    # views on slavery ([1](_URL_2_ & pg=PA173), [2](_URL_1_), [3](_URL_5_)).
+    # ...and this letter is used by most biographers as the cornerstone of Lee's personal views on slavery.
+    result = re.sub(r"[\(\s]*\[\d+\]\([^)]+\)[,)]*", "", text, 0, re.MULTILINE)
+
+    # URL reference where we need to preserve link text but remove URL
+    # At the outbreak of the Civil War, [Leyburn left his church](_URL_19_) and joined the South.
+    # At the outbreak of the Civil War, Leyburn left his church and joined the South.
+    result = re.sub(r"\[([^]]+)\]\([^)]+\)", "\\1", result, 0, re.MULTILINE)
+
+    # lastly remove just dangling _URL_[0-9]_ URL references
+    result = re.sub(r"_URL_\d_", "", result, 0, re.MULTILINE)
+    return result
+
+
+def clean_answer(text):
+    result = cleanup_references(text)
+    result = result.replace("\n", " ")
+    result = re.sub(r"\s\s+", " ", result)
+    result = re.sub(r"BULLET::::-", "", result)
+    return result.strip()
+
+
+def clean_question(text):
+    result = cleanup_references(text)
+    result = result.replace("\n", " ")
+    result = re.sub(r"\s\s+", " ", result)
+    result = result.replace("[deleted]", "")
+    return result.lower().strip()
+
+
+def prepare_support_docs(example):
+    provenances = example["output"][-1]["provenance"]
+    context = "<P> " + " <P> ".join([p["text"] for p in provenances])
+    return {"context": context}
+
+
+def preprocess_eli5(examples, **fn_kwargs):
+    document_cache = fn_kwargs["document_cache"]
+    training = fn_kwargs.get("training", True)
+    extra_answer_threshold = fn_kwargs.get("extra_answer_threshold", 3)
+    include_selftext = fn_kwargs.get("include_selftext", False)
+    exclude_answer_patterns = fn_kwargs.get("exclude_answer_patterns", [])
+
+    questions, contexts, answers = [], [], []
+    for q_id, question, selftext, answer in zip(examples["q_id"], examples["title"], examples["selftext"],
+                                                examples["answers"]):
+        accepted_answer_idx = []
+        if training:
+            accepted_answer_idx = [idx for idx, score in enumerate(answer["score"]) if
+                                   score > extra_answer_threshold]
+        if not training or not accepted_answer_idx:
+            accepted_answer_idx = [0]
+        document = document_cache[q_id]
+        for idx in accepted_answer_idx:
+            skip_answer = any([p.search(answer["text"][idx]) for p in exclude_answer_patterns])
+            if skip_answer:
+                continue
+            if include_selftext:
+                questions.append(clean_question(f"{question} {selftext}"))
+            else:
+                questions.append(clean_question(question))
+            contexts.append(document.lower().strip())
+            answers.append(clean_answer(answer["text"][idx]))
+
+    return {"question": questions, "context": contexts, "answer": answers}
+
+
+def eval_qa_s2s_epoch(model, dataloader, accelerator, args):
+    model.eval()
+    num_eval_steps = math.ceil(len(dataloader))
+    progress_bar = tqdm(range(num_eval_steps), disable=not accelerator.is_local_main_process)
+    total_loss = 0.
+    with torch.no_grad():
+        for step, batch in enumerate(dataloader):
+            outputs = model(**batch)
+            loss = outputs.loss
+            total_loss += loss.item()
+            progress_bar.update(1)
+            progress_bar.set_postfix(loss=round((total_loss / (step + 1)), 3))
+        return total_loss / (step + 1)
+
+
+def train(config):
+    set_seed(42)
+    args = config["args"]
+    eli5 = load_dataset(args.dataset_name)
+
+    support_docs = load_dataset("vblagoje/lfqa_support_docs")
+
+    # Initialize the accelerator. We will let the accelerator handle device placement for us in this example.
+    accelerator = Accelerator()
+    # Make one log on every process with the configuration for debugging.
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        level=logging.INFO,
+    )
+    logger.setLevel(logging.INFO if accelerator.is_local_main_process else logging.ERROR)
+    logger.info(accelerator.state)
+
+    tokenizer = AutoTokenizer.from_pretrained(args.pretrained_model_name)
+    model = AutoModelForSeq2SeqLM.from_pretrained(args.pretrained_model_name)
+
+    # Optimizer
+    # Split weights in two groups, one with weight decay and the other not.
+    no_decay = ["bias", "LayerNorm.weight"]
+    optimizer_grouped_parameters = [
+        {
+            "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
+            "weight_decay": args.weight_decay,
+        },
+        {
+            "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
+            "weight_decay": 0.0,
+        },
+    ]
+    optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, weight_decay=args.weight_decay)
+
+    processed_datasets = {}
+    support_docs_prepared = {}
+    with accelerator.main_process_first():
+        for split in ["train", "validation"]:
+            support_docs_prepared[split] = support_docs[split].map(prepare_support_docs,
+                                                                   batched=False,
+                                                                   cache_file_name=f"./support_docs_{split}.arrow",
+                                                                   load_from_cache_file=not args.overwrite_cache,
+                                                                   desc="Preparing support docs",
+                                                                   )
+        column_names = eli5["train"].column_names
+        for split in ["train", "validation"]:
+            d_cache = dict([(e["id"], e["context"]) for e in tqdm(support_docs_prepared[split],
+                                                                  desc=f"Adding support docs to LFQA {split}")])
+            processed_datasets[split] = eli5[split].map(preprocess_eli5,
+                                                        batched=True,
+                                                        remove_columns=column_names,
+                                                        cache_file_name=f"./processed_datasets_{split}.arrow",
+                                                        load_from_cache_file=not args.overwrite_cache,
+                                                        desc="Preparing dataset for tokenization",
+                                                        fn_kwargs={"document_cache": d_cache,
+                                                                   "training": split == "train",
+                                                                   "exclude_answer_patterns": [re.compile("not sure what you"),
+                                                                                               re.compile("\n\n >")]}
+                                                        )
+
+    padding = "max_length" if args.pad_to_max_length else False
+    # Temporarily set max_target_length for training.
+    max_target_length = args.max_target_length
+
+    label_pad_token_id = -100 if args.ignore_pad_token_for_loss else tokenizer.pad_token_id
+
+    def tokenize_dataset(examples):
+        inputs = ["question: {} context: {}".format(q, c) for q, c in zip(examples["question"], examples["context"])]
+        targets = examples["answer"]
+        model_inputs = tokenizer(inputs, max_length=args.max_source_length, padding=padding, truncation=True)
+
+        # Setup the tokenizer for targets
+        with tokenizer.as_target_tokenizer():
+            labels = tokenizer(targets, max_length=max_target_length, padding=True, truncation=True,
+                               return_tensors="np")
+
+        model_inputs["decoder_input_ids"] = labels["input_ids"][:, :-1].tolist()
+        # replace pad_token_id with label_pad_token_id to avoid loss calculation on those tokens
+        labels["input_ids"] = np.where(labels["input_ids"] == tokenizer.pad_token_id,
+                                       label_pad_token_id, labels["input_ids"])
+
+        model_inputs["labels"] = labels["input_ids"][:, 1:].tolist()
+        return model_inputs
+
+    tokenized_datasets = {}
+    with accelerator.main_process_first():
+        for split, dataset in processed_datasets.items():
+            tokenized_datasets[split] = dataset.map(
+                tokenize_dataset,
+                batched=True,
+                cache_file_name=f"./tokenized_dataset_{split}.arrow",
+                remove_columns=dataset.column_names,
+                load_from_cache_file=not args.overwrite_cache,
+                desc="Running tokenizer on dataset"
+            )
+
+    train_dataset = tokenized_datasets["train"]
+    eval_dataset = tokenized_datasets["validation"]
+    train_dataset.set_format(type='torch')
+    eval_dataset.set_format(type='torch')
+
+    data_collator = DataCollatorWithPadding(tokenizer, "max_length")
+
+    # first epoch we don't shuffle
+    train_dataloader = DataLoader(train_dataset, shuffle=False, batch_size=args.per_device_train_batch_size,
+                                  collate_fn=data_collator)
+    eval_dataloader = DataLoader(eval_dataset, batch_size=args.per_device_eval_batch_size, collate_fn=data_collator)
+
+    # train the model
+    model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(model, optimizer, train_dataloader,
+                                                                              eval_dataloader)
+    # Scheduler and math around the number of training steps.
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+    if args.max_train_steps is None:
+        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
+    else:
+        args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+
+    num_warmup_steps = args.num_warmup_steps if args.num_warmup_steps else math.ceil(args.max_train_steps *
+                                                                                     args.warmup_percentage)
+    scheduler = get_scheduler(
+        name=args.lr_scheduler_type,
+        optimizer=optimizer,
+        num_warmup_steps=num_warmup_steps,
+        num_training_steps=args.max_train_steps,
+    )
+    # Train!
+    total_batch_size = args.per_device_train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+
+    logger.info("***** Running training *****")
+    logger.info(f"  Num examples = {len(train_dataset)}")
+    logger.info(f"  Num eval examples = {len(eval_dataset)}")
+    logger.info(f"  Num Epochs = {args.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_batch_size}")
+    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
+    logger.info(f"  Total optimization steps = {args.max_train_steps}")
+    logger.info(f"  Warmup steps = {num_warmup_steps}")
+    logger.info(f"  Logging training progress every {args.log_freq} optimization steps")
+
+    # Only show the progress bar once on each machine.
+    progress_bar = tqdm(range(args.max_train_steps), disable=not accelerator.is_local_main_process)
+    completed_steps = 0
+    switched_train_dataloader = False
+    for epoch in range(args.num_train_epochs):
+        model.train()
+        if epoch > 0 and not switched_train_dataloader:
+            train_dataloader = DataLoader(train_dataset, batch_size=args.per_device_train_batch_size,
+                                          shuffle=True, collate_fn=data_collator)
+            train_dataloader = accelerator.prepare(train_dataloader)
+            switched_train_dataloader = True
+
+        for step, batch in enumerate(train_dataloader):
+            outputs = model(**batch)
+            loss = torch.mean(outputs.loss)
+            accelerator.backward(loss)
+            if ((step + 1) % args.gradient_accumulation_steps == 0) or (step + 1 == len(train_dataloader)):
+                optimizer.step()
+                scheduler.step()
+                optimizer.zero_grad()
+                progress_bar.update(1)
+                progress_bar.set_postfix(loss=round(loss.item(), 3))
+                completed_steps += 1
+
+            if completed_steps >= args.max_train_steps:
+                break
+
+            if step % (args.log_freq * args.gradient_accumulation_steps) == 0:
+                validation_loss = eval_qa_s2s_epoch(model, eval_dataloader, accelerator, args)
+                model.train()
+                logger.info(f"Train loss {loss.item()} , validation loss {validation_loss}")
+                if args.wandb and accelerator.is_local_main_process:
+                    import wandb
+                    wandb.log({"loss": loss.item(),
+                               "lr": scheduler.get_last_lr()[0],
+                               "validation_loss": validation_loss,
+                               "completed_steps": completed_steps})
+
+        logger.info("Saving model {}".format(args.model_save_name))
+        accelerator.wait_for_everyone()
+        unwrapped_model = accelerator.unwrap_model(model)
+        accelerator.save(unwrapped_model.state_dict(), "{}_{}.bin".format(args.model_save_name, epoch))
+
+        # Calculating the validation loss over epoch
+        validation_loss = eval_qa_s2s_epoch(model, eval_dataloader, accelerator, args)
+
+        logger.info("Epoch: {}".format(epoch))
+        logger.info("Validation loss: {}".format(validation_loss))
+
+
+def main():
+    parser = get_parser()
+    parser.add_argument(
+        "--wandb",
+        action="store_true",
+        help="If true, use W&B logging",
+    )
+    main_args, _ = parser.parse_known_args()
+    config = {"args": main_args}
+    if main_args.wandb:
+        import wandb
+        wandb.init(project="Bart_ELI5")
+    train(config=config)
+
+
+main()
+
+
+