| """ |
| Dataset Downloader Utility for fetching open-source text corpora for training Small Language Models. |
| """ |
|
|
| import argparse |
| import os |
| import urllib.request |
| import sys |
|
|
| sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) |
|
|
| from slm.utils.logger import get_logger |
|
|
| logger = get_logger("slm.scripts.download_dataset") |
|
|
| DATASET_URLS = { |
| "tinystories": { |
| "url": "https://raw.githubusercontent.com/roneneldan/TinyStories/main/TinyStories-valid.txt", |
| "filename": "tinystories_valid.txt", |
| "description": "Clean synthetic story corpus ideal for small language model training." |
| }, |
| "wikitext2": { |
| "url": "https://raw.githubusercontent.com/pytorch/examples/main/word_language_model/data/wikitext-2/train.txt", |
| "filename": "wikitext2_train.txt", |
| "description": "High-quality Wikipedia text corpus (WikiText-2)." |
| }, |
| "tinyshakespeare": { |
| "url": "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt", |
| "filename": "tinyshakespeare.txt", |
| "description": "Classic Shakespeare text corpus (1MB)." |
| } |
| } |
|
|
|
|
| def download_dataset(dataset_key: str, output_dir: str = "data") -> str: |
| """ |
| Downloads specified open text dataset into target directory. |
| """ |
| if dataset_key not in DATASET_URLS: |
| raise ValueError(f"Unknown dataset key '{dataset_key}'. Choose from: {list(DATASET_URLS.keys())}") |
|
|
| info = DATASET_URLS[dataset_key] |
| os.makedirs(output_dir, exist_ok=True) |
| dest_path = os.path.join(output_dir, info["filename"]) |
|
|
| logger.info(f"Downloading {dataset_key} ({info['description']})...") |
| logger.info(f"Source URL: {info['url']}") |
|
|
| def _progress(count, block_size, total_size): |
| percent = int(count * block_size * 100 / max(1, total_size)) |
| sys.stdout.write(f"\rDownloading... {percent}% complete") |
| sys.stdout.flush() |
|
|
| urllib.request.urlretrieve(info["url"], dest_path, reporthook=_progress) |
| print() |
| logger.info(f"Successfully downloaded dataset to {dest_path}") |
| return dest_path |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description="Download open text datasets for Small Language Model training.") |
| parser.add_argument("--name", type=str, default="wikitext2", choices=list(DATASET_URLS.keys()), help="Dataset choice") |
| parser.add_argument("--output_dir", type=str, default="data", help="Output folder") |
| args = parser.parse_args() |
|
|
| download_dataset(args.name, args.output_dir) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
