wikipedia-assistant / training /run_retriever_no_trainer_gpl.py
Rschmaelzle's picture
Duplicate from deepset/wikipedia-assistant
2024325
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)