coderchen01/MMSD2.0

from typing import TypedDict, cast

import pytorch_lightning as pl
from datasets import Dataset, load_dataset
from torch import Tensor
from torch.utils.data import DataLoader
from transformers import CLIPProcessor


class MMSDModelInput(TypedDict):
    pixel_values: Tensor
    input_ids: Tensor
    attention_mask: Tensor
    label: Tensor
    id: list[str]


class MMSDDatasetModule(pl.LightningDataModule):

    def __init__(
        self,
        clip_ckpt_name: str = "openai/clip-vit-base-patch32",
        dataset_version: str = "mmsd-v2",
        max_length: int = 77,
        train_batch_size: int = 32,
        val_batch_size: int = 32,
        test_batch_size: int = 32,
        num_workers: int = 19,
    ) -> None:
        super().__init__()
        self.clip_ckpt_name = clip_ckpt_name
        self.dataset_version = dataset_version
        self.train_batch_size = train_batch_size
        self.val_batch_size = val_batch_size
        self.test_batch_size = test_batch_size
        self.num_workers = num_workers
        self.max_length = max_length

    def setup(self, stage: str) -> None:
        processor = CLIPProcessor.from_pretrained(self.clip_ckpt_name)

        def preprocess(example):
            inputs = processor(
                text=example["text"],
                images=example["image"],
                return_tensors="pt",
                padding="max_length",
                truncation=True,
                max_length=self.max_length,
            )

            return {
                "pixel_values": inputs["pixel_values"],
                "input_ids": inputs["input_ids"],
                "attention_mask": inputs["attention_mask"],
                "label": example["label"],
            }

        self.raw_dataset = cast(
            Dataset,
            load_dataset("coderchen01/MMSD2.0", name=self.dataset_version),
        )
        self.dataset = self.raw_dataset.map(
            preprocess,
            batched=True,
            remove_columns=["text", "image"],
        )

    def train_dataloader(self) -> DataLoader:
        return DataLoader(
            self.dataset["train"],
            batch_size=self.train_batch_size,
            shuffle=True,
            num_workers=self.num_workers,
        )

    def val_dataloader(self) -> DataLoader:
        return DataLoader(
            self.dataset["validation"],
            batch_size=self.val_batch_size,
            num_workers=self.num_workers,
        )

    def test_dataloader(self) -> DataLoader:
        return DataLoader(
            self.dataset["test"],
            batch_size=self.test_batch_size,
            num_workers=self.num_workers,
        )