Fix training bug

.gitignore

@@ -2,6 +2,8 @@
 
 wandb/
 
+checkpoints/
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

configs/dataset_config.py

@@ -51,14 +51,10 @@ visual_datasets = [
 language_datasets = [
     dict(
         type="dolly",
-        ann_paths=[
-            "data/dolly/databricks-dolly-15k.jsonl",
-        ],
+        ann_path="data/dolly/databricks-dolly-15k.jsonl",
     ),
     dict(
         type="alpaca_gpt4",
-        ann_paths=[
-            "data/alpaca_gpt4/alpaca_gpt4_data.json",
-        ],
+        ann_path="data/alpaca_gpt4/alpaca_gpt4_data.json",
     ),
 ]

mmgpt/models/builder.py

@@ -1,5 +1,6 @@
 from .open_flamingo import create_model_and_transforms as create_open_flamingo_model_and_transforms
-
+import torch.nn as nn
+from transformers import LlamaTokenizer, LlamaForCausalLM
 
 def create_model_and_transforms(
     model_name: str,
@@ -24,3 +25,50 @@ def create_model_and_transforms(
     # TODO: support BLIP2
     else:
         raise ValueError(f"Unknown model name: {model_name}")
+
+# only for debugging
+def create_toy_model_and_transforms(
+    model_name: str,
+    clip_vision_encoder_path: str,
+    clip_vision_encoder_pretrained: str,
+    lang_encoder_path: str,
+    tokenizer_path: str,
+    tuning_config,
+    pretrained_model_path,
+    **kwargs,
+):
+    print("init toy vision encoder")
+    import torchvision
+
+    image_processor = torchvision.transforms.Compose(
+        [
+            torchvision.transforms.Resize((224, 224)),
+            torchvision.transforms.ToTensor(),
+        ]
+    )
+    print("init tokenizer")
+    text_tokenizer = LlamaTokenizer.from_pretrained(tokenizer_path)
+    # add Flamingo special tokens to the tokenizer
+    text_tokenizer.add_special_tokens({"additional_special_tokens": ["<|endofchunk|>", "<image>"]})
+    if text_tokenizer.pad_token is None:
+        # Issue: GPT models don't have a pad token, which we use to
+        # modify labels for the loss.
+        text_tokenizer.add_special_tokens({"pad_token": "<PAD>"})
+
+    class ToyModel(nn.Module):
+        def __init__(self, *args, **kwargs):
+            super().__init__()
+            self.input_embeddings = nn.Embedding(38000, 512)
+            self.layer = nn.Linear(512, 512)
+            self.config = {"hidden_size": 512}
+
+        def forward(self, lang_x, **kwargs):
+            x = self.input_embeddings(lang_x)
+            x = self.layer(x)
+            loss = x.sum()
+
+            return (loss,)
+
+    model = ToyModel()
+
+    return model, image_processor, text_tokenizer

mmgpt/train/instruction_finetune.py

@@ -21,6 +21,7 @@ from transformers import (
 )
 
 from mmgpt import create_model_and_transforms
+from mmgpt.models.builder import create_toy_model_and_transforms
 from mmgpt.datasets import InfiniteSampler, build_dataset
 from mmgpt.train.distributed import init_distributed_device, world_info_from_env
 from mmgpt.train.train_utils import AverageMeter, get_autocast, get_cast_dtype, get_checkpoint
@@ -185,7 +186,7 @@ def main():
     )
 
     # build language dataset and dataloader for multi-modality training
-    if dataset_config.get('language_datasets') is not None and len(args.language_datasets) > 0:
+    if dataset_config.get('language_datasets') is not None and len(dataset_config.language_datasets) > 0:
         lang_dataset = build_dataset(
             dataset_config=dataset_config.language_datasets,
             tokenizer=tokenizer,