| import os |
| import sys |
| import json |
| from datasets import load_dataset |
| from tqdm import tqdm |
|
|
| |
| project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) |
| sys.path.append(project_root) |
|
|
| from configs.model_config import PathConfig |
|
|
| def prepare_directories(): |
| os.makedirs(PathConfig.RAW_DATA_DIR, exist_ok=True) |
| os.makedirs(PathConfig.ALIGNMENT_DIR, exist_ok=True) |
| os.makedirs(PathConfig.SELF_ADAPT_DIR, exist_ok=True) |
| os.makedirs(PathConfig.DEBUG_LOCAL_DIR, exist_ok=True) |
|
|
| def download_pretrain_data(): |
| """Downloads a rich blend of datasets for robust Pretraining.""" |
| print("\n[1/4] Downloading Pretraining Data Blend (Stories, Wiki, Web)...") |
| |
| datasets_to_fetch = [ |
| {"name": "TinyStories", "path": "roneneldan/TinyStories", "split": "train", "samples": 250_000, "out": "tinystories.txt"}, |
| {"name": "WikiText-2", "path": "wikitext", "config": "wikitext-2-raw-v1", "split": "train", "samples": 100_000, "out": "wikitext.txt"}, |
| {"name": "OpenWebText", "path": "Skylion007/openwebtext", "split": "train", "samples": 50_000, "out": "openwebtext.txt"} |
| ] |
| |
| for ds_info in datasets_to_fetch: |
| print(f" -> Fetching {ds_info['name']}...") |
| if "config" in ds_info: |
| dataset = load_dataset(ds_info["path"], ds_info["config"], split=ds_info["split"], streaming=True) |
| else: |
| dataset = load_dataset(ds_info["path"], split=ds_info["split"], streaming=True) |
| |
| out_path = os.path.join(PathConfig.RAW_DATA_DIR, ds_info["out"]) |
| records_written = 0 |
| |
| with open(out_path, "w", encoding="utf-8") as f: |
| for item in tqdm(dataset, desc=f"Writing {ds_info['name']}", total=ds_info["samples"]): |
| if records_written >= ds_info["samples"]: |
| break |
| text = item["text"].strip() |
| if text: |
| f.write(text + "\n<|eos|>\n") |
| records_written += 1 |
| |
| print(f" Saved {records_written} documents to {out_path}") |
|
|
| def download_sft_data(num_samples=2000): |
| """Downloads Stanford Alpaca and formats it for SFT.""" |
| print(f"\n[2/4] Downloading Alpaca (SFT) | Target: {num_samples} samples...") |
| dataset = load_dataset("tatsu-lab/alpaca", split="train") |
| |
| records_written = 0 |
| with open(PathConfig.SFT_DATA_PATH, "w", encoding="utf-8") as f: |
| for item in tqdm(dataset, desc="Formatting SFT"): |
| if records_written >= num_samples: |
| break |
| |
| prompt = item["instruction"] |
| if item["input"].strip(): |
| prompt += f"\nInput: {item['input']}" |
| |
| record = { |
| "prompt": prompt, |
| "response": item["output"] |
| } |
| f.write(json.dumps(record) + "\n") |
| records_written += 1 |
| |
| print(f"-> Saved real SFT data to: {PathConfig.SFT_DATA_PATH}") |
|
|
| def download_dpo_data(num_samples=1000): |
| """Downloads Orca DPO pairs and formats them for DPO alignment.""" |
| print(f"\n[3/4] Downloading Orca DPO Pairs | Target: {num_samples} samples...") |
| dataset = load_dataset("Intel/orca_dpo_pairs", split="train") |
| |
| records_written = 0 |
| with open(PathConfig.PREFERENCES_DATA_PATH, "w", encoding="utf-8") as f: |
| for item in tqdm(dataset, desc="Formatting DPO"): |
| if records_written >= num_samples: |
| break |
| |
| record = { |
| "prompt": item["question"], |
| "chosen": item["chosen"], |
| "rejected": item["rejected"] |
| } |
| f.write(json.dumps(record) + "\n") |
| records_written += 1 |
| |
| print(f"-> Saved real DPO data to: {PathConfig.PREFERENCES_DATA_PATH}") |
|
|
| def download_squad_adaptation_data(num_samples=500): |
| """Downloads SQuAD v1.1 for SEAL knowledge incorporation.""" |
| print(f"\n[4/4] Downloading SQuAD v1.1 (SEAL) | Target: {num_samples} passages...") |
| dataset = load_dataset("squad", split="train") |
| |
| passage_groups = {} |
| for item in tqdm(dataset, desc="Grouping SQuAD Passages"): |
| passage = item["context"] |
| if passage not in passage_groups: |
| passage_groups[passage] = [] |
| |
| passage_groups[passage].append({ |
| "q": item["question"], |
| "a": item["answers"]["text"][0] |
| }) |
| |
| records_written = 0 |
| with open(PathConfig.SQUAD_MINI_PATH, "w", encoding="utf-8") as f: |
| for passage, qa_pairs in passage_groups.items(): |
| if records_written >= num_samples: |
| break |
| |
| record = { |
| "passage": passage, |
| "qa_pairs": qa_pairs[:3] |
| } |
| f.write(json.dumps(record) + "\n") |
| records_written += 1 |
| |
| print(f"-> Saved real SQuAD data to: {PathConfig.SQUAD_MINI_PATH}") |
|
|
| if __name__ == "__main__": |
| print("=== STARTING REAL-WORLD DATA ACQUISITION ===") |
| prepare_directories() |
| download_pretrain_data() |
| download_sft_data() |
| download_dpo_data() |
| download_squad_adaptation_data() |
| print("\n=== DATA ACQUISITION COMPLETE: Ready for training on real data! ===") |
|
|
