Upload bert.py

bert.py

@@ -0,0 +1,137 @@
+import torch
+import random
+from transformers import AutoTokenizer, BertForSequenceClassification
+from datasets import load_dataset
+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, load_metric
+from torch import nn
+from torch.nn import CrossEntropyLoss
+import collections
+import re
+
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+
+tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
+model = BertForSequenceClassification.from_pretrained("bert-base-uncased").to(device)
+
+def preprocess_training_examples(examples):
+    questions = [q.strip() for q in examples["question"]]
+    answers = []
+    labels = []
+    final_questions = []
+    # Generate close-looking answers from existing context
+    for i in range(len(questions)):
+        context = examples["context"][i]
+        words = context.split()
+        final_questions.append(questions[i])
+        original_answer = examples["answers"][i]["text"][0]
+        original_words = original_answer.split()
+        answers.append(original_answer)
+        labels.append(1)
+        answer_start = examples["answers"][i]["answer_start"][0]
+        answer_end = answer_start + len(original_answer)
+
+        begin_context = context[:answer_start]
+        end_context = context[answer_end:]
+
+        end = 1 if len(original_words) == 1 else 3
+        case_ind = random.randint(0, end)
+        pre_words = begin_context.rsplit(maxsplit=1)
+        post_words = end_context.split(maxsplit=1)
+        if case_ind == 0 and pre_words: # Append left
+            words = pre_words
+            wrong_context = " ".join([words[-1], original_answer])
+            final_questions.append(questions[i])
+            answers.append(wrong_context)
+            labels.append(0)
+        elif case_ind == 1 and post_words: # Append right
+            words = post_words
+            wrong_context = " ".join([original_answer, words[0]])
+            final_questions.append(questions[i])
+            answers.append(wrong_context)
+            labels.append(0)
+        elif case_ind == 3: # Drop left
+            wrong_context = " ".join(original_words[1:])
+            final_questions.append(questions[i])
+            answers.append(wrong_context)
+            labels.append(0)
+        elif case_ind == 4: # Drop right
+            wrong_context = " ".join(original_words[:len(original_words) - 1])
+            final_questions.append(questions[i])
+            answers.append(wrong_context)
+            labels.append(0)
+
+    inputs = tokenizer(
+            final_questions,
+            answers,
+            padding="max_length",
+        )
+    inputs["labels"] = labels
+    return inputs
+
+raw_datasets = load_dataset("squad")
+raw_train = raw_datasets["train"]
+raw_eval = raw_datasets["validation"]
+
+train_dataset = raw_train.map(
+            preprocess_training_examples,
+            batched=True,
+            remove_columns=raw_train.column_names,
+        )
+
+eval_dataset = raw_eval.map(
+        preprocess_training_examples,
+        batched=True,
+        remove_columns=raw_train.column_names,
+    )
+
+batch_size = 8
+
+# train_dataset = train_dataset.with_format("torch")
+
+args = TrainingArguments(
+        "right_span_bert",
+        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=2,
+        weight_decay=0.01,
+        push_to_hub=True,
+        fp16=True
+    )
+
+def compute_metrics(eval_pred):
+   load_accuracy = evaluate.load("accuracy")
+   load_f1 = evaluate.load("f1")
+   logits, labels = eval_pred
+   predictions = np.argmax(logits, axis=-1)
+   accuracy = load_accuracy.compute(predictions=predictions, references=labels)["accuracy"]
+   f1 = load_f1.compute(predictions=predictions, references=labels)["f1"]
+   return {"accuracy": accuracy, "f1": f1}
+
+trainer = Trainer(
+    model,
+    args,
+    train_dataset=train_dataset,
+    eval_dataset=eval_dataset,
+    data_collator=default_data_collator,
+    tokenizer=tokenizer,
+    compute_metrics=compute_metrics
+    )
+
+trainer.train()
+
+res = trainer.evaluate()
+print(res)