Update trainer script

ft.py

@@ -4,85 +4,126 @@ import torch
 from dl import load_dataset
 from tqdm import tqdm
 
-# Enable cudnn optimizations
-torch.backends.cudnn.benchmark = True
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+class FT:
+    def __init__(self):
+        # Enable cudnn optimizations
+        torch.backends.cudnn.benchmark = True
 
-# load tokenizer and model
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-tokenizer = BartTokenizer.from_pretrained('facebook/bart-base')
-model = BartForConditionalGeneration.from_pretrained('facebook/bart-base')
-model.to(device)
 
-# set up optimizer
+        # load tokenizer and model
-optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5)
+        self.tokenizer = BartTokenizer.from_pretrained("facebook/bart-base")
+        self.model = BartForConditionalGeneration.from_pretrained("facebook/bart-base")
+        self.model.to(self.device)
 
-# Initialize Amp. This should be optional and should not affect computation if not available
+        # set up optimizer
-try:
+        self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=1e-5)
+
+        try:
             from torch.cuda.amp import GradScaler, autocast
-    scaler = GradScaler()
+
-except ImportError:
+            self.scaler = GradScaler()
+        except ImportError:
             # If Amp is not available, we'll simply define a dummy context manager
             class autocast:
                 def __enter__(self):
                     pass
+
                 def __exit__(self, *args):
                     pass
-    scaler = None  # We won't use a scaler if we don't have Amp
 
-def train_model(dataloader):
+            self.scaler = None  # We won't use a scaler if we don't have Amp
-    model.train()
+
+    def train_model(self, dataloader):
+        self.model.train()
         total_loss = 0
         print("Training model...")
         for batch in tqdm(dataloader):
-        optimizer.zero_grad()
+            self.optimizer.zero_grad()
 
-        inputs = tokenizer(batch[1], return_tensors="pt", padding=True, truncation=True, max_length=512)
+            inputs = self.tokenizer(
-        inputs.to(device)
+                batch[1],
-        labels = tokenizer(batch[0], return_tensors="pt", padding=True, truncation=True, max_length=512)
+                return_tensors="pt",
-        labels.to(device)
+                padding=True,
+                truncation=True,
+                max_length=512,
+            )
+            inputs.to(self.device)
+            labels = self.tokenizer(
+                batch[0],
+                return_tensors="pt",
+                padding=True,
+                truncation=True,
+                max_length=512,
+            )
+            labels.to(self.device)
 
-        outputs = model(input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"], labels=labels["input_ids"])
+            outputs = self.model(
+                input_ids=inputs["input_ids"],
+                attention_mask=inputs["attention_mask"],
+                labels=labels["input_ids"],
+            )
 
             loss = outputs.loss
             loss.backward()
 
-        optimizer.step()
+            self.optimizer.step()
             total_loss += loss.item()
 
         avg_train_loss = total_loss / len(dataloader)
         return avg_train_loss
 
-def test_model(dataloader):
+    def test_model(self, dataloader):
-    model.eval()
+        self.model.eval()
         total_loss = 0
         print("Testing model...")
         for batch in tqdm(dataloader):
             with torch.no_grad():
-            inputs = tokenizer(batch[1], return_tensors="pt", padding=True, truncation=True, max_length=512)
+                inputs = self.tokenizer(
-            inputs.to(device)
+                    batch[1],
-            labels = tokenizer(batch[0], return_tensors="pt", padding=True, truncation=True, max_length=512)
+                    return_tensors="pt",
-            labels.to(device)
+                    padding=True,
-            outputs = model(input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"], labels=labels["input_ids"])
+                    truncation=True,
+                    max_length=512,
+                )
+                inputs.to(self.device)
+                labels = self.tokenizer(
+                    batch[0],
+                    return_tensors="pt",
+                    padding=True,
+                    truncation=True,
+                    max_length=512,
+                )
+                labels.to(self.device)
+                outputs = self.model(
+                    input_ids=inputs["input_ids"],
+                    attention_mask=inputs["attention_mask"],
+                    labels=labels["input_ids"],
+                )
                 loss = outputs.loss
                 total_loss += loss.item()
 
         avg_test_loss = total_loss / len(dataloader)
         return avg_test_loss
 
-def train():
+    def train(self):
         train_dataloader, test_dataloader = load_dataset(
             "../datasets/deu_mixed-typical_2011_1M/deu_mixed-typical_2011_1M-sentences.txt",
-        100, 100, 1, test_ratio=0.2
+            100,
+            100,
+            1,
+            test_ratio=0.2,
         )
         num_epochs = 3
         for epoch in range(num_epochs):
-        avg_train_loss = train_model(train_dataloader)
+            avg_train_loss = self.train_model(train_dataloader)
             print(f"Train loss for epoch {epoch+1}: {avg_train_loss}")
 
-        avg_test_loss = test_model(test_dataloader)
+            avg_test_loss = self.test_model(test_dataloader)
             print(f"Test loss for epoch {epoch+1}: {avg_test_loss}")
 
-if __name__ == "__main__":
-    train()
 
+if __name__ == "__main__":
+    trainer = FT()
+    trainer.train()