wehub-resource-sync
81 lines
4.1 KiB
Python
81 lines
4.1 KiB
Python
import os
|
|
import argparse
|
|
import requests
|
|
from tqdm import tqdm
|
|
from typing import List
|
|
|
|
# Base URL for the dataset files
|
|
BASE_URL = "https://huggingface.co/datasets/monology/pile-uncopyrighted/resolve/main"
|
|
VAL_URL = f"{BASE_URL}/val.jsonl.zst" # URL for the validation dataset
|
|
TRAIN_URLS = [f"{BASE_URL}/train/{i:02d}.jsonl.zst" for i in range(65)] # URLs for 65 training files (adjust the range if needed)
|
|
|
|
def download_file(url: str, file_name: str) -> None:
|
|
"""
|
|
Downloads a file from the given URL and saves it with the specified file name.
|
|
Displays a progress bar using tqdm.
|
|
|
|
Args:
|
|
url (str): The URL of the file to download.
|
|
file_name (str): The local path where the file will be saved.
|
|
"""
|
|
print(f"Downloading: {file_name}...")
|
|
response = requests.get(url, stream=True) # Stream the file content
|
|
total_size = int(response.headers.get('content-length', 0)) # Get total file size if available
|
|
block_size = 1024 # Size of each block for the progress bar
|
|
with open(file_name, 'wb') as f: # Open file for writing in binary mode
|
|
for chunk in tqdm(response.iter_content(block_size), total=total_size // block_size, desc="Downloading", leave=True):
|
|
f.write(chunk) # Write each chunk to the file
|
|
|
|
def download_dataset(val_url: str, train_urls: List[str], val_dir: str, train_dir: str, max_train_files: int) -> None:
|
|
"""
|
|
Manages downloading of the dataset, including both validation and training files.
|
|
|
|
Args:
|
|
val_url (str): URL for the validation dataset.
|
|
train_urls (list): List of URLs for the training dataset files.
|
|
val_dir (str): Directory where the validation file will be stored.
|
|
train_dir (str): Directory where the training files will be stored.
|
|
max_train_files (int): Maximum number of training files to download.
|
|
"""
|
|
# Define the path for the validation file
|
|
val_file_path = os.path.join(val_dir, "val.jsonl.zst")
|
|
if not os.path.exists(val_file_path): # Check if the validation file already exists
|
|
print(f"Validation file not found. Downloading from {val_url}...")
|
|
download_file(val_url, val_file_path) # Download the validation file
|
|
else:
|
|
print("Validation data already present. Skipping download.")
|
|
|
|
# Loop through the training file URLs and download if not already present
|
|
for idx, url in enumerate(train_urls[:max_train_files]): # Limit to max_train_files
|
|
file_name = f"{idx:02d}.jsonl.zst" # Format file name (e.g., 00.jsonl.zst)
|
|
file_path = os.path.join(train_dir, file_name) # Construct the full file path
|
|
if not os.path.exists(file_path): # Check if the file already exists
|
|
print(f"Training file {file_name} not found. Downloading...")
|
|
download_file(url, file_path) # Download the training file
|
|
else:
|
|
print(f"Training file {file_name} already present. Skipping download.")
|
|
|
|
def main() -> None:
|
|
"""
|
|
Main function to parse arguments and orchestrate the dataset download process.
|
|
"""
|
|
# Parse command-line arguments using argparse
|
|
parser = argparse.ArgumentParser(description="Download PILE dataset.") # Description of the script
|
|
parser.add_argument('--train_max', type=int, default=1, help="Max number of training files to download.") # Max training files
|
|
parser.add_argument('--train_dir', default="data/train", help="Directory for storing training data.") # Training directory
|
|
parser.add_argument('--val_dir', default="data/val", help="Directory for storing validation data.") # Validation directory
|
|
|
|
args = parser.parse_args() # Parse the arguments provided by the user
|
|
|
|
# Ensure directories for training and validation data exist
|
|
os.makedirs(args.train_dir, exist_ok=True) # Create training directory if it doesn't exist
|
|
os.makedirs(args.val_dir, exist_ok=True) # Create validation directory if it doesn't exist
|
|
|
|
# Start downloading the dataset
|
|
download_dataset(VAL_URL, TRAIN_URLS, args.val_dir, args.train_dir, args.train_max)
|
|
|
|
print("Dataset downloaded successfully.") # Indicate successful download
|
|
|
|
if __name__ == "__main__":
|
|
# Entry point of the script
|
|
main() |