Merge branch 'train-baseline'

.gitignore

@@ -1 +1,4 @@
-.venv
+.venv
+dataset
+output
+big_vision_repo

Readme.md

@@ -0,0 +1,19 @@
+# Dataset Structure
+
+/custom_vqa_project/
+│
+├── /dataset/
+│   ├── /images/
+│   │   ├── train/
+│   │   │   ├── image1.jpg
+│   │   │   ├── image2.jpg
+│   │   └── val/
+│   │       ├── image3.jpg
+│   │       └── image4.jpg
+│   ├── train.json  # Metadata for the training set
+│   └── val.json    # Metadata for the validation set
+│
+├── /scripts/
+│   └── train.py   # Your fine-tuning script
+│
+└── README.md

article_base_train.py

@@ -0,0 +1,186 @@
+import os, time, math
+import pandas as pd
+from datasets import Dataset
+from transformers import PaliGemmaProcessor, PaliGemmaForConditionalGeneration, BitsAndBytesConfig, TrainingArguments, Trainer
+import torch
+from PIL import Image
+from peft import get_peft_model, LoraConfig
+import argparse
+
+
+# Function to load custom dataset from CSV
+def load_custom_dataset_from_csv(csv_file, image_folder):
+    # Load CSV data using pandas
+    data = pd.read_csv(csv_file)
+    
+    # Prepare dataset format for Hugging Face
+    questions = data['question'].tolist()
+    images = [os.path.join(image_folder, img) for img in data['image'].tolist()]
+    answers = data['answer'].tolist()
+    
+    # Create a Hugging Face dataset from the loaded CSV
+    return Dataset.from_dict({
+        'question': questions,
+        'image': images,
+        'answer': answers
+    })
+
+
+# Function to load custom dataset from Parquet
+def load_custom_dataset_from_parquet(parquet_file, image_folder):
+    # Load Parquet data using pandas
+    data = pd.read_parquet(parquet_file)
+    
+    # Prepare dataset format for Hugging Face
+    questions = data['question'].tolist()
+    images = [os.path.join(image_folder, img) for img in data['image'].tolist()]
+    answers = data['answer'].tolist()
+    
+    # Create a Hugging Face dataset from the loaded Parquet
+    return Dataset.from_dict({
+        'question': questions,
+        'image': images,
+        'answer': answers
+    })
+
+
+# Choose the appropriate loader based on metadata_type argument
+def load_dataset_by_type(metadata_type, dataset_dir, image_folder):
+    if metadata_type == "csv":
+        return load_custom_dataset_from_csv(
+            os.path.join(dataset_dir, 'train_samples.csv'),
+            image_folder
+        )
+    elif metadata_type == "parquet":
+        return load_custom_dataset_from_parquet(
+            os.path.join(dataset_dir, 'train.parquet'),
+            image_folder
+        )
+    else:
+        raise ValueError("Unsupported metadata type. Use 'csv' or 'parquet'.")
+
+
+def load_model_and_args(use_qlora, model_id, device, output_dir):
+    if use_qlora:
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.bfloat16
+        )
+        lora_config = LoraConfig(
+            r=8,
+            target_modules=["q_proj", "o_proj", "k_proj", "v_proj", "gate_proj", "up_proj", "down_proj"],
+            task_type="CAUSAL_LM"
+        )
+        
+        model = PaliGemmaForConditionalGeneration.from_pretrained(model_id, quantization_config=bnb_config, device_map={"": 0})
+        model = get_peft_model(model, lora_config)
+        model.print_trainable_parameters()
+        
+        # TODO: Customize training setting
+        args = TrainingArguments(
+            output_dir=os.path.join(output_dir, f"{math.floor(time.time())}"),
+            num_train_epochs=2,
+            remove_unused_columns=False,
+            per_device_train_batch_size=1,
+            gradient_accumulation_steps=4,
+            warmup_steps=2,
+            learning_rate=2e-5,
+            weight_decay=1e-6,
+            logging_steps=100,
+            optim="adamw_hf",
+            save_strategy="steps",
+            save_steps=1000,
+            save_total_limit=1,
+            bf16=True,
+            report_to=["tensorboard"],
+            dataloader_pin_memory=False
+        )
+        
+        return model, args
+    else:
+        model = PaliGemmaForConditionalGeneration.from_pretrained(model_id, torch_dtype=torch.bfloat16).to(device)
+        for param in model.vision_tower.parameters():
+            param.requires_grad = False
+
+        for param in model.multi_modal_projector.parameters():
+            param.requires_grad = True
+        
+        # TODO: Customize training setting
+        args = TrainingArguments(
+            output_dir=os.path.join(output_dir, f"{math.floor(time.time())}"),
+            num_train_epochs=2,
+            remove_unused_columns=False,
+            per_device_train_batch_size=4,
+            gradient_accumulation_steps=4,
+            warmup_steps=2,
+            learning_rate=2e-5,
+            weight_decay=1e-6,
+            logging_steps=100,
+            optim="paged_adamw_8bit",
+            save_strategy="steps",
+            save_steps=1000,
+            save_total_limit=1,
+            bf16=True,
+            report_to=["tensorboard"],
+            dataloader_pin_memory=False
+        )
+
+        return model, args
+
+
+# Main training function
+def main(args):
+    dataset_dir = args.dataset_dir
+    model_id = args.model_id
+    output_dir = args.output_dir
+    metadata_type = args.metadata_type
+    
+    # Load custom datasetsㄴ
+    # dataset = load_custom_dataset_from_csv(
+    #     os.path.join(dataset_dir, 'train_samples.csv'),
+    #     os.path.join(dataset_dir, 'images/train')) # TODO: change to appropriate path
+    dataset = load_dataset_by_type(metadata_type, dataset_dir, os.path.join(dataset_dir, 'images/train'))
+    train_val_split = dataset.train_test_split(test_size=0.1)
+    
+    train_ds = train_val_split['train']
+    val_ds = train_val_split['test']
+
+    processor = PaliGemmaProcessor.from_pretrained(model_id)
+    device = "cuda"
+
+    model, args = load_model_and_args(args.use_qlora, model_id, device, output_dir)
+
+    # Custom collate function
+    def collate_fn(examples):
+        texts = [example["question"] for example in examples]
+        labels = [example['answer'] for example in examples]
+        images = [Image.open(example['image']).convert("RGB") for example in examples]
+        tokens = processor(text=texts, images=images, suffix=labels, return_tensors="pt", padding="longest")
+        tokens = tokens.to(torch.bfloat16).to(device)
+        return tokens
+    
+    trainer = Trainer(
+        model=model,
+        train_dataset=train_ds,
+        eval_dataset=val_ds,
+        data_collator=collate_fn,
+        args=args
+    )
+    
+    trainer.train()
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Train a model with custom dataset")
+    parser.add_argument('--dataset_dir', type=str, default='./dataset', help='Path to the folder containing the images')
+    parser.add_argument('--model_id', type=str, default='google/paligemma-3b-pt-224', help='Model ID to use for training')
+    parser.add_argument('--output_dir', type=str, default='./output', help='Directory to save the output')
+    parser.add_argument('--use_qlora', type=bool, default=False, help='Use QLoRA for training')
+    parser.add_argument('--metadata_type', type=str, default='parquet', choices=['csv', 'parquet'], help='Metadata format (csv or parquet)')
+    return parser.parse_args()
+
+    
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)

article_base_train_test.py

@@ -1,80 +0,0 @@
-from huggingface_hub import notebook_login
-from datasets import load_dataset
-from transformers import PaliGemmaProcessor, PaliGemmaForConditionalGeneration, BitsAndBytesConfig, TrainingArguments, Trainer
-import torch
-from peft import get_peft_model, LoraConfig
-
-
-def main():
-  ds = load_dataset('HuggingFaceM4/VQAv2', split="train", trust_remote_code=True) 
-  cols_remove = ["question_type", "answers", "answer_type", "image_id", "question_id"] 
-  ds = ds.remove_columns(cols_remove)
-  ds = ds.train_test_split(test_size=0.1)
-  train_ds = ds["train"]
-  val_ds = ds["test"]
-  
-  model_id = "google/paligemma-3b-pt-224"
-  processor = PaliGemmaProcessor.from_pretrained(model_id)
-  image_token = processor.tokenizer.convert_tokens_to_ids("<image>")
-  device = "cuda"
-
-  bnb_config = BitsAndBytesConfig(
-        load_in_4bit=True,
-        bnb_4bit_quant_type="nf4",
-        bnb_4bit_compute_type=torch.bfloat16
-  )
-  lora_config = LoraConfig(
-      r=8, 
-      target_modules=["q_proj", "o_proj", "k_proj", "v_proj", "gate_proj", "up_proj", "down_proj"],
-      task_type="CAUSAL_LM",
-  )
-  model = PaliGemmaForConditionalGeneration.from_pretrained(model_id, quantization_config=bnb_config, device_map={"":0})
-  model = get_peft_model(model, lora_config)
-  model.print_trainable_parameters()
-  #trainable params: 11,298,816 || all params: 2,934,634,224 || trainable%: 0.38501616002417344
-  
-  args=TrainingArguments(
-            num_train_epochs=2,
-            remove_unused_columns=False,
-            per_device_train_batch_size=16,
-            gradient_accumulation_steps=4,
-            warmup_steps=2,
-            learning_rate=2e-5,
-            weight_decay=1e-6,
-            adam_beta2=0.999,
-            logging_steps=100,
-            # optim="adamw_hf",
-            optim="paged_adamw_8bit", # for QLoRA
-            save_strategy="steps",
-            save_steps=1000,
-            push_to_hub=True,
-            save_total_limit=1,
-            bf16=True,
-            report_to=["tensorboard"],
-            dataloader_pin_memory=False
-        )
-  
-  def collate_fn(examples):
-    texts = ["answer " + example["question"] for example in examples]
-    labels= [example['multiple_choice_answer'] for example in examples] # 우리는 label 이 필요 없을듯?
-    images = [example["image"].convert("RGB") for example in examples]
-    tokens = processor(text=texts, images=images, suffix=labels,
-                      return_tensors="pt", padding="longest")
-
-    tokens = tokens.to(torch.bfloat16).to(device)
-    return tokens
-  
-  trainer = Trainer(
-    model=model,
-    train_dataset=train_ds,
-    eval_dataset=val_ds,
-    data_collator=collate_fn,
-    args=args
-  )
-  
-  trainer.train()
-
-
-if __name__ == "__main__":
-  notebook_login()
-  main()

article_base_tutorial.ipynb

@@ -254,7 +254,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Inference for test"
+    "Save Model"
    ]
   },
   {
@@ -262,7 +262,20 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": []
+   "source": [
+    "save_path = \"./fine_tuned_model\"\n",
+    "model.save_pretrained(save_path)\n",
+    "processor.save_pretrained(save_path)\n",
+    "\n",
+    "print(f\"Model saved locally at {save_path}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Inference for test"
+   ]
   }
  ],
  "metadata": {

requirements.txt

@@ -0,0 +1,74 @@
+accelerate==0.34.2
+aiohappyeyeballs==2.4.2
+aiohttp==3.10.6
+aiosignal==1.3.1
+appnope==0.1.4
+asttokens==2.4.1
+attrs==24.2.0
+certifi==2024.8.30
+charset-normalizer==3.3.2
+comm==0.2.2
+datasets==3.0.1
+debugpy==1.8.6
+decorator==5.1.1
+dill==0.3.8
+executing==2.1.0
+filelock==3.16.1
+frozenlist==1.4.1
+fsspec==2024.6.1
+huggingface-hub==0.25.1
+idna==3.10
+ipykernel==6.29.5
+ipython==8.27.0
+ipywidgets==8.1.5
+jedi==0.19.1
+Jinja2==3.1.4
+jupyter_client==8.6.3
+jupyter_core==5.7.2
+jupyterlab_widgets==3.0.13
+MarkupSafe==2.1.5
+matplotlib-inline==0.1.7
+mpmath==1.3.0
+multidict==6.1.0
+multiprocess==0.70.16
+nest-asyncio==1.6.0
+networkx==3.3
+numpy==2.1.1
+packaging==24.1
+pandas==2.2.3
+parso==0.8.4
+peft==0.13.0
+pexpect==4.9.0
+pillow==10.4.0
+pip==24.0
+platformdirs==4.3.6
+prompt_toolkit==3.0.48
+psutil==6.0.0
+ptyprocess==0.7.0
+pure_eval==0.2.3
+pyarrow==17.0.0
+Pygments==2.18.0
+python-dateutil==2.9.0.post0
+pytz==2024.2
+PyYAML==6.0.2
+pyzmq==26.2.0
+regex==2024.9.11
+requests==2.32.3
+safetensors==0.4.5
+setuptools==75.1.0
+six==1.16.0
+stack-data==0.6.3
+sympy==1.13.3
+tokenizers==0.20.0
+torch==2.4.1
+tornado==6.4.1
+tqdm==4.66.5
+traitlets==5.14.3
+transformers==4.45.1
+typing_extensions==4.12.2
+tzdata==2024.2
+urllib3==2.2.3
+wcwidth==0.2.13
+widgetsnbextension==4.0.13
+xxhash==3.5.0
+yarl==1.13.0

test_inference.py

@@ -0,0 +1,24 @@
+from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
+from PIL import Image
+
+
+def main():
+  model_id = "google/paligemma-3b-pt-224"
+  # model_path = "output/1727488022/checkpoint-112"
+  model_path = "output/1727490265/checkpoint-450"
+  model = PaliGemmaForConditionalGeneration.from_pretrained(model_path)
+  processor = AutoProcessor.from_pretrained(model_id)
+
+  # prompt = "Analyze image from a critic's point of view."
+  prompt = "Please construct a formal analysis paragraph that is coherent and focuses solely on visual characteristic."
+  image_file_path = "dataset/images/manual_test/starry_night.jpg"
+  raw_image = Image.open(image_file_path)
+  inputs = processor(prompt, raw_image, return_tensors="pt")
+  output = model.generate(**inputs, max_new_tokens=20)
+
+  # Starry Night
+  print("Response: ", processor.decode(output[0], skip_special_tokens=True)[len(prompt):])
+  
+ 
+if __name__ == "__main__":
+  main()