Upload meta.py

meta.py

@@ -1,12 +1,13 @@
 import torch.nn as nn
 import torch
-from transformers import AutoTokenizer, BertForSequenceClassification, PreTrainedModel, PretrainedConfig
+from transformers import AutoTokenizer, BertForSequenceClassification, PreTrainedModel, PretrainedConfig, AutoModelForQuestionAnswering, 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
 
@@ -16,7 +17,6 @@ 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()
@@ -52,27 +52,36 @@ class HuggingWrapper(PreTrainedModel):
         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)
-
-accelerator = Accelerator()
 optimizer = AdamW(model.parameters())
-
-num_epoch = 2
-
+model, optimizer = accelerator.prepare(model, optimizer)
+batch_size = 2
 raw_datasets = load_dataset("squad")
 raw_train = raw_datasets["train"]
-batch_size = 2
+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(training_data))))
+        batch_data = raw_train.select(range(start, min(end, len(raw_train))))
         with torch.no_grad():
             res = generate(batch_data)
             prediction = []
@@ -95,10 +104,18 @@ for epoch in range(num_epoch):
         labels = torch.LongTensor(labels).to(device)
         classifier_out = model(questions=batch_data["question"] , answers=prediction, logits=predicted_logit, labels=labels)
         loss = classifier_out.loss
-        print(loss.item())
-        loss.backward()
+        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("some_fake_bert")