run_retriever_no_trainer_gpl.py

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)