File size: 8,157 Bytes
ad68b7f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 | """
Consolidate the actual training data (litdata binary chunks) into single folders.
Pretraining: litdata_3b + litdata_english -> consolidated_pretrain_litdata/
Finetuning: finetune/ + finetune_english/ -> consolidated_finetune/
Token counts are computed directly from the litdata index.json metadata.
"""
import json
import shutil
import os
from pathlib import Path
ROOT = Path(__file__).resolve().parent.parent.parent # LUNA root
DATA = ROOT / "Base" / "data"
DATASETS = ROOT / "Base" / "Datasets"
def read_litdata_index(litdata_dir):
"""Read index.json and return chunks list + config."""
with open(litdata_dir / "index.json", "r") as f:
index = json.load(f)
return index["chunks"], index["config"]
def count_tokens_from_index(chunks):
"""Sum up all tokens from chunk dim fields."""
return sum(c["dim"] for c in chunks)
# ══════════════════════════════════════════════════════════════════
# 1. PRETRAINING: litdata_3b + litdata_english
# ══════════════════════════════════════════════════════════════════
print("=" * 65)
print(" PRETRAINING DATA CONSOLIDATION")
print("=" * 65)
# Read both indexes
chunks_3b, config = read_litdata_index(DATA / "litdata_3b")
chunks_en, _ = read_litdata_index(DATA / "litdata_english")
tokens_3b = count_tokens_from_index(chunks_3b)
tokens_en = count_tokens_from_index(chunks_en)
tokens_pretrain = tokens_3b + tokens_en
print(f"\n litdata_3b: {len(chunks_3b):>4} chunks | {tokens_3b:>15,} tokens")
print(f" litdata_english: {len(chunks_en):>4} chunks | {tokens_en:>15,} tokens")
print(f" {'─' * 55}")
print(f" TOTAL PRETRAIN: {len(chunks_3b)+len(chunks_en):>4} chunks | {tokens_pretrain:>15,} tokens")
print(f" ({tokens_pretrain / 1e9:.3f} B tokens)")
# Create consolidated litdata folder
out_pretrain = DATA / "consolidated_pretrain_litdata"
out_pretrain.mkdir(parents=True, exist_ok=True)
# Copy chunks from litdata_3b (keep original names as chunk-0-{0..N})
print(f"\n Copying litdata_3b chunks ({len(chunks_3b)} files)...")
new_chunks = []
for i, chunk in enumerate(chunks_3b):
src = DATA / "litdata_3b" / chunk["filename"]
new_name = f"chunk-0-{i}.bin"
dst = out_pretrain / new_name
if not dst.exists():
shutil.copy2(str(src), str(dst))
new_chunks.append({**chunk, "filename": new_name})
if (i + 1) % 50 == 0 or i == len(chunks_3b) - 1:
print(f" {i+1}/{len(chunks_3b)} copied")
# Copy chunks from litdata_english (renumber continuing from 3b)
offset = len(chunks_3b)
print(f"\n Copying litdata_english chunks ({len(chunks_en)} files)...")
for i, chunk in enumerate(chunks_en):
src = DATA / "litdata_english" / chunk["filename"]
new_name = f"chunk-0-{offset + i}.bin"
dst = out_pretrain / new_name
if not dst.exists():
shutil.copy2(str(src), str(dst))
new_chunks.append({**chunk, "filename": new_name})
print(f" {i+1}/{len(chunks_en)} copied")
# Write combined index.json
combined_index = {"chunks": new_chunks, "config": config}
with open(out_pretrain / "index.json", "w") as f:
json.dump(combined_index, f, indent=2)
print(f"\n Saved: {out_pretrain}")
# ══════════════════════════════════════════════════════════════════
# 2. FINETUNING: finetune/ + finetune_english/
# ══════════════════════════════════════════════════════════════════
print(f"\n{'=' * 65}")
print(" FINETUNING DATA CONSOLIDATION")
print("=" * 65)
# Load all finetune JSONs
ft_v1_train = json.loads((DATASETS / "finetune" / "train.json").read_text(encoding="utf-8"))
ft_v1_val = json.loads((DATASETS / "finetune" / "val.json").read_text(encoding="utf-8"))
ft_en_train = json.loads((DATASETS / "finetune_english" / "train.json").read_text(encoding="utf-8"))
ft_en_val = json.loads((DATASETS / "finetune_english" / "val.json").read_text(encoding="utf-8"))
ft_v1_all = ft_v1_train + ft_v1_val
ft_en_all = ft_en_train + ft_en_val
# Tag sources
for item in ft_v1_all:
item["source"] = "finetune_v1"
for item in ft_en_all:
item["source"] = "finetune_english"
# Count tokens using the same tokenizer
from tokenizers import Tokenizer
tokenizer = Tokenizer.from_file(
str(ROOT / "Base" / "checkpoints" / "EleutherAI" / "pythia-160m" / "tokenizer.json")
)
def finetune_text(item):
parts = []
if item.get("instruction"): parts.append(item["instruction"])
if item.get("input"): parts.append(item["input"])
if item.get("output"): parts.append(item["output"])
return " ".join(parts)
def count_tokens_batch(texts, label=""):
total = 0
BATCH = 10000
for i in range(0, len(texts), BATCH):
batch = texts[i:i+BATCH]
encoded = tokenizer.encode_batch(batch, add_special_tokens=False)
total += sum(len(e.ids) for e in encoded)
return total
print(f"\n finetune_v1: train={len(ft_v1_train):>7,} val={len(ft_v1_val):>6,} total={len(ft_v1_all):>7,}")
print(f" finetune_english: train={len(ft_en_train):>7,} val={len(ft_en_val):>6,} total={len(ft_en_all):>7,}")
print("\n Counting finetune tokens...")
ft_v1_tokens = count_tokens_batch([finetune_text(x) for x in ft_v1_all], "v1")
ft_en_tokens = count_tokens_batch([finetune_text(x) for x in ft_en_all], "english")
ft_total = ft_v1_tokens + ft_en_tokens
print(f"\n finetune_v1 tokens: {ft_v1_tokens:>12,}")
print(f" finetune_english tokens: {ft_en_tokens:>12,}")
print(f" {'─' * 55}")
print(f" TOTAL FINETUNE: {ft_total:>12,} tokens")
# Save consolidated finetune
out_finetune = DATASETS / "consolidated_finetune"
out_finetune.mkdir(parents=True, exist_ok=True)
all_finetune = ft_v1_all + ft_en_all
with open(out_finetune / "all_finetune_data.json", "w", encoding="utf-8") as f:
json.dump(all_finetune, f, ensure_ascii=False, indent=2)
print(f"\n Saved: {out_finetune / 'all_finetune_data.json'}")
print(f" Total samples: {len(all_finetune):,}")
# ══════════════════════════════════════════════════════════════════
# FINAL SUMMARY
# ══════════════════════════════════════════════════════════════════
print(f"\n{'=' * 65}")
print(" FINAL SUMMARY")
print("=" * 65)
summary = f"""
PRETRAINING (Base/data/consolidated_pretrain_litdata/)
litdata_3b : {len(chunks_3b):>4} chunks | {tokens_3b:>15,} tokens
litdata_english : {len(chunks_en):>4} chunks | {tokens_en:>15,} tokens
───────────────────────────────────────────────────────
TOTAL : {len(chunks_3b)+len(chunks_en):>4} chunks | {tokens_pretrain:>15,} tokens ({tokens_pretrain/1e9:.3f}B)
FINETUNING (Base/Datasets/consolidated_finetune/)
finetune_v1 : {len(ft_v1_all):>7,} samples | {ft_v1_tokens:>12,} tokens
finetune_english : {len(ft_en_all):>7,} samples | {ft_en_tokens:>12,} tokens
───────────────────────────────────────────────────────
TOTAL : {len(all_finetune):>7,} samples | {ft_total:>12,} tokens
GRAND TOTAL TOKENS: {tokens_pretrain + ft_total:,}
"""
print(summary)
# Save summary
with open(out_pretrain / "SUMMARY.txt", "w", encoding="utf-8") as f:
f.write(summary)
with open(out_finetune / "SUMMARY.txt", "w", encoding="utf-8") as f:
f.write(summary)
print("Done! Summary saved to both consolidated folders.")
|