Adding training data module

.gitignore

@@ -0,0 +1 @@
+.vscode/

src/config.py

@@ -0,0 +1,40 @@
+import pydantic
+
+
+class DataConfig(pydantic.BaseModel):
+    buffer_size: int = 1000
+    data_len: int = 100000
+    train_len: int = 90000
+    small_dataset: str = "laion/220k-gpt4vision-captions-from-livis"
+    large_dataset: str = "laion/laion400m"
+    dataset: str = small_dataset
+
+
+class ModelConfig(pydantic.BaseModel):
+    text_model: str = "microsoft/xtremedistil-l6-h256-uncased"  # 51 mb
+    vision_model: str = "edgenext_small"  # 20 mb
+    projection_layers: int = 3
+    embed_dim: int = 256
+    transformer_embed_dim: int = 768
+    max_len: int = 77  # maximum length of text in CLIP
+    cls_type: bool = True
+    freeze_vision_base: bool = False
+    freeze_text_base: bool = False
+
+
+class TrainerConfig(pydantic.BaseModel):
+    epochs: int = 20
+    batch_size: int = 256
+    learning_rate: float = 5e-4
+    accumulate_grad_batches: int = 1
+    temperature: float = 1.0
+    vision_freeze_layers: int = 2
+    lambda_1: float = 1.0
+    lambda_2: float = 1.0
+
+    val_check_interval: int = 1000
+
+    run_openai_clip: bool = False
+
+    model_config: ModelConfig = ModelConfig()
+    data_config: DataConfig = DataConfig()

src/data.py

@@ -0,0 +1,109 @@
+import io
+import multiprocessing as mp
+from typing import Optional, Union
+
+import datasets
+from PIL import Image
+import requests
+import torch
+from torch.utils.data import Dataset, DataLoader
+from torchvision import transforms
+from transformers import AutoTokenizer
+
+from src import config
+
+
+class Tokenizer:
+    def __init__(self, model_name: str, max_len: int) -> None:
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.max_len = max_len
+
+    def __call__(self, x: Union[str, list[str]]) -> dict[str, torch.LongTensor]:
+        return self.tokenizer(
+            x, max_length=self.max_len, truncation=True, padding=True, return_tensors="pt"
+        )
+
+    def decode(self, x: dict[str, torch.LongTensor]) -> list[str]:
+        return [
+            self.tokenizer.decode(sentence[:sentence_len])
+            for sentence, sentence_len in zip(x["input_ids"], x["attention_mask"].sum(axis=-1))
+        ]
+
+
+def _get_image_and_caption(item: dict[str, str]) -> Optional[tuple[Image.Image, str]]:
+    image_url = item["url"]
+    caption = item["caption"]
+    try:
+        response = requests.get(image_url, timeout=1)
+        response.raise_for_status()  # Raise HTTPError for bad responses (4xx and 5xx)
+        image = Image.open(io.BytesIO(response.content))
+        return image, caption
+    except (requests.RequestException, IOError):
+        return None
+
+
+class CollateFn:
+    def __init__(self, tokenizer: Tokenizer, transform: transforms.Compose):
+        self.tokenizer = tokenizer
+        self.transform = transform
+
+    def __call__(
+        self, batch: list[Optional[tuple[str, torch.FloatTensor]]]
+    ) -> tuple[dict[str, torch.LongTensor], torch.FloatTensor]:
+        filtered_batch = [data for data in map(_get_image_and_caption, batch) if data is not None]
+        x, y = zip(*filtered_batch)
+        tokenized_text = self.tokenizer(list(x))
+        return tokenized_text, torch.stack([self.transform(image) for image in y])
+
+
+def _get_dataloaders(
+    train_ds: Dataset,
+    valid_ds: Dataset,
+    training_config: config.TrainerConfig,
+    collate_fn: CollateFn,
+) -> tuple[DataLoader, DataLoader]:
+    common_params = {
+        "batch_size": training_config.batch_size,
+        "pin_memory": True,
+        "num_workers": mp.cpu_count(),
+        "collate_fn": collate_fn,
+    }
+    train_loader = DataLoader(
+        train_ds,
+        shuffle=True,
+        drop_last=True,
+        **common_params,
+    )
+    valid_loader = DataLoader(
+        valid_ds,
+        shuffle=False,
+        drop_last=False,
+        **common_params,
+    )
+    return train_loader, valid_loader
+
+
+def get_dataset(
+    transform: transforms.Compose,
+    tokenizer: Tokenizer,
+    hyper_parameters: config.TrainerConfig,
+    num_workers: int,
+) -> tuple[DataLoader, DataLoader]:
+    dataset = datasets.load_dataset(
+        hyper_parameters.data_config.dataset, split="train", streaming=True
+    )
+    full_dataset = dataset.shuffle(
+        seed=42, buffer_size=hyper_parameters.data_config.buffer_size
+    ).take(hyper_parameters.data_config.data_len)
+    train_dataset = full_dataset.take(hyper_parameters.data_config.train_len)
+    valid_dataset = full_dataset.skip(hyper_parameters.data_config.train_len)
+
+    collate_fn = CollateFn(tokenizer, transform)
+
+    return _get_dataloaders(
+        train_ds=train_dataset,
+        valid_ds=valid_dataset,
+        training_config=hyper_parameters,
+        collate_fn=collate_fn,
+        num_workers=num_workers,
+    )