fine-tune/ft.py

#! /usr/bin/env python3
from transformers import BartForConditionalGeneration, BartTokenizer, AdamW
import torch
from dl import load_dataset
from tqdm import tqdm


class FT:
    def __init__(self):
        # Enable cudnn optimizations
        torch.backends.cudnn.benchmark = True

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        # load tokenizer and model
        self.tokenizer = BartTokenizer.from_pretrained("facebook/bart-base")
        self.model = BartForConditionalGeneration.from_pretrained("facebook/bart-base")
        self.model.to(self.device)

        # set up optimizer
        self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=1e-5)

        try:
            from torch.cuda.amp import GradScaler, autocast

            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

            self.scaler = None  # We won't use a scaler if we don't have Amp

    def train_model(self, dataloader):
        self.model.train()
        total_loss = 0
        print("Training model...")
        for batch in tqdm(dataloader):
            self.optimizer.zero_grad()

            inputs = self.tokenizer(
                batch[1],
                return_tensors="pt",
                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 = self.model(
                input_ids=inputs["input_ids"],
                attention_mask=inputs["attention_mask"],
                labels=labels["input_ids"],
            )

            loss = outputs.loss
            loss.backward()

            self.optimizer.step()
            total_loss += loss.item()

        avg_train_loss = total_loss / len(dataloader)
        return avg_train_loss

    def test_model(self, dataloader):
        self.model.eval()
        total_loss = 0
        print("Testing model...")
        for batch in tqdm(dataloader):
            with torch.no_grad():
                inputs = self.tokenizer(
                    batch[1],
                    return_tensors="pt",
                    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 = 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(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,
        )
        num_epochs = 3
        for epoch in range(num_epochs):
            avg_train_loss = self.train_model(train_dataloader)
            print(f"Train loss for epoch {epoch+1}: {avg_train_loss}")

            avg_test_loss = self.test_model(test_dataloader)
            print(f"Test loss for epoch {epoch+1}: {avg_test_loss}")


if __name__ == "__main__":
    trainer = FT()
    trainer.train()