Upload 12 files

.gitattributes

@@ -33,3 +33,9 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+test_image_0.png filter=lfs diff=lfs merge=lfs -text
+test_image_1.png filter=lfs diff=lfs merge=lfs -text
+test_image_2.png filter=lfs diff=lfs merge=lfs -text
+test_image_3.png filter=lfs diff=lfs merge=lfs -text
+test_image_4.png filter=lfs diff=lfs merge=lfs -text
+test_image_5.png filter=lfs diff=lfs merge=lfs -text

.python-version

@@ -0,0 +1 @@
+3.13

__init__.py

@@ -0,0 +1,2 @@
+
+import model

checkpoint-epoch=01-loss=0.13.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2190cb9a3ba864e44fd8e28cb57595b043baf5b2ee32b4386c9b2d637164a24e
+size 1774312061

inference.py

@@ -0,0 +1,7 @@
+import model
+import datasets
+from PIL import Image
+vlm = model.ImageNetCaptionModel.load_from_checkpoint('checkpoint-epoch=01-loss=0.13.ckpt')
+
+image = Image.open("test_image_5.png")
+print(vlm.generate(image=image))

model.py

@@ -0,0 +1,336 @@
+from pathlib import Path
+import comet_ml
+import datasets
+import evaluate
+import lightning as L
+import torch
+from timm import create_model, data
+from tokenizers import Tokenizer
+from torch import nn
+from torch.utils.data import DataLoader
+from transformers import (
+    GPT2LMHeadModel,
+)
+from lightning.pytorch.loggers import TensorBoardLogger
+from lightning.pytorch.callbacks import ModelCheckpoint
+
+
+eos_token_id = 50256  # obtained from gpt model
+
+
+class Projection(nn.Module):
+    def __init__(self, in_features, out_features):
+        super().__init__()
+        self.network = nn.Sequential(
+            nn.Linear(in_features, in_features * 3),
+            nn.GELU(),
+            nn.Linear(in_features * 3, out_features),
+        )
+
+    def forward(self, input):
+        return self.network(input)
+
+
+class ImageNetCaptionModel(L.LightningModule):
+    def __init__(self):
+        super().__init__()
+        # backbone model to extract image feature token
+        self.backbone = create_model(
+            "vit_mediumd_patch16_reg4_gap_384", pretrained=True
+        )
+        self.llm = GPT2LMHeadModel.from_pretrained("gpt2")
+
+        self.image_start_token = "<image_start>"
+        self.image_end_token = "<image_end>"
+        self.tokenizer = Tokenizer.from_pretrained("gpt2")
+        self.tokenizer.add_special_tokens(
+            [self.image_start_token, self.image_end_token]
+        )
+        self.image_start_token_id = self.tokenizer.token_to_id(self.image_start_token)
+        self.image_end_token_id = self.tokenizer.token_to_id(self.image_end_token)
+        self.eos_token = eos_token_id
+
+        self.llm.resize_token_embeddings(self.tokenizer.get_vocab_size())
+        self.embedding = self.llm.get_input_embeddings()
+
+        self.projection = Projection(
+            in_features=512, out_features=self.llm.config.hidden_size
+        )
+
+        self.bleu_metric = evaluate.load("bleu")
+        self.meteor_metric = evaluate.load("meteor")
+
+        ## freeze backbone and gpt models.
+        for param in self.backbone.parameters():
+            param.requires_grad = False
+
+        for param in self.llm.parameters():
+            param.requires_grad = True
+
+    def get_tokenizer(self):
+        return self.tokenizer
+
+    def forward(self, image=None, input_caption=None, **kwargs):
+        image_feature = self.backbone.forward_features(image)
+        projection = self.projection(image_feature)
+        input_caption_embedding = self.embedding(input=input_caption)
+
+        # concat start_image_token + projection + end_image_token + input_caption
+        image_start_token, image_end_token = self.get_image_seperation_token(
+            image=image
+        )
+        input_embedding = torch.cat(
+            [image_start_token, projection, image_end_token, input_caption_embedding],
+            dim=1,
+        )
+        attention_mask = torch.ones(
+            input_embedding.size()[:-1], dtype=torch.long, device=image.device
+        )
+
+        labels = torch.full(
+            (input_embedding.size(0), input_embedding.size(1)),
+            -100,
+            dtype=torch.long,
+            device=image.device,
+        )
+        labels[:, projection.size(1) + 2 :] = input_caption  # align text labels
+
+        llm_output = self.llm(
+            inputs_embeds=input_embedding, attention_mask=attention_mask, labels=labels
+        )
+        return llm_output
+
+    def training_step(self, batch, batch_idx):
+        output = self.forward(**batch)
+        self.log("loss", output.loss.item())
+        return output.loss
+
+    def validation_step(self, batch, batch_idx):
+        if batch_idx < 5:
+            pred = self.predict_step(batch=batch, batch_idx=batch_idx)
+            print(
+                "evaluation ",
+                "pred",
+                pred,
+                "original caption",
+                batch["original_caption_enriched"],
+            )
+            bleu = self.bleu_metric.compute(
+                predictions=pred, references=batch["original_caption_enriched"]
+            )
+            self.log("bleu", bleu["bleu"])
+            self.log("precision", bleu["brevity_penalty"])
+            metor = self.meteor_metric.compute(
+                predictions=pred, references=batch["original_caption_enriched"]
+            )
+            print(metor)
+            self.log_dict(metor)
+
+    def get_image_seperation_token(self, image):
+        image_start_embedding = self.embedding(
+            torch.tensor([self.image_start_token_id], device=image.device)
+        )
+        image_end_embedding = self.embedding(
+            torch.tensor([self.image_end_token_id], device=image.device)
+        )
+        image_start_token = image_start_embedding.unsqueeze(0).repeat(len(image), 1, 1)
+        image_end_token = image_end_embedding.unsqueeze(0).repeat(len(image), 1, 1)
+
+        return image_start_token, image_end_token
+
+    def configure_optimizers(self):
+        proj_params = [p for p in self.projection.parameters() if p.requires_grad]
+        llm_params = [p for p in self.llm.parameters() if p.requires_grad]
+
+        optimizer = torch.optim.AdamW(
+            [
+                {"params": proj_params, "lr": 1e-4, "weight_decay": 0.01},
+                {"params": llm_params, "lr": 5e-6, "weight_decay": 0.01},
+            ]
+        )
+
+        return optimizer
+
+    def predict_step(self, batch, batch_idx, dataloader_idx=0):
+        image = batch["image"]
+        image_feature = self.backbone.forward_features(image)
+        projection = self.projection(image_feature)
+
+        image_start_embedding = self.embedding(
+            torch.tensor([self.image_start_token_id], device=image.device)
+        )
+        image_end_embedding = self.embedding(
+            torch.tensor([self.image_end_token_id], device=image.device)
+        )
+        input_start_image_embedding_batch = image_start_embedding.unsqueeze(0).repeat(
+            len(image), 1, 1
+        )
+        input_end_image_embedding_batch = image_end_embedding.unsqueeze(0).repeat(
+            len(image), 1, 1
+        )
+
+        input_embedding = torch.cat(
+            [
+                input_start_image_embedding_batch,
+                projection,
+                input_end_image_embedding_batch,
+            ],
+            dim=1,
+        )
+        attention_mask = torch.ones(
+            input_embedding.size()[:-1], dtype=torch.long, device=image.device
+        )
+
+        outputs = self.llm.generate(
+            inputs_embeds=input_embedding,
+            attention_mask=attention_mask,
+            eos_token_id=0,
+            max_new_tokens=30,
+            do_sample=True,  # add randomness
+            top_p=0.9,  # nucleus sampling
+            temperature=0.7,
+        )
+
+        # Convert tensor to list of lists for decode_batch
+        if outputs.dim() == 2:
+            # outputs is [batch_size, sequence_length], convert to list of lists
+            outputs_list = outputs.tolist()
+        else:
+            # outputs is already a list/sequence
+            outputs_list = outputs
+
+        return self.tokenizer.decode_batch(outputs_list, skip_special_tokens=True)
+    
+    def generate(self, image):
+        data_config = data.resolve_model_data_config(
+        create_model("vit_mediumd_patch16_reg4_gap_384", pretrained=True)
+        )
+        transforms = data.create_transform(**data_config, is_training=False)
+        image = transforms(image)
+        
+        return self.predict_step(batch={"image":image.unsqueeze(0)},batch_idx=0)[0]
+        
+
+
+def collate_fn(batch):
+    collected = {"image": [], "input_caption": [], "original_caption_enriched": []}
+
+    for data in batch:
+        collected["image"].append(torch.tensor(data["image"], dtype=torch.float))
+        collected["input_caption"].append(
+            torch.tensor(data["input_caption"], dtype=torch.long)
+        )
+        collected["original_caption_enriched"].append(data["original_caption_enriched"])
+
+    return {
+        "image": torch.stack(collected["image"], dim=0),
+        "input_caption": torch.stack(collected["input_caption"], dim=0),
+        "original_caption_enriched": collected["original_caption_enriched"],
+    }
+
+
+def agument(tokenizer: Tokenizer):
+    data_config = data.resolve_model_data_config(
+        create_model("vit_mediumd_patch16_reg4_gap_384", pretrained=True)
+    )
+    transforms = data.create_transform(**data_config, is_training=False)
+
+    def transform(data):
+        ids = tokenizer.encode(data["caption_enriched"])
+
+        # Handle sequences based on length
+        if len(ids.ids) <= 59:
+            # For short sequences, just append EOS
+            ids.ids.append(eos_token_id)
+        else:
+            # For long sequences, truncate to 59 tokens and append EOS
+            ids.ids = ids.ids[:59]
+            ids.ids.append(eos_token_id)
+
+            # Pad to exactly 60 tokens
+            ids.ids = ids.ids[:60]  # Ensure we don't exceed 60
+        ids.pad(60)
+
+        decoded = tokenizer.decode(ids.ids, skip_special_tokens=True)
+        print("original", data["caption_enriched"], "decoded", decoded)
+
+        data["input_caption"] = torch.tensor(ids.ids, dtype=torch.long)
+
+        data["original_caption_enriched"] = data["caption_enriched"]
+        data["image"] = transforms(data["image"])
+        return data
+
+    return transform
+
+
+def is_valid_image(example):
+    try:
+        # Try opening the image
+        if example["image"].mode == "RGB":
+            return True
+
+        return False
+    except Exception as e:
+        # ValueError will catch the MAX_TEXT_CHUNK error
+        print("false", example["image"])
+        print("Exception:", e)
+        return False
+
+
+def train(
+    root_path: Path,
+    dataset: datasets.Dataset,
+    num_loader_worker: int = 0,
+    batch_size=16,
+    logger=None,
+):
+    # dataset = datasets.load_dataset("visual-layer/imagenet-1k-vl-enriched", split="validation").shuffle(seed=42).select(range(20000)).train_test_split(test_size=0.1)
+    test_ds = dataset["test"]
+    train_ds = dataset["train"]
+
+    model = ImageNetCaptionModel()
+
+    tokenizer = model.get_tokenizer()
+
+    # Apply transformation to both datasets
+    train_ds = train_ds.filter(is_valid_image)
+    train_ds = train_ds.map(agument(tokenizer=tokenizer))
+
+    test_ds = test_ds.filter(is_valid_image)
+    test_ds = test_ds.map(agument(tokenizer=tokenizer))
+
+    train_data_loader = DataLoader(
+        dataset=train_ds,
+        drop_last=True,
+        batch_size=batch_size,
+        collate_fn=collate_fn,
+        num_workers=num_loader_worker,
+    )
+    evaluation_data_loader = DataLoader(
+        dataset=test_ds,
+        drop_last=True,
+        batch_size=batch_size,
+        collate_fn=collate_fn,
+        num_workers=num_loader_worker,
+    )
+
+    if logger is None:
+        logger = TensorBoardLogger(save_dir=str(root_path), version=1, name="logs")
+    checkpoint_callback = ModelCheckpoint(
+        dirpath=root_path / "checkpoint",
+        filename="checkpoint-{epoch:02d}-{loss:.2f}",
+        every_n_epochs=1,
+        save_top_k=-1,
+    )
+    print("path", root_path)
+    trainer = L.Trainer(
+        logger=logger,
+        max_epochs=2,
+        default_root_dir=root_path,
+        callbacks=[checkpoint_callback],
+    )
+    trainer.fit(
+        model=model,
+        train_dataloaders=train_data_loader,
+        val_dataloaders=evaluation_data_loader,
+    )

pyproject.toml

@@ -0,0 +1,27 @@
+[project]
+name = "imagenet-caption"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.13"
+dependencies = [
+    "comet-ml>=3.52.1",
+    "datasets[vision]>=4.0.0",
+    "evaluate[bleu,meteor]>=0.4.6",
+    "lightning>=2.5.5",
+    "modal>=1.1.4",
+    "nltk>=3.9.1",
+    "numpy>=2.3.3",
+    "pillow>=11.3.0",
+    "tensorboard>=2.20.0",
+    "tensorboardx>=2.6.4",
+    "timm>=1.0.19",
+    "tokenizers>=0.22.0",
+    "torch>=2.8.0",
+    "transformers[torch]>=4.56.1",
+]
+
+[dependency-groups]
+dev = [
+    "ipykernel>=6.30.1",
+]