from transformers import LukePreTrainedModel, LukeModel, AutoTokenizer, TrainingArguments, default_data_collator, Trainer, AutoModelForQuestionAnswering from transformers.modeling_outputs import ModelOutput from typing import Optional, Tuple, Union import numpy as np from tqdm import tqdm import evaluate import torch from dataclasses import dataclass from datasets import load_dataset, concatenate_datasets from torch import nn from torch.nn import CrossEntropyLoss import collections PEFT = False tf32 = True fp16= True train = False test = True trained_model = "LUKE_squadshift" train_checkpoint = None base_tokenizer = "roberta-base" base_model = "studio-ousia/luke-base" # base_tokenizer = "xlnet-base-cased" # base_model = "xlnet-base-cased" # base_tokenizer = "bert-base-cased" # base_model = "SpanBERT/spanbert-base-cased" torch.backends.cuda.matmul.allow_tf32 = tf32 torch.backends.cudnn.allow_tf32 = tf32 device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") if tf32: trained_model += "_tf32" # https://github.com/huggingface/transformers/blob/v4.34.1/src/transformers/models/luke/modeling_luke.py#L319-L353 # Taken from HF repository, easier to include additional features -- Currently identical to LukeForQuestionAnswering by HF @dataclass class LukeQuestionAnsweringModelOutput(ModelOutput): """ Outputs of question answering models. Args: loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): Total span extraction loss is the sum of a Cross-Entropy for the start and end positions. start_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length)`): Span-start scores (before SoftMax). end_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length)`): Span-end scores (before SoftMax). hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. entity_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of shape `(batch_size, entity_length, hidden_size)`. Entity hidden-states of the model at the output of each layer plus the initial entity embedding outputs. attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. """ loss: Optional[torch.FloatTensor] = None start_logits: torch.FloatTensor = None end_logits: torch.FloatTensor = None hidden_states: Optional[Tuple[torch.FloatTensor]] = None entity_hidden_states: Optional[Tuple[torch.FloatTensor]] = None attentions: Optional[Tuple[torch.FloatTensor]] = None class AugmentedLukeForQuestionAnswering(LukePreTrainedModel): def __init__(self, config): super().__init__(config) # This is 2. self.num_labels = config.num_labels self.luke = LukeModel(config, add_pooling_layer=False) ''' Any improvement to the model are expected here. Additional features, anything... ''' self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels) # Initialize weights and apply final processing self.post_init() def forward( self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.FloatTensor] = None, token_type_ids: Optional[torch.LongTensor] = None, position_ids: Optional[torch.FloatTensor] = None, entity_ids: Optional[torch.LongTensor] = None, entity_attention_mask: Optional[torch.FloatTensor] = None, entity_token_type_ids: Optional[torch.LongTensor] = None, entity_position_ids: Optional[torch.LongTensor] = None, head_mask: Optional[torch.FloatTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, start_positions: Optional[torch.LongTensor] = None, end_positions: Optional[torch.LongTensor] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, ) -> Union[Tuple, LukeQuestionAnsweringModelOutput]: r""" start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): Labels for position (index) of the start of the labelled span for computing the token classification loss. Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence are not taken into account for computing the loss. end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): Labels for position (index) of the end of the labelled span for computing the token classification loss. Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence are not taken into account for computing the loss. """ return_dict = return_dict if return_dict is not None else self.config.use_return_dict outputs = self.luke( input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, entity_ids=entity_ids, entity_attention_mask=entity_attention_mask, entity_token_type_ids=entity_token_type_ids, entity_position_ids=entity_position_ids, head_mask=head_mask, inputs_embeds=inputs_embeds, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=True, ) sequence_output = outputs.last_hidden_state logits = self.qa_outputs(sequence_output) start_logits, end_logits = logits.split(1, dim=-1) start_logits = start_logits.squeeze(-1) end_logits = end_logits.squeeze(-1) total_loss = None if start_positions is not None and end_positions is not None: # If we are on multi-GPU, split add a dimension if len(start_positions.size()) > 1: start_positions = start_positions.squeeze(-1) if len(end_positions.size()) > 1: end_positions = end_positions.squeeze(-1) # sometimes the start/end positions are outside our model inputs, we ignore these terms ignored_index = start_logits.size(1) start_positions.clamp_(0, ignored_index) end_positions.clamp_(0, ignored_index) loss_fct = CrossEntropyLoss(ignore_index=ignored_index) start_loss = loss_fct(start_logits, start_positions) end_loss = loss_fct(end_logits, end_positions) total_loss = (start_loss + end_loss) / 2 if not return_dict: return tuple( v for v in [ total_loss, start_logits, end_logits, outputs.hidden_states, outputs.entity_hidden_states, outputs.attentions, ] if v is not None ) return LukeQuestionAnsweringModelOutput( loss=total_loss, start_logits=start_logits, end_logits=end_logits, hidden_states=outputs.hidden_states, entity_hidden_states=outputs.entity_hidden_states, attentions=outputs.attentions, ) # Get data to train model - squadshift is designed as a validation/testing set, so there are multiple answers, take the shortest def get_squadshifts_training(): wiki = load_dataset("squadshifts", "new_wiki")["test"] nyt = load_dataset("squadshifts", "nyt")["test"] reddit = load_dataset("squadshifts", "reddit")["test"] raw_dataset = concatenate_datasets([wiki, nyt, reddit]) updated = raw_dataset.map(validation_to_train) return updated def validation_to_train(example): answers = example["answers"] answer_text = answers["text"] index_min = min(range(len(answer_text)), key=lambda x : len(answer_text.__getitem__(x))) answers["text"] = answers["text"][index_min:index_min+1] answers["answer_start"] = answers["answer_start"][index_min:index_min+1] return example if __name__ == "__main__": # Setting up tokenizer and helper functions # Work-around for FastTokenizer - RoBERTa and LUKE share the same subword vocab, and we are not using entities functions of LUKE-tokenizer anyways tokenizer = AutoTokenizer.from_pretrained(base_tokenizer) # Necessary initialization max_length = 500 stride = 128 batch_size = 8 n_best = 20 max_answer_length = 30 metric = evaluate.load("squad") raw_datasets = load_dataset("squad") raw_train = raw_datasets["train"] raw_validation = raw_datasets["validation"] def compute_metrics(start_logits, end_logits, features, examples): example_to_features = collections.defaultdict(list) for idx, feature in enumerate(features): example_to_features[feature["example_id"]].append(idx) predicted_answers = [] for example in tqdm(examples): example_id = example["id"] context = example["context"] answers = [] # Loop through all features associated with that example for feature_index in example_to_features[example_id]: start_logit = start_logits[feature_index] end_logit = end_logits[feature_index] offsets = features[feature_index]["offset_mapping"] start_indexes = np.argsort(start_logit)[-1 : -n_best - 1 : -1].tolist() end_indexes = np.argsort(end_logit)[-1 : -n_best - 1 : -1].tolist() for start_index in start_indexes: for end_index in end_indexes: # Skip answers that are not fully in the context if offsets[start_index] is None or offsets[end_index] is None: continue # Skip answers with a length that is either < 0 or > max_answer_length if ( end_index < start_index or end_index - start_index + 1 > max_answer_length ): continue answer = { "text": context[offsets[start_index][0] : offsets[end_index][1]], "logit_score": start_logit[start_index] + end_logit[end_index], } answers.append(answer) # Select the answer with the best score if len(answers) > 0: best_answer = max(answers, key=lambda x: x["logit_score"]) predicted_answers.append( {"id": example_id, "prediction_text": best_answer["text"]} ) else: predicted_answers.append({"id": example_id, "prediction_text": ""}) theoretical_answers = [{"id": ex["id"], "answers": ex["answers"]} for ex in examples] return metric.compute(predictions=predicted_answers, references=theoretical_answers) def preprocess_training_examples(examples): questions = [q.strip() for q in examples["question"]] inputs = tokenizer( questions, examples["context"], max_length=max_length, truncation="only_second", stride=stride, return_overflowing_tokens=True, return_offsets_mapping=True, padding="max_length", ) offset_mapping = inputs.pop("offset_mapping") sample_map = inputs.pop("overflow_to_sample_mapping") answers = examples["answers"] start_positions = [] end_positions = [] for i, offset in enumerate(offset_mapping): sample_idx = sample_map[i] answer = answers[sample_idx] start_char = answer["answer_start"][0] end_char = answer["answer_start"][0] + len(answer["text"][0]) sequence_ids = inputs.sequence_ids(i) # Find the start and end of the context idx = 0 while sequence_ids[idx] != 1: idx += 1 context_start = idx while sequence_ids[idx] == 1: idx += 1 context_end = idx - 1 # If the answer is not fully inside the context, label is (0, 0) if offset[context_start][0] > start_char or offset[context_end][1] < end_char: start_positions.append(0) end_positions.append(0) else: # Otherwise it's the start and end token positions idx = context_start while idx <= context_end and offset[idx][0] <= start_char: idx += 1 start_positions.append(idx - 1) idx = context_end while idx >= context_start and offset[idx][1] >= end_char: idx -= 1 end_positions.append(idx + 1) inputs["start_positions"] = start_positions inputs["end_positions"] = end_positions return inputs def preprocess_validation_examples(examples): questions = [q.strip() for q in examples["question"]] inputs = tokenizer( questions, examples["context"], max_length=max_length, truncation="only_second", stride=stride, return_overflowing_tokens=True, return_offsets_mapping=True, padding="max_length", ) sample_map = inputs.pop("overflow_to_sample_mapping") example_ids = [] for i in range(len(inputs["input_ids"])): sample_idx = sample_map[i] example_ids.append(examples["id"][sample_idx]) sequence_ids = inputs.sequence_ids(i) offset = inputs["offset_mapping"][i] inputs["offset_mapping"][i] = [ o if sequence_ids[k] == 1 else None for k, o in enumerate(offset) ] inputs["example_id"] = example_ids return inputs if train: model = AutoModelForQuestionAnswering.from_pretrained(base_model).to(device) # For squadshift raw_train = get_squadshifts_training() train_dataset = raw_train.map( preprocess_training_examples, batched=True, remove_columns=raw_train.column_names, ) validation_dataset = raw_validation.map( preprocess_validation_examples, batched=True, remove_columns=raw_validation.column_names, ) # --------------- PEFT -------------------- # One epoch without PEFT took about 2h on my computer with CUDA - performance of PEFT kinda ass though if PEFT: from peft import get_peft_config, get_peft_model, LoraConfig, TaskType # ---- For all linear layers ---- import re pattern = r'\((\w+)\): Linear' linear_layers = re.findall(pattern, str(model.modules)) target_modules = list(set(linear_layers)) # If using peft, can consider increaisng r for better performance peft_config = LoraConfig( task_type=TaskType.QUESTION_ANS, inference_mode=False, r=8, lora_alpha=32, lora_dropout=0.1, target_modules=target_modules, bias='all' ) model = get_peft_model(model, peft_config) model.print_trainable_parameters() trained_model += "_PEFT" # ------------------------------------------ # args = TrainingArguments( trained_model, evaluation_strategy = "no", save_strategy="epoch", learning_rate=2e-5, per_device_train_batch_size=batch_size, per_device_eval_batch_size=batch_size, num_train_epochs=3, weight_decay=0.01, push_to_hub=True, fp16=fp16 ) trainer = Trainer( model, args, train_dataset=train_dataset, eval_dataset=validation_dataset, data_collator=default_data_collator, tokenizer=tokenizer ) trainer.train(train_checkpoint) if test: model = AutoModelForQuestionAnswering.from_pretrained(trained_model).to(device) interval = len(raw_datasets["validation"]) // 100 exact_match = 0 f1 = 0 with torch.no_grad(): for i in range(1, 101): start = interval * (i - 1) end = interval * i small_eval_set = raw_datasets["validation"].select(range(start ,end)) eval_set = small_eval_set.map( preprocess_validation_examples, batched=True, remove_columns=raw_datasets["validation"].column_names ) eval_set_for_model = eval_set.remove_columns(["example_id", "offset_mapping"]) eval_set_for_model.set_format("torch") batch = {k: eval_set_for_model[k].to(device) for k in eval_set_for_model.column_names} outputs = model(**batch) start_logits = outputs.start_logits.cpu().numpy() end_logits = outputs.end_logits.cpu().numpy() res = compute_metrics(start_logits, end_logits, eval_set, small_eval_set) exact_match += res['exact_match'] f1 += res["f1"] print("F1 score: {}".format(f1 / 100)) print("Exact match: {}".format(exact_match / 100)) # XLNET # F1 score: 91.54154256653278 # Exact match: 84.86666666666666 # SpanBERT # F1 score: 92.160285362531 # Exact match: 85.73333333333333 # LUKE SQUADSHIFT (SQUAD then SQUADSHIFT) # F1 score: 91.27683543983473 # Exact match: 84.96190476190473