Upload meta.py

meta.py

@@ -0,0 +1,123 @@
+import torch.nn as nn
+import torch
+from transformers import AutoTokenizer, BertForSequenceClassification, PreTrainedModel, PretrainedConfig, get_scheduler
+from transformers.modeling_outputs import SequenceClassifierOutput
+from torch.nn import CrossEntropyLoss
+from torch.optim import AdamW
+from LUKE_pipe import generate
+from datasets import load_dataset
+from accelerate import Accelerator
+from tqdm import tqdm
+
+MAX_BEAM = 10
+tf32 = True
+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")
+
+class ClassifierAdapter(nn.Module):
+    def __init__(self, l1=3):
+        super().__init__()
+        self.linear1 = nn.Linear(l1, 1)
+        self.tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
+        self.bert = BertForSequenceClassification.from_pretrained("botcon/right_span_bert")
+        self.relu = nn.ReLU()
+
+    def forward(self, questions, answers, logits):
+        beam_size = len(answers[0])
+        samples = len(questions)
+        questions = [question for _ in range(len(answers[0])) for question in questions]
+        answers = [answer for beam in answers for answer in beam]
+        input = self.tokenizer(
+            questions,
+            answers,
+            padding="max_length",
+            return_tensors="pt"
+        ).to(device)
+        bert_logits = self.bert(**input).logits
+        bert_logits = bert_logits.reshape(samples, beam_size, 2)
+        logits = torch.FloatTensor(logits).to(device).unsqueeze(-1)
+        logits = torch.cat((logits, bert_logits), dim=-1)
+        logits = self.relu(logits)
+        out = torch.squeeze(self.linear1(logits), dim=-1)
+        return out
+
+class HuggingWrapper(PreTrainedModel):
+    config_class = PretrainedConfig()
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = ClassifierAdapter()
+
+    def forward(self, **kwargs):
+        labels = kwargs.pop("labels")
+        output = self.model(**kwargs)
+        loss_fn = CrossEntropyLoss(ignore_index=MAX_BEAM)
+        loss = loss_fn(output, labels)
+        return SequenceClassifierOutput(logits=output, loss=loss)
+    
+accelerator = Accelerator(mixed_precision="fp16")
+model = HuggingWrapper.from_pretrained("botcon/special_bert").to(device)
+optimizer = AdamW(model.parameters())
+model, optimizer = accelerator.prepare(model, optimizer)
+batch_size = 2
+raw_datasets = load_dataset("squad")
+raw_train = raw_datasets["train"]
+num_updates = len(raw_train) // batch_size
+num_epoch = 2
+num_training_steps = num_updates * num_epoch
+lr_scheduler = get_scheduler(
+    "linear",
+    optimizer=optimizer,
+    num_warmup_steps=0,
+    num_training_steps=num_training_steps,
+)
+
+progress_bar = tqdm(range(num_training_steps))
+
+for epoch in range(num_epoch):
+    start = 0
+    end = batch_size
+    steps = 0
+    cumu_loss = 0
+    training_data = raw_train
+    model.train()
+    while start < len(training_data):
+        optimizer.zero_grad()
+        batch_data = raw_train.select(range(start, min(end, len(raw_train))))
+        with torch.no_grad():
+            res = generate(batch_data)
+            prediction = []
+            predicted_logit = []
+            labels = []
+            for i in range(len(res)):
+                x = res[i]
+                ground_answer = batch_data["answers"][i]["text"][0]
+                predicted_text = x["prediction_text"]
+                found = False
+                for k in range(len(predicted_text)):
+                    if predicted_text[k] == ground_answer:
+                        labels.append(k)
+                        found = True
+                        break
+                if not found:
+                    labels.append(MAX_BEAM)
+                prediction.append(predicted_text)
+                predicted_logit.append(x["logits"])
+        labels = torch.LongTensor(labels).to(device)
+        classifier_out = model(questions=batch_data["question"] , answers=prediction, logits=predicted_logit, labels=labels)
+        loss = classifier_out.loss
+        if not torch.isnan(loss).item():
+            cumu_loss += loss.item()
+            steps += 1
+        accelerator.backward(loss)
+        optimizer.step()
+        lr_scheduler.step()
+        progress_bar.update(1)
+        start += batch_size
+        end += batch_size
+        # every 100 steps
+        if steps % 100 == 0:
+            print("Cumu loss: {}".format(cumu_loss / 100))
+            cumu_loss = 0
+
+model.push_to_hub("Adapter Bert")