Delete main.py

main.py

@@ -1,116 +0,0 @@
-from datasets import load_dataset
-from linear_mapping import LinearMapping, LinearMappingProcessor, LinearMappingConfig, Transform
-import torch
-from torchvision.io import ImageReadMode, read_image
-from transformers import Trainer, TrainingArguments
-import os
-from PIL import Image
-os.environ["WANDB_DISABLED"] = "true"
-
-DATA_DIR = os.path.join(os.getcwd(), "coco")
-CAPTION_COLUMN = "caption"
-IMAGE_COLUMN = "image_path"
-
-
-def main():
-    ds = load_dataset("ydshieh/coco_dataset_script", "2017", DATA_DIR)
-    config = LinearMappingConfig()
-    processor = LinearMappingProcessor(config)
-
-    def collate_fn(batch):
-        return {
-            'pixel_values': torch.stack([x['pixel_values'] for x in batch]),
-            'input_ids': torch.tensor([x['input_ids'] for x in batch], dtype=torch.long),
-            'attention_mask': torch.stack([x["attention_mask"] for x in batch]),
-        }
-
-    def tokenize_fn(examples):
-        texts = list(examples[CAPTION_COLUMN])
-        if config.add_image_token:
-            texts = list(processor.tokenizer.cls_token + text for text in texts)
-        inputs = processor.tokenizer(
-            texts, padding="max_length", max_length=77,
-            return_tensors="pt", truncation=True
-        )
-        examples["input_ids"] = inputs.input_ids
-        examples["attention_mask"] = inputs.attention_mask
-        return examples
-
-    image_transformations = Transform(
-        config.image_resize,
-        [0.48145466, 0.4578275, 0.40821073],
-        [0.26862954, 0.26130258, 0.27577711]
-    )
-    image_transformations = torch.jit.script(image_transformations)
-
-    def transform_images(examples):
-        images = [read_image(image_file, mode=ImageReadMode.RGB) for image_file in examples[IMAGE_COLUMN]]
-        examples["pixel_values"] = [image_transformations(image) for image in images]
-
-        examples["attention_mask"] = torch.cat([
-            torch.ones(len(images), config.prefix_length),
-            torch.tensor(examples["attention_mask"])
-        ], dim=1).to(dtype=torch.long)
-        return examples
-
-    def preprocess_fn(examples):
-
-        texts = list(examples[CAPTION_COLUMN])
-
-        images = [read_image(image_file, mode=ImageReadMode.RGB) for image_file in examples[IMAGE_COLUMN]]
-        inputs = processor(
-            texts=texts, images=images, padding="max_length", truncation=True, max_length=77, return_tensors="pt"
-        )
-        return inputs
-
-    def filter_corrupt_images(examples):
-        """remove problematic images"""
-        valid_images = []
-        for image_file in examples[IMAGE_COLUMN]:
-            try:
-                Image.open(image_file)
-                valid_images.append(True)
-            except Exception:
-                valid_images.append(False)
-        return valid_images
-
-    train_dataset = ds["train"]
-
-    train_dataset = train_dataset.filter(
-        function=filter_corrupt_images,
-        batched=True
-    )
-    train_dataset = train_dataset.map(
-        function=tokenize_fn,
-        batched=True,
-        remove_columns=[col for col in train_dataset.column_names if col != IMAGE_COLUMN and col != CAPTION_COLUMN],
-        load_from_cache_file=True
-    )
-    train_dataset.set_transform(transform_images)
-
-    training_args = TrainingArguments(
-        learning_rate=5e-4,
-        lr_scheduler_type='cosine',
-        output_dir='clip-gpt2-image-captioner',
-        do_train=True,
-        logging_steps=50,
-        num_train_epochs=5,
-        logging_dir='runs',
-        remove_unused_columns=False,
-        max_grad_norm=1.0,
-        per_device_train_batch_size=16,
-        save_total_limit=3,
-        warmup_steps=500
-    )
-    model = LinearMapping(config)
-    trainer = Trainer(
-        model=model,
-        args=training_args,
-        train_dataset=train_dataset,
-        data_collator=collate_fn
-    )
-    trainer.train()
-
-
-if __name__ == '__main__':
-    main()