support download allava4v dataset

Author: jiapingW

datasets/README.md

@@ -0,0 +1,5 @@
+## Store Comprehensive Datasets Download Scripts
+
+| DatasetName | Github | Huggingface | command |
+| -------- | -------- | -------- | -------- |
+| ALLaVA-4V | [link](https://github.com/FreedomIntelligence/ALLaVA) | [link](https://huggingface.co/datasets/FreedomIntelligence/ALLaVA-4V) | download_laion.sh |

datasets/download_laion.sh

@@ -0,0 +1,36 @@
+
+
+laion_root="allava_laion"
+
+mkdir $laion_root
+cd $laion_root
+
+
+# 1. download annotation files
+## 1.1 caption
+wget -c -O ALLaVA-Caption-LAION-4V.json https://huggingface.co/datasets/FreedomIntelligence/ALLaVA-4V/resolve/main/allava_laion/ALLaVA-Caption-LAION-4V.json?download=true
+
+## 1.2 instruction
+wget -c -O ALLaVA-Instruct-LAION-4V.json https://huggingface.co/datasets/FreedomIntelligence/ALLaVA-4V/resolve/main/allava_laion/ALLaVA-Instruct-LAION-4V.json?download=true
+
+
+# 2. download and upzip images
+mkdir image_chunks
+
+## 2.1 download
+for ((i=0; i<10; i++))
+do
+    wget -c -O image_chunks/images_$i.zip https://huggingface.co/datasets/FreedomIntelligence/ALLaVA-4V/resolve/main/allava_laion/image_chunks/images_$i.zip?download=true &
+done
+
+mkdir -p images/
+wait
+
+## 2.2 unzip
+for ((i=0; i<10; i++))
+do
+    unzip -j -o image_chunks/images_$i.zip -d images/ & # wait patiently, it takes a while...
+done
+
+wait
+echo "All done!"

scripts/prepare_data.py

@@ -1,12 +1,14 @@
 import argparse
 import json
 import os
+import subprocess
 from pathlib import Path
 from typing import Dict, Tuple
 
-from datasets import concatenate_datasets, load_dataset
 from tqdm import tqdm
 
+from datasets import concatenate_datasets, config, load_dataset
+
 """
 This script will convert the ultrachat/sharegpt dataset to the following schema in jsonl format:
 {
@@ -88,7 +90,49 @@ def parse_args():
     return parser.parse_args()
 
 
-def process_ultrachat_row(row: Dict) -> Tuple[Dict, int]:
+def get_cache_dir(dataset_name):
+    cache_dir = None
+    if dataset_name == "sharegpt4v":
+        raise Exception("Don't Support Download sharegpt4v.")
+    elif dataset_name == "allava4v":
+        cache_dir = os.path.join(
+            config.HF_DATASETS_CACHE, "FreedomIntelligence", "ALLaVA"
+        )
+    else:
+        raise Exception(f"Don't support {dataset_name}")
+    return cache_dir
+
+
+def download_vlm_dataset(dataset_name: str) -> None:
+    """Download VLM's dataset such as sharegpt4v and allava4v"""
+    if dataset_name == "sharegpt4v":
+        raise Exception("Don't Support Download sharegpt4v.")
+    elif dataset_name == "allava4v":
+        cache_dir = get_cache_dir(dataset_name)
+        os.makedirs(cache_dir, exist_ok=True)
+        script_path = os.path.join(
+            os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
+            "datasets",
+            "download_laion.sh",
+        )
+        os.chmod(script_path, 0o755)
+        if not os.path.exists(os.path.join(cache_dir, "allava_laion")):
+            result = subprocess.run(
+                ["bash", script_path],
+                cwd=cache_dir,
+                capture_output=True,
+                text=True,
+            )
+            if result.returncode != 0:
+                raise RuntimeError(f"Download image dataset failed: {result.stderr}")
+            print("##### allava4v dataset Download Complete #####")
+        else:
+            print("##### allava4v dataset has existed.")
+    else:
+        raise Exception(f"Don't support {dataset_name}")
+
+
+def process_ultrachat_row(row: Dict, dataset_name: str = None) -> Tuple[Dict, int]:
     """Process a row from the ultrachat dataset.
 
     The function expects a row with the following schema:
@@ -110,7 +154,7 @@ def process_ultrachat_row(row: Dict) -> Tuple[Dict, int]:
     return row, 0
 
 
-def process_sharegpt_row(row: Dict) -> Tuple[Dict, int]:
+def process_sharegpt_row(row: Dict, dataset_name: str = None) -> Tuple[Dict, int]:
     """
     sharegpt dataset schema:
     {
@@ -138,7 +182,7 @@ def process_sharegpt_row(row: Dict) -> Tuple[Dict, int]:
     return row, skipped_count
 
 
-def process_sharegpt4v_row(row) -> Dict:
+def process_sharegpt4v_row(row, dataset_name: str = None) -> Dict:
     """
     sharegpt4v dataset schema:
     {
@@ -153,8 +197,9 @@ def process_sharegpt4v_row(row) -> Dict:
         ]
     }
     """
+    cache_dir = get_cache_dir(dataset_name)
     conversations = row["conversations"]
-    image = f'FreedomIntelligence/ALLaVA-4V/{row["image"]}'
+    image = os.path.join(cache_dir, f"{row["image"]}")
     if not os.path.exists(image):
         print(f"Image path {image} does not exist, skipping this sample.")
         return None, None
@@ -194,7 +239,7 @@ def process_and_save_ds(train_ds, test_ds, output_path, proc_fn, dataset_name):
     with open(train_output_jsonl_path, "w") as f:
         for item in tqdm(train_ds, desc=f"Processing {dataset_name} dataset"):
             if proc_fn is not None:
-                row, skipped_count = proc_fn(item)
+                row, skipped_count = proc_fn(item, dataset_name)
                 if row is None:
                     continue
                 total_skipped_count += skipped_count
@@ -207,7 +252,7 @@ def process_and_save_ds(train_ds, test_ds, output_path, proc_fn, dataset_name):
         with open(test_output_jsonl_path, "w") as f:
             for item in tqdm(test_ds, desc=f"Processing {dataset_name} test dataset"):
                 if proc_fn is not None:
-                    row, skipped_count = proc_fn(item)
+                    row, skipped_count = proc_fn(item, dataset_name)
                     if row is None:
                         continue
                     total_skipped_count += skipped_count
@@ -292,11 +337,14 @@ def main():
         proc_fn = process_sharegpt_row
     elif args.dataset == "sharegpt4v":
         ds = load_dataset("Lin-Chen/ShareGPT4V", "ShareGPT4V")["train"]
+        raise Exception("Not supported sharegpt4v now")
+        download_vlm_dataset(args.dataset)
         proc_fn = process_sharegpt4v_row
     elif args.dataset == "allava4v":
         ds = load_dataset("FreedomIntelligence/ALLaVA-4V", name="allava_laion")[
             "instruct"
         ]
+        download_vlm_dataset(args.dataset)
         proc_fn = process_sharegpt4v_row
     elif args.dataset == "opc":
         if args.opc_subset == "all":
@@ -318,7 +366,6 @@ def main():
         raise ValueError(
             f"This script only supports ultrachat, sharegpt, sharegpt4v, allava4v, opc, and perfect-blend-gptoss-20B datasets for demo purpose, if you wish to use other datasets, please modify this script."
         )
-
     # filter and split dataset
     if args.sample_size is not None and args.sample_size < len(ds):
         ds = ds.select(range(args.sample_size))

scripts/prepare_hidden_states.py

@@ -43,10 +43,10 @@
 
 import torch
 import torch.distributed as dist
-from datasets import load_dataset
 from tqdm import tqdm
 from transformers import AutoConfig, AutoProcessor, AutoTokenizer
 
+from datasets import load_dataset
 from specforge.args import SGLangBackendArgs
 from specforge.data import build_eagle3_dataset, prepare_dp_dataloaders
 from specforge.distributed import (

scripts/train_eagle3.py

@@ -10,14 +10,14 @@
 import torch.distributed as dist
 import torch.nn as nn
 from accelerate.utils import set_seed
-from datasets import load_dataset
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from torch.distributed.fsdp import MixedPrecision, ShardingStrategy, StateDictType
 from torch.optim import Optimizer
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 from transformers import AutoProcessor, AutoTokenizer
 
+from datasets import load_dataset
 from specforge import (
     AutoDraftModelConfig,
     AutoEagle3DraftModel,

specforge/data/preprocessing.py

@@ -27,10 +27,11 @@
 from typing import Dict, List, Optional, Tuple, Union
 
 import torch
-from datasets import Dataset as HFDataset
 from tqdm import tqdm
 from transformers import ImageProcessingMixin, PreTrainedTokenizer
 
+from datasets import Dataset as HFDataset
+
 try:
     from qwen_vl_utils import process_vision_info
 

specforge/data/utils.py

@@ -23,10 +23,13 @@
 
 import torch
 import torch.distributed as dist
-from datasets import Dataset
 from torch.utils.data import DataLoader, DistributedSampler
 
+<<<<<<< HEAD
 from specforge.distributed import get_draft_sp_group
+=======
+from datasets import Dataset
+>>>>>>> 9837a17 (support download allava4v dataset)
 
 
 class DataCollatorWithPadding: