ASTERIZER commited on
Commit Β·
ad68b7f
1
Parent(s): 5a142b7
LUNA 100M: cloud-ready training pipeline
Browse files- Config-driven train.py with auto/manual hardware detection
- train_config.yaml with auto_config toggle
- setup_and_train.sh: one-command RunPod deployment
- fetch_data.py: HuggingFace/GDrive dataset downloader
- benchmark_runpod.py: local benchmark + cloud cost estimator
- Tokenizer files, data preprocessing scripts, GGUF converter
- .gitattributes +0 -2
- .gitignore +60 -0
- Base/checkpoints/EleutherAI/pythia-160m/config.json +24 -0
- Base/checkpoints/EleutherAI/pythia-160m/tokenizer.json +0 -0
- Base/checkpoints/EleutherAI/pythia-160m/tokenizer_config.json +9 -0
- Base/configs/finetune_100m_english_instruct.yaml +55 -0
- Base/configs/finetune_100m_instruct.yaml +56 -0
- Base/configs/pretrain_100m_10m.yaml +49 -0
- Base/configs/pretrain_100m_3b.yaml +50 -0
- Base/configs/pretrain_100m_3b_full.yaml +53 -0
- Base/configs/pretrain_100m_english.yaml +58 -0
- Base/scripts/audit_5b_pretrain.py +494 -0
- Base/scripts/audit_english_clean.py +707 -0
- Base/scripts/audit_litdata.py +475 -0
- Base/scripts/build_100m_english.py +667 -0
- Base/scripts/build_5b_pretrain.py +761 -0
- Base/scripts/build_english_corpus.py +349 -0
- Base/scripts/build_instruct_dataset.py +769 -0
- Base/scripts/chat.py +68 -0
- Base/scripts/clean_and_merge_pretrain.py +497 -0
- Base/scripts/consolidate_datasets.py +194 -0
- Base/scripts/consolidate_litdata.py +187 -0
- Base/scripts/dedup_5b_pretrain.py +298 -0
- Base/scripts/filter_datasets.py +151 -0
- Base/scripts/merge_and_build.py +701 -0
- Base/scripts/prepare_finetune_data.py +88 -0
- Base/scripts/prepare_litdata.py +275 -0
- Base/scripts/reclean_3b.py +593 -0
- Base/scripts/reclean_english.py +610 -0
- Base/scripts/reclean_litdata.py +444 -0
- Base/scripts/smart_cleanup_english.py +379 -0
- Base/scripts/validate_model.py +137 -0
- README.md +119 -2
- benchmark_runpod.py +412 -0
- fetch_data.py +162 -0
- quantisations/convert_to_gguf.py +195 -0
- requirements.txt +6 -0
- setup_and_train.sh +94 -0
- train.py +608 -0
- train_config.yaml +67 -0
.gitattributes
DELETED
|
@@ -1,2 +0,0 @@
|
|
| 1 |
-
# Auto detect text files and perform LF normalization
|
| 2 |
-
* text=auto
|
|
|
|
|
|
|
|
|
.gitignore
ADDED
|
@@ -0,0 +1,60 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# βββ Python βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 2 |
+
__pycache__/
|
| 3 |
+
*.py[cod]
|
| 4 |
+
*.pyo
|
| 5 |
+
*.pyd
|
| 6 |
+
*.egg-info/
|
| 7 |
+
*.egg
|
| 8 |
+
dist/
|
| 9 |
+
build/
|
| 10 |
+
.eggs/
|
| 11 |
+
*.spec
|
| 12 |
+
|
| 13 |
+
# βββ Virtual environments βββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 14 |
+
.venv/
|
| 15 |
+
venv/
|
| 16 |
+
env/
|
| 17 |
+
.env
|
| 18 |
+
|
| 19 |
+
# βββ Large binary data β NEVER commit these βββββββββββββββββββββββββββββββββββ
|
| 20 |
+
# LitData tokenized chunks (stream from cloud storage instead)
|
| 21 |
+
Base/data/
|
| 22 |
+
Data preprocess/
|
| 23 |
+
|
| 24 |
+
# Raw / processed datasets (host on HuggingFace or GDrive)
|
| 25 |
+
Base/Datasets/
|
| 26 |
+
|
| 27 |
+
# Model checkpoints and weights (host on HuggingFace model hub)
|
| 28 |
+
Base/out/
|
| 29 |
+
|
| 30 |
+
# GGUF quantisations
|
| 31 |
+
quantisations/*.gguf
|
| 32 |
+
quantisations/*.bin
|
| 33 |
+
|
| 34 |
+
# Zip archives
|
| 35 |
+
*.zip
|
| 36 |
+
|
| 37 |
+
# βββ Tokenizer cache (download at runtime) ββββββββββββββββββββββββββββββββββββ
|
| 38 |
+
Base/checkpoints/EleutherAI/pythia-160m/.cache/
|
| 39 |
+
|
| 40 |
+
# βββ Logs / temp ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 41 |
+
*.log
|
| 42 |
+
temp_*.txt
|
| 43 |
+
logs/
|
| 44 |
+
tensorboard/
|
| 45 |
+
lightning_logs/
|
| 46 |
+
out/
|
| 47 |
+
runpod_cost_estimate.json
|
| 48 |
+
|
| 49 |
+
# βββ OS / editor ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 50 |
+
.DS_Store
|
| 51 |
+
Thumbs.db
|
| 52 |
+
*.swp
|
| 53 |
+
*.swo
|
| 54 |
+
.idea/
|
| 55 |
+
.vscode/
|
| 56 |
+
*.code-workspace
|
| 57 |
+
|
| 58 |
+
# βββ Temporary / generated ββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 59 |
+
temp_prompt_*.txt
|
| 60 |
+
runpod_cost_estimate.json
|
Base/checkpoints/EleutherAI/pythia-160m/config.json
ADDED
|
@@ -0,0 +1,24 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"architectures": [
|
| 3 |
+
"GPTNeoXForCausalLM"
|
| 4 |
+
],
|
| 5 |
+
"bos_token_id": 0,
|
| 6 |
+
"eos_token_id": 0,
|
| 7 |
+
"hidden_act": "gelu",
|
| 8 |
+
"hidden_size": 768,
|
| 9 |
+
"initializer_range": 0.02,
|
| 10 |
+
"intermediate_size": 3072,
|
| 11 |
+
"layer_norm_eps": 1e-05,
|
| 12 |
+
"max_position_embeddings": 2048,
|
| 13 |
+
"model_type": "gpt_neox",
|
| 14 |
+
"num_attention_heads": 12,
|
| 15 |
+
"num_hidden_layers": 12,
|
| 16 |
+
"rotary_emb_base": 10000,
|
| 17 |
+
"rotary_pct": 0.25,
|
| 18 |
+
"tie_word_embeddings": false,
|
| 19 |
+
"torch_dtype": "float16",
|
| 20 |
+
"transformers_version": "4.24.0",
|
| 21 |
+
"use_cache": true,
|
| 22 |
+
"use_parallel_residual": true,
|
| 23 |
+
"vocab_size": 50304
|
| 24 |
+
}
|
Base/checkpoints/EleutherAI/pythia-160m/tokenizer.json
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
Base/checkpoints/EleutherAI/pythia-160m/tokenizer_config.json
ADDED
|
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"add_prefix_space": false,
|
| 3 |
+
"bos_token": "<|endoftext|>",
|
| 4 |
+
"eos_token": "<|endoftext|>",
|
| 5 |
+
"name_or_path": "EleutherAI/gpt-neox-20b",
|
| 6 |
+
"special_tokens_map_file": "/admin/home-hailey/.cache/huggingface/hub/models--EleutherAI--gpt-neox-20b/snapshots/4e49eadb5d14bd22f314ec3f45b69a87b88c7691/special_tokens_map.json",
|
| 7 |
+
"tokenizer_class": "GPTNeoXTokenizer",
|
| 8 |
+
"unk_token": "<|endoftext|>"
|
| 9 |
+
}
|
Base/configs/finetune_100m_english_instruct.yaml
ADDED
|
@@ -0,0 +1,55 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Instruction Fine-Tuning Config for English-Pretrained 100M Model
|
| 2 |
+
# Full fine-tuning (all weights) on massive diverse instruct dataset (~100M tokens)
|
| 3 |
+
# Longer sequences than v1 (512 vs 256) for richer, more detailed outputs
|
| 4 |
+
# Checkpoint from English continued pretraining
|
| 5 |
+
|
| 6 |
+
# Base checkpoint to fine-tune from (output of pretrain_100m_english.yaml)
|
| 7 |
+
checkpoint_dir: Base/out/pretrain/custom-100m-english/final_raw
|
| 8 |
+
|
| 9 |
+
# Output directory
|
| 10 |
+
out_dir: Base/out/finetune/custom-100m-english-instruct
|
| 11 |
+
|
| 12 |
+
# Precision
|
| 13 |
+
precision: bf16-mixed
|
| 14 |
+
|
| 15 |
+
# Data β new massive diverse dataset
|
| 16 |
+
data:
|
| 17 |
+
class_path: litgpt.data.JSON
|
| 18 |
+
init_args:
|
| 19 |
+
json_path: Base/Datasets/finetune_english
|
| 20 |
+
mask_prompt: true
|
| 21 |
+
val_split_fraction: null
|
| 22 |
+
prompt_style: alpaca
|
| 23 |
+
num_workers: 4
|
| 24 |
+
|
| 25 |
+
# Training
|
| 26 |
+
train:
|
| 27 |
+
save_interval: 1000
|
| 28 |
+
log_interval: 10
|
| 29 |
+
global_batch_size: 64
|
| 30 |
+
micro_batch_size: 8
|
| 31 |
+
lr_warmup_steps: 200
|
| 32 |
+
epochs: 2
|
| 33 |
+
max_seq_length: 512
|
| 34 |
+
min_lr: 5.0e-7
|
| 35 |
+
|
| 36 |
+
# Evaluation
|
| 37 |
+
eval:
|
| 38 |
+
interval: 500
|
| 39 |
+
max_new_tokens: 200
|
| 40 |
+
max_iters: 100
|
| 41 |
+
initial_validation: true
|
| 42 |
+
final_validation: true
|
| 43 |
+
|
| 44 |
+
# Optimizer β slightly lower LR for larger, more diverse data
|
| 45 |
+
optimizer:
|
| 46 |
+
class_path: torch.optim.AdamW
|
| 47 |
+
init_args:
|
| 48 |
+
lr: 1.0e-5
|
| 49 |
+
weight_decay: 0.01
|
| 50 |
+
betas:
|
| 51 |
+
- 0.9
|
| 52 |
+
- 0.95
|
| 53 |
+
|
| 54 |
+
# Logging
|
| 55 |
+
logger_name: csv
|
Base/configs/finetune_100m_instruct.yaml
ADDED
|
@@ -0,0 +1,56 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Instruction Fine-Tuning Config for Custom 100M Model
|
| 2 |
+
# Full fine-tuning (all weights, no LoRA) on English instruction datasets
|
| 3 |
+
# ~41K samples, 3 epochs, short sequences (max ~160 tokens)
|
| 4 |
+
|
| 5 |
+
# Base checkpoint to fine-tune from
|
| 6 |
+
checkpoint_dir: Base/out/pretrain/custom-100m-3b-full/final
|
| 7 |
+
|
| 8 |
+
# Output directory
|
| 9 |
+
out_dir: Base/out/finetune/custom-100m-instruct
|
| 10 |
+
|
| 11 |
+
# Precision
|
| 12 |
+
precision: bf16-mixed
|
| 13 |
+
|
| 14 |
+
# Data
|
| 15 |
+
data:
|
| 16 |
+
class_path: litgpt.data.JSON
|
| 17 |
+
init_args:
|
| 18 |
+
json_path: Base/Datasets/finetune
|
| 19 |
+
mask_prompt: true
|
| 20 |
+
val_split_fraction: null
|
| 21 |
+
prompt_style: alpaca
|
| 22 |
+
num_workers: 4
|
| 23 |
+
|
| 24 |
+
# Training
|
| 25 |
+
train:
|
| 26 |
+
save_interval: 500
|
| 27 |
+
log_interval: 10
|
| 28 |
+
global_batch_size: 64
|
| 29 |
+
micro_batch_size: 16
|
| 30 |
+
lr_warmup_steps: 50
|
| 31 |
+
epochs: 3
|
| 32 |
+
max_seq_length: 256
|
| 33 |
+
min_lr: 1.0e-6
|
| 34 |
+
|
| 35 |
+
# Evaluation
|
| 36 |
+
eval:
|
| 37 |
+
interval: 200
|
| 38 |
+
max_new_tokens: 100
|
| 39 |
+
max_iters: 100
|
| 40 |
+
initial_validation: true
|
| 41 |
+
final_validation: true
|
| 42 |
+
|
| 43 |
+
# Optimizer
|
| 44 |
+
optimizer:
|
| 45 |
+
class_path: torch.optim.AdamW
|
| 46 |
+
init_args:
|
| 47 |
+
lr: 2.0e-5
|
| 48 |
+
weight_decay: 0.01
|
| 49 |
+
betas:
|
| 50 |
+
- 0.9
|
| 51 |
+
- 0.95
|
| 52 |
+
|
| 53 |
+
# Logging
|
| 54 |
+
logger_name: csv
|
| 55 |
+
|
| 56 |
+
seed: 42
|
Base/configs/pretrain_100m_10m.yaml
ADDED
|
@@ -0,0 +1,49 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# pretrain_100m_10m.yaml
|
| 2 |
+
# Small test run config: ~100M param custom model on ~10M tokens
|
| 3 |
+
# Use this to verify training pipeline works before committing to full 3B run
|
| 4 |
+
|
| 5 |
+
model_name: pythia-14m
|
| 6 |
+
model_config:
|
| 7 |
+
name: custom-pythia-like-100m
|
| 8 |
+
block_size: 1024
|
| 9 |
+
n_layer: 10
|
| 10 |
+
n_embd: 768
|
| 11 |
+
n_head: 12
|
| 12 |
+
vocab_size: 50254
|
| 13 |
+
padding_multiple: 128
|
| 14 |
+
rotary_percentage: 0.25
|
| 15 |
+
|
| 16 |
+
out_dir: Base/out/pretrain/custom-100m-10m-test
|
| 17 |
+
|
| 18 |
+
precision: bf16-mixed
|
| 19 |
+
|
| 20 |
+
data:
|
| 21 |
+
class_path: litgpt.data.LitData
|
| 22 |
+
init_args:
|
| 23 |
+
data_path: Base/data/litdata_10m
|
| 24 |
+
|
| 25 |
+
train:
|
| 26 |
+
save_interval: 500
|
| 27 |
+
log_interval: 5
|
| 28 |
+
global_batch_size: 16
|
| 29 |
+
micro_batch_size: 4
|
| 30 |
+
lr_warmup_steps: 50
|
| 31 |
+
max_tokens: 10000000
|
| 32 |
+
max_seq_length: 1024
|
| 33 |
+
tie_embeddings: true
|
| 34 |
+
max_norm: 1.0
|
| 35 |
+
min_lr: 6.0e-05
|
| 36 |
+
|
| 37 |
+
optimizer:
|
| 38 |
+
class_path: torch.optim.AdamW
|
| 39 |
+
init_args:
|
| 40 |
+
lr: 6.0e-04
|
| 41 |
+
weight_decay: 0.1
|
| 42 |
+
betas: [0.9, 0.95]
|
| 43 |
+
eps: 1.0e-08
|
| 44 |
+
|
| 45 |
+
devices: 1
|
| 46 |
+
num_nodes: 1
|
| 47 |
+
tokenizer_dir: Base/checkpoints/EleutherAI/pythia-160m
|
| 48 |
+
logger_name: tensorboard
|
| 49 |
+
seed: 1337
|
Base/configs/pretrain_100m_3b.yaml
ADDED
|
@@ -0,0 +1,50 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# pretrain_100m_3b.yaml
|
| 2 |
+
# Full pretraining config: ~100M param custom model on ~27M filtered tokens
|
| 3 |
+
# With the strict quality filter (language_score > 0.99), only ~27M tokens qualify
|
| 4 |
+
# Training for ~270M tokens = ~10 epochs over the data
|
| 5 |
+
|
| 6 |
+
model_name: pythia-14m
|
| 7 |
+
model_config:
|
| 8 |
+
name: custom-pythia-like-100m
|
| 9 |
+
block_size: 1024
|
| 10 |
+
n_layer: 10
|
| 11 |
+
n_embd: 768
|
| 12 |
+
n_head: 12
|
| 13 |
+
vocab_size: 50254
|
| 14 |
+
padding_multiple: 128
|
| 15 |
+
rotary_percentage: 0.25
|
| 16 |
+
|
| 17 |
+
out_dir: Base/out/pretrain/custom-100m-3b
|
| 18 |
+
|
| 19 |
+
precision: bf16-mixed
|
| 20 |
+
|
| 21 |
+
data:
|
| 22 |
+
class_path: litgpt.data.LitData
|
| 23 |
+
init_args:
|
| 24 |
+
data_path: Base/data/litdata_3b
|
| 25 |
+
|
| 26 |
+
train:
|
| 27 |
+
save_interval: 2000
|
| 28 |
+
log_interval: 10
|
| 29 |
+
global_batch_size: 64
|
| 30 |
+
micro_batch_size: 4
|
| 31 |
+
lr_warmup_steps: 50
|
| 32 |
+
max_tokens: 27000000
|
| 33 |
+
max_seq_length: 1024
|
| 34 |
+
tie_embeddings: true
|
| 35 |
+
max_norm: 1.0
|
| 36 |
+
min_lr: 6.0e-05
|
| 37 |
+
|
| 38 |
+
optimizer:
|
| 39 |
+
class_path: torch.optim.AdamW
|
| 40 |
+
init_args:
|
| 41 |
+
lr: 6.0e-04
|
| 42 |
+
weight_decay: 0.1
|
| 43 |
+
betas: [0.9, 0.95]
|
| 44 |
+
eps: 1.0e-08
|
| 45 |
+
|
| 46 |
+
devices: 1
|
| 47 |
+
num_nodes: 1
|
| 48 |
+
tokenizer_dir: Base/checkpoints/EleutherAI/pythia-160m
|
| 49 |
+
logger_name: tensorboard
|
| 50 |
+
seed: 1337
|
Base/configs/pretrain_100m_3b_full.yaml
ADDED
|
@@ -0,0 +1,53 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# pretrain_100m_3b_full.yaml
|
| 2 |
+
# Full bulk pretraining: ~100M param model on 3B tokens (score >= 0.96)
|
| 3 |
+
# Optimized for RTX 4060 Ti 16GB β targeting ~90% VRAM utilization
|
| 4 |
+
# Checkpoints every 1000 steps for resume safety
|
| 5 |
+
|
| 6 |
+
model_name: pythia-14m
|
| 7 |
+
model_config:
|
| 8 |
+
name: custom-pythia-like-100m
|
| 9 |
+
block_size: 1024
|
| 10 |
+
n_layer: 10
|
| 11 |
+
n_embd: 768
|
| 12 |
+
n_head: 12
|
| 13 |
+
vocab_size: 50254
|
| 14 |
+
padding_multiple: 128
|
| 15 |
+
rotary_percentage: 0.25
|
| 16 |
+
|
| 17 |
+
# Separate output β does NOT override old checkpoints
|
| 18 |
+
out_dir: Base/out/pretrain/custom-100m-3b-full
|
| 19 |
+
|
| 20 |
+
precision: bf16-mixed
|
| 21 |
+
|
| 22 |
+
data:
|
| 23 |
+
class_path: litgpt.data.LitData
|
| 24 |
+
init_args:
|
| 25 |
+
data_path: Base/data/litdata_3b
|
| 26 |
+
|
| 27 |
+
train:
|
| 28 |
+
save_interval: 1000
|
| 29 |
+
log_interval: 10
|
| 30 |
+
# micro_batch_size=12 pushes VRAM to ~14GB (90% of 16GB)
|
| 31 |
+
# global_batch=120 = 10 gradient accumulation steps
|
| 32 |
+
global_batch_size: 120
|
| 33 |
+
micro_batch_size: 12
|
| 34 |
+
lr_warmup_steps: 500
|
| 35 |
+
max_tokens: 3000000000
|
| 36 |
+
max_seq_length: 1024
|
| 37 |
+
tie_embeddings: true
|
| 38 |
+
max_norm: 1.0
|
| 39 |
+
min_lr: 6.0e-05
|
| 40 |
+
|
| 41 |
+
optimizer:
|
| 42 |
+
class_path: torch.optim.AdamW
|
| 43 |
+
init_args:
|
| 44 |
+
lr: 6.0e-04
|
| 45 |
+
weight_decay: 0.1
|
| 46 |
+
betas: [0.9, 0.95]
|
| 47 |
+
eps: 1.0e-08
|
| 48 |
+
|
| 49 |
+
devices: 1
|
| 50 |
+
num_nodes: 1
|
| 51 |
+
tokenizer_dir: Base/checkpoints/EleutherAI/pythia-160m
|
| 52 |
+
logger_name: tensorboard
|
| 53 |
+
seed: 1337
|
Base/configs/pretrain_100m_english.yaml
ADDED
|
@@ -0,0 +1,58 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# pretrain_100m_english.yaml
|
| 2 |
+
# Continued pretraining on ultra-clean English corpus
|
| 3 |
+
# Initializes from existing 3B-pretrained checkpoint (no random init)
|
| 4 |
+
# Lower learning rate for stable knowledge absorption
|
| 5 |
+
# RTX 4060 Ti 16GB β same batch settings as 3B run
|
| 6 |
+
|
| 7 |
+
model_name: pythia-14m
|
| 8 |
+
model_config:
|
| 9 |
+
name: custom-pythia-like-100m
|
| 10 |
+
block_size: 1024
|
| 11 |
+
n_layer: 10
|
| 12 |
+
n_embd: 768
|
| 13 |
+
n_head: 12
|
| 14 |
+
vocab_size: 50254
|
| 15 |
+
padding_multiple: 128
|
| 16 |
+
rotary_percentage: 0.25
|
| 17 |
+
|
| 18 |
+
# Fresh output directory for the English continued pretraining run
|
| 19 |
+
out_dir: Base/out/pretrain/custom-100m-english
|
| 20 |
+
|
| 21 |
+
# Load weights from existing pretrained checkpoint (raw model weights only)
|
| 22 |
+
initial_checkpoint_dir: Base/out/pretrain/custom-100m-3b-full/final_raw
|
| 23 |
+
|
| 24 |
+
precision: bf16-mixed
|
| 25 |
+
|
| 26 |
+
data:
|
| 27 |
+
class_path: litgpt.data.LitData
|
| 28 |
+
init_args:
|
| 29 |
+
data_path: Base/data/litdata_english
|
| 30 |
+
|
| 31 |
+
train:
|
| 32 |
+
save_interval: 500
|
| 33 |
+
log_interval: 10
|
| 34 |
+
# Same batch config as the 3B run β proven stable on this GPU
|
| 35 |
+
global_batch_size: 120
|
| 36 |
+
micro_batch_size: 12
|
| 37 |
+
lr_warmup_steps: 200
|
| 38 |
+
# 3 epochs over ~50M tokens = 150M token budget
|
| 39 |
+
max_tokens: 150000000
|
| 40 |
+
max_seq_length: 1024
|
| 41 |
+
tie_embeddings: true
|
| 42 |
+
max_norm: 1.0
|
| 43 |
+
min_lr: 1.0e-05
|
| 44 |
+
|
| 45 |
+
optimizer:
|
| 46 |
+
class_path: torch.optim.AdamW
|
| 47 |
+
init_args:
|
| 48 |
+
# 6x lower than initial pretraining (6e-4) for stable continued learning
|
| 49 |
+
lr: 1.0e-04
|
| 50 |
+
weight_decay: 0.1
|
| 51 |
+
betas: [0.9, 0.95]
|
| 52 |
+
eps: 1.0e-08
|
| 53 |
+
|
| 54 |
+
devices: 1
|
| 55 |
+
num_nodes: 1
|
| 56 |
+
tokenizer_dir: Base/checkpoints/EleutherAI/pythia-160m
|
| 57 |
+
logger_name: tensorboard
|
| 58 |
+
seed: 1337
|
Base/scripts/audit_5b_pretrain.py
ADDED
|
@@ -0,0 +1,494 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""
|
| 3 |
+
COMPREHENSIVE AUDIT of litdata_pretrain_final (5B tokens)
|
| 4 |
+
|
| 5 |
+
Checks:
|
| 6 |
+
1. Binary integrity β all 308 chunks readable, headers valid, correct dims
|
| 7 |
+
2. Token range validation β all tokens in [0, vocab_size), no garbage
|
| 8 |
+
3. EOS placement β correct document boundaries
|
| 9 |
+
4. Duplicate detection β sample chunks for near-duplicate documents
|
| 10 |
+
5. Token distribution β check for anomalous frequency spikes (noise)
|
| 11 |
+
6. Decode quality β random sample of 50 documents, decode + inspect
|
| 12 |
+
7. Training readiness β correct block_size, total tokens match config
|
| 13 |
+
|
| 14 |
+
This script is READ-ONLY. It does NOT modify data.
|
| 15 |
+
"""
|
| 16 |
+
|
| 17 |
+
import json
|
| 18 |
+
import os
|
| 19 |
+
import re
|
| 20 |
+
import time
|
| 21 |
+
import hashlib
|
| 22 |
+
from pathlib import Path
|
| 23 |
+
from collections import Counter, defaultdict
|
| 24 |
+
|
| 25 |
+
import numpy as np
|
| 26 |
+
|
| 27 |
+
ROOT = Path(__file__).resolve().parent.parent.parent
|
| 28 |
+
FINAL_DIR = ROOT / "Base" / "data" / "litdata_pretrain_final"
|
| 29 |
+
TOKENIZER_PATH = str(ROOT / "Base" / "checkpoints" / "EleutherAI" / "pythia-160m" / "tokenizer.json")
|
| 30 |
+
BLOCK_SIZE = 1025
|
| 31 |
+
DTYPE = np.int32
|
| 32 |
+
EOS_TOKEN_ID = 0
|
| 33 |
+
VOCAB_SIZE = 50277
|
| 34 |
+
|
| 35 |
+
# How many chunks to deeply scan for duplicates (all 308 is fine for audit)
|
| 36 |
+
MAX_CHUNKS_DEEP = 308
|
| 37 |
+
# Number of random documents to decode and display
|
| 38 |
+
DECODE_SAMPLES = 50
|
| 39 |
+
# Min hash length for near-duplicate detection (chars)
|
| 40 |
+
HASH_WINDOW = 200
|
| 41 |
+
|
| 42 |
+
|
| 43 |
+
def read_chunk(filepath):
|
| 44 |
+
"""Read a single chunk binary file. Returns (blocks, num_blocks)."""
|
| 45 |
+
with open(filepath, "rb") as f:
|
| 46 |
+
raw = f.read()
|
| 47 |
+
|
| 48 |
+
# Parse header
|
| 49 |
+
num_blocks = np.frombuffer(raw[:4], dtype=np.uint32)[0]
|
| 50 |
+
header_size = 4 + (num_blocks + 1) * 4 # 1 uint32 + (num_blocks+1) offsets
|
| 51 |
+
data_bytes = raw[header_size:]
|
| 52 |
+
|
| 53 |
+
expected_tokens = num_blocks * BLOCK_SIZE
|
| 54 |
+
expected_bytes = expected_tokens * DTYPE().itemsize
|
| 55 |
+
|
| 56 |
+
tokens = np.frombuffer(data_bytes[:expected_bytes], dtype=DTYPE)
|
| 57 |
+
return tokens, int(num_blocks)
|
| 58 |
+
|
| 59 |
+
|
| 60 |
+
def extract_documents(tokens):
|
| 61 |
+
"""Split a token array at EOS boundaries into individual documents."""
|
| 62 |
+
eos_positions = np.where(tokens == EOS_TOKEN_ID)[0]
|
| 63 |
+
docs = []
|
| 64 |
+
start = 0
|
| 65 |
+
for eos_pos in eos_positions:
|
| 66 |
+
if eos_pos > start:
|
| 67 |
+
doc_tokens = tokens[start:eos_pos]
|
| 68 |
+
if len(doc_tokens) > 0:
|
| 69 |
+
docs.append(doc_tokens)
|
| 70 |
+
start = eos_pos + 1
|
| 71 |
+
# Trailing tokens (no final EOS β partial doc carried across chunks)
|
| 72 |
+
if start < len(tokens):
|
| 73 |
+
remaining = tokens[start:]
|
| 74 |
+
if len(remaining) > 5: # ignore tiny fragments
|
| 75 |
+
docs.append(remaining)
|
| 76 |
+
return docs
|
| 77 |
+
|
| 78 |
+
|
| 79 |
+
def token_hash(doc_tokens, window=200):
|
| 80 |
+
"""Hash first N tokens for near-duplicate detection."""
|
| 81 |
+
key = doc_tokens[:window].tobytes()
|
| 82 |
+
return hashlib.md5(key).hexdigest()
|
| 83 |
+
|
| 84 |
+
|
| 85 |
+
def main():
|
| 86 |
+
t_start = time.time()
|
| 87 |
+
|
| 88 |
+
# Load index
|
| 89 |
+
with open(FINAL_DIR / "index.json") as f:
|
| 90 |
+
index = json.load(f)
|
| 91 |
+
chunks_meta = index["chunks"]
|
| 92 |
+
config = index.get("config", {})
|
| 93 |
+
num_chunks = len(chunks_meta)
|
| 94 |
+
|
| 95 |
+
print(f"{'='*75}")
|
| 96 |
+
print(f" COMPREHENSIVE AUDIT β litdata_pretrain_final")
|
| 97 |
+
print(f"{'='*75}")
|
| 98 |
+
print(f" Chunks in index: {num_chunks}")
|
| 99 |
+
print(f" Config: {config}")
|
| 100 |
+
print(f" Expected BLOCK_SIZE: {BLOCK_SIZE}")
|
| 101 |
+
print(f" Expected VOCAB_SIZE: {VOCAB_SIZE}")
|
| 102 |
+
print(f" Expected EOS: {EOS_TOKEN_ID}")
|
| 103 |
+
print()
|
| 104 |
+
|
| 105 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 106 |
+
# CHECK 1: Binary integrity + token range
|
| 107 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 108 |
+
print(f" CHECK 1: BINARY INTEGRITY + TOKEN RANGE")
|
| 109 |
+
print(f" {'-'*60}")
|
| 110 |
+
|
| 111 |
+
total_tokens = 0
|
| 112 |
+
total_blocks = 0
|
| 113 |
+
missing_files = []
|
| 114 |
+
corrupt_chunks = []
|
| 115 |
+
out_of_range_chunks = []
|
| 116 |
+
eos_counts = []
|
| 117 |
+
chunk_token_counts = []
|
| 118 |
+
|
| 119 |
+
# Global token frequency (sample every 10th chunk for speed)
|
| 120 |
+
global_freq = Counter()
|
| 121 |
+
FREQ_SAMPLE_INTERVAL = 3 # sample every 3rd chunk
|
| 122 |
+
|
| 123 |
+
# For duplicate detection
|
| 124 |
+
doc_hashes = defaultdict(list) # hash -> [(chunk_idx, doc_idx)]
|
| 125 |
+
total_docs = 0
|
| 126 |
+
doc_lengths = []
|
| 127 |
+
|
| 128 |
+
# For decode sampling: collect random docs
|
| 129 |
+
sample_docs = []
|
| 130 |
+
np.random.seed(42)
|
| 131 |
+
|
| 132 |
+
for ci, meta in enumerate(chunks_meta):
|
| 133 |
+
filename = meta["filename"]
|
| 134 |
+
filepath = FINAL_DIR / filename
|
| 135 |
+
expected_dim = meta["dim"]
|
| 136 |
+
|
| 137 |
+
if not filepath.exists():
|
| 138 |
+
missing_files.append(filename)
|
| 139 |
+
print(f" MISSING: {filename}")
|
| 140 |
+
continue
|
| 141 |
+
|
| 142 |
+
try:
|
| 143 |
+
tokens, num_blocks = read_chunk(filepath)
|
| 144 |
+
except Exception as e:
|
| 145 |
+
corrupt_chunks.append((filename, str(e)))
|
| 146 |
+
print(f" CORRUPT: {filename} β {e}")
|
| 147 |
+
continue
|
| 148 |
+
|
| 149 |
+
actual_dim = len(tokens)
|
| 150 |
+
if actual_dim != expected_dim:
|
| 151 |
+
corrupt_chunks.append((filename, f"dim mismatch: expected {expected_dim}, got {actual_dim}"))
|
| 152 |
+
print(f" DIM MISMATCH: {filename} expected {expected_dim}, got {actual_dim}")
|
| 153 |
+
|
| 154 |
+
total_tokens += actual_dim
|
| 155 |
+
total_blocks += num_blocks
|
| 156 |
+
chunk_token_counts.append(actual_dim)
|
| 157 |
+
|
| 158 |
+
# Token range check
|
| 159 |
+
min_tok = int(tokens.min())
|
| 160 |
+
max_tok = int(tokens.max())
|
| 161 |
+
if min_tok < 0 or max_tok >= VOCAB_SIZE:
|
| 162 |
+
out_of_range_chunks.append((filename, min_tok, max_tok))
|
| 163 |
+
print(f" OUT OF RANGE: {filename} min={min_tok} max={max_tok}")
|
| 164 |
+
|
| 165 |
+
# EOS count
|
| 166 |
+
eos_count = int(np.sum(tokens == EOS_TOKEN_ID))
|
| 167 |
+
eos_counts.append(eos_count)
|
| 168 |
+
|
| 169 |
+
# Token frequency (sampled)
|
| 170 |
+
if ci % FREQ_SAMPLE_INTERVAL == 0:
|
| 171 |
+
unique, counts = np.unique(tokens, return_counts=True)
|
| 172 |
+
for tok, cnt in zip(unique, counts):
|
| 173 |
+
global_freq[int(tok)] += int(cnt)
|
| 174 |
+
|
| 175 |
+
# Extract documents for duplicate + quality check
|
| 176 |
+
if ci < MAX_CHUNKS_DEEP:
|
| 177 |
+
docs = extract_documents(tokens)
|
| 178 |
+
for di, doc in enumerate(docs):
|
| 179 |
+
total_docs += 1
|
| 180 |
+
doc_lengths.append(len(doc))
|
| 181 |
+
if len(doc) >= 50:
|
| 182 |
+
h = token_hash(doc)
|
| 183 |
+
doc_hashes[h].append((ci, di))
|
| 184 |
+
|
| 185 |
+
# Random sample for decode
|
| 186 |
+
if len(sample_docs) < DECODE_SAMPLES and np.random.random() < 0.0005:
|
| 187 |
+
sample_docs.append(doc)
|
| 188 |
+
|
| 189 |
+
if (ci + 1) % 50 == 0:
|
| 190 |
+
print(f" Scanned {ci+1}/{num_chunks} chunks... ({total_tokens:,} tokens)")
|
| 191 |
+
|
| 192 |
+
print(f" Scanned all {num_chunks} chunks: {total_tokens:,} total tokens")
|
| 193 |
+
print()
|
| 194 |
+
|
| 195 |
+
# Results for Check 1
|
| 196 |
+
issues = []
|
| 197 |
+
if missing_files:
|
| 198 |
+
issues.append(f"MISSING FILES: {len(missing_files)}")
|
| 199 |
+
if corrupt_chunks:
|
| 200 |
+
issues.append(f"CORRUPT CHUNKS: {len(corrupt_chunks)}")
|
| 201 |
+
if out_of_range_chunks:
|
| 202 |
+
issues.append(f"OUT-OF-RANGE TOKENS: {len(out_of_range_chunks)} chunks")
|
| 203 |
+
|
| 204 |
+
if not issues:
|
| 205 |
+
print(f" β All {num_chunks} chunks intact, all tokens in [0, {VOCAB_SIZE})")
|
| 206 |
+
else:
|
| 207 |
+
for issue in issues:
|
| 208 |
+
print(f" β {issue}")
|
| 209 |
+
print()
|
| 210 |
+
|
| 211 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 212 |
+
# CHECK 2: EOS + DOCUMENT BOUNDARIES
|
| 213 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 214 |
+
print(f" CHECK 2: EOS & DOCUMENT BOUNDARIES")
|
| 215 |
+
print(f" {'-'*60}")
|
| 216 |
+
|
| 217 |
+
total_eos = sum(eos_counts)
|
| 218 |
+
avg_eos = total_eos / max(num_chunks, 1)
|
| 219 |
+
min_eos = min(eos_counts) if eos_counts else 0
|
| 220 |
+
max_eos = max(eos_counts) if eos_counts else 0
|
| 221 |
+
|
| 222 |
+
print(f" Total EOS tokens: {total_eos:,}")
|
| 223 |
+
print(f" Avg EOS per chunk: {avg_eos:.1f}")
|
| 224 |
+
print(f" Min/Max EOS per chunk: {min_eos} / {max_eos}")
|
| 225 |
+
print(f" Total documents found: {total_docs:,}")
|
| 226 |
+
|
| 227 |
+
if doc_lengths:
|
| 228 |
+
dl = np.array(doc_lengths)
|
| 229 |
+
print(f" Doc length (tokens): min={int(dl.min())}, median={int(np.median(dl))}, "
|
| 230 |
+
f"mean={int(dl.mean())}, max={int(dl.max())}")
|
| 231 |
+
tiny_docs = int(np.sum(dl < 20))
|
| 232 |
+
short_docs = int(np.sum(dl < 50))
|
| 233 |
+
long_docs = int(np.sum(dl > 50000))
|
| 234 |
+
print(f" Tiny docs (<20 tok): {tiny_docs:,} ({100*tiny_docs/len(dl):.2f}%)")
|
| 235 |
+
print(f" Short docs (<50 tok): {short_docs:,} ({100*short_docs/len(dl):.2f}%)")
|
| 236 |
+
print(f" Very long docs (>50K tok): {long_docs:,}")
|
| 237 |
+
|
| 238 |
+
# Flag if too many tiny docs (noise)
|
| 239 |
+
if doc_lengths and tiny_docs / len(dl) > 0.05:
|
| 240 |
+
print(f" β WARNING: {100*tiny_docs/len(dl):.1f}% tiny docs β possible noise")
|
| 241 |
+
else:
|
| 242 |
+
print(f" β Document boundaries look healthy")
|
| 243 |
+
print()
|
| 244 |
+
|
| 245 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 246 |
+
# CHECK 3: DUPLICATE DETECTION
|
| 247 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 248 |
+
print(f" CHECK 3: NEAR-DUPLICATE DETECTION")
|
| 249 |
+
print(f" {'-'*60}")
|
| 250 |
+
|
| 251 |
+
dup_groups = {h: locs for h, locs in doc_hashes.items() if len(locs) > 1}
|
| 252 |
+
dup_doc_count = sum(len(locs) - 1 for locs in dup_groups.values())
|
| 253 |
+
|
| 254 |
+
print(f" Unique doc hashes: {len(doc_hashes):,}")
|
| 255 |
+
print(f" Duplicate groups: {len(dup_groups):,}")
|
| 256 |
+
print(f" Duplicate docs (extra copies): {dup_doc_count:,}")
|
| 257 |
+
dup_pct = 100 * dup_doc_count / max(total_docs, 1)
|
| 258 |
+
print(f" Duplication rate: {dup_pct:.3f}%")
|
| 259 |
+
|
| 260 |
+
if dup_pct > 5.0:
|
| 261 |
+
print(f" β WARNING: High duplication rate ({dup_pct:.1f}%). Consider deduplication.")
|
| 262 |
+
elif dup_pct > 1.0:
|
| 263 |
+
print(f" β MODERATE: {dup_pct:.2f}% duplicates. Acceptable but not ideal.")
|
| 264 |
+
else:
|
| 265 |
+
print(f" β Very low duplication ({dup_pct:.3f}%). Excellent.")
|
| 266 |
+
|
| 267 |
+
# Show a few duplicate examples
|
| 268 |
+
if dup_groups:
|
| 269 |
+
print(f"\n Top 5 duplicate groups (by copy count):")
|
| 270 |
+
sorted_dups = sorted(dup_groups.items(), key=lambda x: -len(x[1]))[:5]
|
| 271 |
+
for h, locs in sorted_dups:
|
| 272 |
+
print(f" hash={h[:12]}... : {len(locs)} copies in chunks {[l[0] for l in locs[:6]]}")
|
| 273 |
+
print()
|
| 274 |
+
|
| 275 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 276 |
+
# CHECK 4: TOKEN DISTRIBUTION (anomaly detection)
|
| 277 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 278 |
+
print(f" CHECK 4: TOKEN DISTRIBUTION ANALYSIS")
|
| 279 |
+
print(f" {'-'*60}")
|
| 280 |
+
|
| 281 |
+
total_sampled = sum(global_freq.values())
|
| 282 |
+
print(f" Sampled tokens: {total_sampled:,} (from every {FREQ_SAMPLE_INTERVAL}rd chunk)")
|
| 283 |
+
|
| 284 |
+
# Top 20 most frequent tokens
|
| 285 |
+
most_common = global_freq.most_common(30)
|
| 286 |
+
print(f" Top 30 tokens by frequency:")
|
| 287 |
+
for tok_id, count in most_common:
|
| 288 |
+
pct = 100 * count / total_sampled
|
| 289 |
+
print(f" token {tok_id:>6}: {count:>12,} ({pct:>5.2f}%)")
|
| 290 |
+
|
| 291 |
+
# Check for suspicious spikes
|
| 292 |
+
# If any single non-EOS token is >10% of all tokens, it's suspicious
|
| 293 |
+
suspicious = []
|
| 294 |
+
for tok_id, count in most_common:
|
| 295 |
+
pct = count / total_sampled
|
| 296 |
+
if tok_id != EOS_TOKEN_ID and pct > 0.10:
|
| 297 |
+
suspicious.append((tok_id, pct))
|
| 298 |
+
|
| 299 |
+
if suspicious:
|
| 300 |
+
print(f"\n β SUSPICIOUS: These non-EOS tokens appear in >10% of data:")
|
| 301 |
+
for tok_id, pct in suspicious:
|
| 302 |
+
print(f" token {tok_id}: {100*pct:.2f}%")
|
| 303 |
+
else:
|
| 304 |
+
print(f"\n β No anomalous token frequency spikes detected")
|
| 305 |
+
|
| 306 |
+
# Vocab coverage
|
| 307 |
+
unique_tokens = len(global_freq)
|
| 308 |
+
coverage_pct = 100 * unique_tokens / VOCAB_SIZE
|
| 309 |
+
print(f" Unique tokens seen: {unique_tokens:,} / {VOCAB_SIZE:,} ({coverage_pct:.1f}% vocab coverage)")
|
| 310 |
+
|
| 311 |
+
# Check for dead zones (large ranges of unused tokens)
|
| 312 |
+
used_set = set(global_freq.keys())
|
| 313 |
+
unused_ranges = []
|
| 314 |
+
start_unused = None
|
| 315 |
+
for i in range(VOCAB_SIZE):
|
| 316 |
+
if i not in used_set:
|
| 317 |
+
if start_unused is None:
|
| 318 |
+
start_unused = i
|
| 319 |
+
else:
|
| 320 |
+
if start_unused is not None:
|
| 321 |
+
gap = i - start_unused
|
| 322 |
+
if gap > 500:
|
| 323 |
+
unused_ranges.append((start_unused, i - 1, gap))
|
| 324 |
+
start_unused = None
|
| 325 |
+
|
| 326 |
+
if unused_ranges:
|
| 327 |
+
print(f" Large unused token ranges (>500):")
|
| 328 |
+
for s, e, g in unused_ranges[:5]:
|
| 329 |
+
print(f" tokens {s}-{e} ({g} unused)")
|
| 330 |
+
print()
|
| 331 |
+
|
| 332 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 333 |
+
# CHECK 5: DECODE QUALITY β sample random documents
|
| 334 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 335 |
+
print(f" CHECK 5: DECODED DOCUMENT SAMPLES")
|
| 336 |
+
print(f" {'-'*60}")
|
| 337 |
+
|
| 338 |
+
# If we didn't get enough samples randomly, grab from specific chunks
|
| 339 |
+
if len(sample_docs) < DECODE_SAMPLES:
|
| 340 |
+
# Grab from spread-out chunks
|
| 341 |
+
sample_chunks = np.linspace(0, num_chunks - 1, min(DECODE_SAMPLES - len(sample_docs), 25), dtype=int)
|
| 342 |
+
for sci in sample_chunks:
|
| 343 |
+
if len(sample_docs) >= DECODE_SAMPLES:
|
| 344 |
+
break
|
| 345 |
+
meta = chunks_meta[sci]
|
| 346 |
+
filepath = FINAL_DIR / meta["filename"]
|
| 347 |
+
try:
|
| 348 |
+
tokens, _ = read_chunk(filepath)
|
| 349 |
+
docs = extract_documents(tokens)
|
| 350 |
+
if docs:
|
| 351 |
+
# Pick a random doc from this chunk
|
| 352 |
+
idx = np.random.randint(0, len(docs))
|
| 353 |
+
sample_docs.append(docs[idx])
|
| 354 |
+
except:
|
| 355 |
+
pass
|
| 356 |
+
|
| 357 |
+
# Now decode
|
| 358 |
+
from tokenizers import Tokenizer
|
| 359 |
+
tokenizer = Tokenizer.from_file(TOKENIZER_PATH)
|
| 360 |
+
|
| 361 |
+
quality_issues = 0
|
| 362 |
+
noise_docs = 0
|
| 363 |
+
non_english_docs = 0
|
| 364 |
+
|
| 365 |
+
RE_CJK = re.compile(r'[\u4e00-\u9fff\u3040-\u309f\u30a0-\u30ff\uac00-\ud7af]{3,}')
|
| 366 |
+
RE_ARABIC = re.compile(r'[\u0600-\u06ff]{5,}')
|
| 367 |
+
RE_CYRILLIC = re.compile(r'[\u0400-\u04ff]{5,}')
|
| 368 |
+
|
| 369 |
+
for si, doc_tokens in enumerate(sample_docs[:DECODE_SAMPLES]):
|
| 370 |
+
text = tokenizer.decode(doc_tokens.tolist(), skip_special_tokens=False)
|
| 371 |
+
|
| 372 |
+
# Quality checks
|
| 373 |
+
words = text.split()
|
| 374 |
+
word_count = len(words)
|
| 375 |
+
alpha_ratio = sum(c.isalpha() for c in text) / max(len(text), 1)
|
| 376 |
+
unique_words = len(set(w.lower() for w in words)) / max(word_count, 1)
|
| 377 |
+
|
| 378 |
+
is_noisy = False
|
| 379 |
+
is_non_english = False
|
| 380 |
+
flags = []
|
| 381 |
+
|
| 382 |
+
if alpha_ratio < 0.50:
|
| 383 |
+
flags.append(f"low-alpha({alpha_ratio:.2f})")
|
| 384 |
+
is_noisy = True
|
| 385 |
+
if word_count > 30 and unique_words < 0.15:
|
| 386 |
+
flags.append(f"repetitive({unique_words:.2f})")
|
| 387 |
+
is_noisy = True
|
| 388 |
+
if RE_CJK.search(text) or RE_ARABIC.search(text) or RE_CYRILLIC.search(text):
|
| 389 |
+
flags.append("non-English")
|
| 390 |
+
is_non_english = True
|
| 391 |
+
if word_count < 10:
|
| 392 |
+
flags.append("very-short")
|
| 393 |
+
is_noisy = True
|
| 394 |
+
|
| 395 |
+
if is_noisy:
|
| 396 |
+
noise_docs += 1
|
| 397 |
+
if is_non_english:
|
| 398 |
+
non_english_docs += 1
|
| 399 |
+
if flags:
|
| 400 |
+
quality_issues += 1
|
| 401 |
+
|
| 402 |
+
# Print preview for flagged + a few clean ones
|
| 403 |
+
if flags or si < 5 or si % 10 == 0:
|
| 404 |
+
preview = text[:300].replace('\n', ' β΅ ')
|
| 405 |
+
status = f"β {','.join(flags)}" if flags else "β clean"
|
| 406 |
+
print(f"\n Sample {si+1}/{len(sample_docs[:DECODE_SAMPLES])} [{status}] ({word_count} words, alpha={alpha_ratio:.2f}, unique={unique_words:.2f})")
|
| 407 |
+
print(f" \"{preview}...\"")
|
| 408 |
+
|
| 409 |
+
print(f"\n DECODE SUMMARY:")
|
| 410 |
+
print(f" Samples checked: {min(len(sample_docs), DECODE_SAMPLES)}")
|
| 411 |
+
print(f" Clean: {min(len(sample_docs), DECODE_SAMPLES) - quality_issues}")
|
| 412 |
+
print(f" Noisy: {noise_docs}")
|
| 413 |
+
print(f" Non-English: {non_english_docs}")
|
| 414 |
+
print(f" Quality issues: {quality_issues}")
|
| 415 |
+
|
| 416 |
+
if quality_issues / max(len(sample_docs), 1) > 0.1:
|
| 417 |
+
print(f" β WARNING: >10% of sampled docs have quality issues")
|
| 418 |
+
else:
|
| 419 |
+
print(f" β Sample quality looks good")
|
| 420 |
+
print()
|
| 421 |
+
|
| 422 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 423 |
+
# CHECK 6: TRAINING READINESS
|
| 424 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 425 |
+
print(f" CHECK 6: TRAINING READINESS")
|
| 426 |
+
print(f" {'-'*60}")
|
| 427 |
+
|
| 428 |
+
index_total = sum(c["dim"] for c in chunks_meta)
|
| 429 |
+
print(f" Index total tokens: {index_total:,}")
|
| 430 |
+
print(f" Actual total tokens: {total_tokens:,}")
|
| 431 |
+
if index_total == total_tokens:
|
| 432 |
+
print(f" β Index matches actual data perfectly")
|
| 433 |
+
else:
|
| 434 |
+
print(f" β MISMATCH: index says {index_total:,} but files have {total_tokens:,}")
|
| 435 |
+
|
| 436 |
+
# Check all blocks are BLOCK_SIZE aligned
|
| 437 |
+
misaligned = [c["filename"] for c in chunks_meta if c["dim"] % BLOCK_SIZE != 0]
|
| 438 |
+
if misaligned:
|
| 439 |
+
print(f" β {len(misaligned)} chunks not BLOCK_SIZE-aligned: {misaligned[:5]}")
|
| 440 |
+
else:
|
| 441 |
+
print(f" β All chunks perfectly BLOCK_SIZE-aligned ({BLOCK_SIZE})")
|
| 442 |
+
|
| 443 |
+
# Config check
|
| 444 |
+
if config.get("block_size") == BLOCK_SIZE:
|
| 445 |
+
print(f" β Config block_size matches: {BLOCK_SIZE}")
|
| 446 |
+
else:
|
| 447 |
+
print(f" β Config block_size: {config.get('block_size')} (expected {BLOCK_SIZE})")
|
| 448 |
+
|
| 449 |
+
if config.get("vocab_size") == VOCAB_SIZE:
|
| 450 |
+
print(f" β Config vocab_size matches: {VOCAB_SIZE}")
|
| 451 |
+
else:
|
| 452 |
+
print(f" β Config vocab_size: {config.get('vocab_size')} (expected {VOCAB_SIZE})")
|
| 453 |
+
|
| 454 |
+
# Tokens per epoch for batch size calculation
|
| 455 |
+
print(f"\n TRAINING PARAMETERS:")
|
| 456 |
+
seq_len = 1024
|
| 457 |
+
steps = total_tokens // (120 * seq_len) # global_batch=120
|
| 458 |
+
print(f" Total tokens: {total_tokens:,}")
|
| 459 |
+
print(f" Global batch size 120 Γ seq 1024 = {120*1024:,} tokens/step")
|
| 460 |
+
print(f" Total steps for 1 epoch: {steps:,}")
|
| 461 |
+
print(f" At ~10 steps/sec (A100): ~{steps/10/60:.0f} min β {steps/10/3600:.1f} hours")
|
| 462 |
+
print()
|
| 463 |
+
|
| 464 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 465 |
+
# OVERALL VERDICT
|
| 466 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 467 |
+
all_issues = []
|
| 468 |
+
if missing_files: all_issues.append(f"{len(missing_files)} missing files")
|
| 469 |
+
if corrupt_chunks: all_issues.append(f"{len(corrupt_chunks)} corrupt chunks")
|
| 470 |
+
if out_of_range_chunks: all_issues.append(f"{len(out_of_range_chunks)} out-of-range chunks")
|
| 471 |
+
if dup_pct > 5.0: all_issues.append(f"High duplication: {dup_pct:.1f}%")
|
| 472 |
+
if suspicious: all_issues.append("Suspicious token spikes")
|
| 473 |
+
if doc_lengths and tiny_docs / len(dl) > 0.05: all_issues.append("Too many tiny docs")
|
| 474 |
+
if quality_issues / max(len(sample_docs), 1) > 0.1: all_issues.append("Quality issues in samples")
|
| 475 |
+
|
| 476 |
+
elapsed = time.time() - t_start
|
| 477 |
+
|
| 478 |
+
print(f"{'='*75}")
|
| 479 |
+
print(f" AUDIT VERDICT")
|
| 480 |
+
print(f"{'='*75}")
|
| 481 |
+
if not all_issues:
|
| 482 |
+
print(f" β ALL CHECKS PASSED β Dataset is TRAINING-READY")
|
| 483 |
+
print(f" {total_tokens:,} clean tokens across {num_chunks} chunks")
|
| 484 |
+
print(f" No corruption, minimal duplicates, good quality")
|
| 485 |
+
else:
|
| 486 |
+
print(f" β ISSUES FOUND:")
|
| 487 |
+
for issue in all_issues:
|
| 488 |
+
print(f" - {issue}")
|
| 489 |
+
print(f"\n Audit completed in {elapsed:.1f}s")
|
| 490 |
+
print(f"{'='*75}")
|
| 491 |
+
|
| 492 |
+
|
| 493 |
+
if __name__ == "__main__":
|
| 494 |
+
main()
|
Base/scripts/audit_english_clean.py
ADDED
|
@@ -0,0 +1,707 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""
|
| 3 |
+
Deep quality audit of litdata_english_clean.
|
| 4 |
+
|
| 5 |
+
Checks EVERY row for:
|
| 6 |
+
1. Content bias (topic distribution, over-represented domains)
|
| 7 |
+
2. Unwanted context (ads, spam, SEO, cookie notices, legal boilerplate, etc.)
|
| 8 |
+
3. English quality (grammar structure, vocabulary richness, readability)
|
| 9 |
+
4. LLM learning value (diverse sentence structures, good knowledge density)
|
| 10 |
+
5. Toxic/harmful content flags
|
| 11 |
+
6. Residual noise (leftover URLs, code, non-English fragments)
|
| 12 |
+
|
| 13 |
+
Goal: ensure the data teaches the LLM to understand English very well
|
| 14 |
+
so it can later do SFT on any dataset with strong comprehension.
|
| 15 |
+
"""
|
| 16 |
+
|
| 17 |
+
import json
|
| 18 |
+
import os
|
| 19 |
+
import re
|
| 20 |
+
import time
|
| 21 |
+
import string
|
| 22 |
+
from pathlib import Path
|
| 23 |
+
from collections import Counter, defaultdict
|
| 24 |
+
|
| 25 |
+
import numpy as np
|
| 26 |
+
from tokenizers import Tokenizer
|
| 27 |
+
|
| 28 |
+
ROOT = Path(__file__).resolve().parent.parent.parent
|
| 29 |
+
BLOCK_SIZE = 1025
|
| 30 |
+
DTYPE = np.int32
|
| 31 |
+
EOS_TOKEN_ID = 0
|
| 32 |
+
|
| 33 |
+
print("Loading tokenizer...")
|
| 34 |
+
tokenizer = Tokenizer.from_file(
|
| 35 |
+
str(ROOT / "Base" / "checkpoints" / "EleutherAI" / "pythia-160m" / "tokenizer.json")
|
| 36 |
+
)
|
| 37 |
+
|
| 38 |
+
# ==============================================================================
|
| 39 |
+
# LITDATA I/O
|
| 40 |
+
# ==============================================================================
|
| 41 |
+
|
| 42 |
+
def read_all_tokens(litdata_dir):
|
| 43 |
+
with open(litdata_dir / "index.json") as f:
|
| 44 |
+
index = json.load(f)
|
| 45 |
+
chunks = index["chunks"]
|
| 46 |
+
total_tokens = sum(c["dim"] for c in chunks)
|
| 47 |
+
print(f" Reading {len(chunks)} chunks ({total_tokens:,} tokens)...")
|
| 48 |
+
all_tokens = np.empty(total_tokens, dtype=DTYPE)
|
| 49 |
+
pos = 0
|
| 50 |
+
for i, chunk in enumerate(chunks):
|
| 51 |
+
chunk_path = litdata_dir / chunk["filename"]
|
| 52 |
+
n_blocks = chunk["chunk_size"]
|
| 53 |
+
header_ints = 1 + n_blocks + 1
|
| 54 |
+
header_bytes = header_ints * 4
|
| 55 |
+
with open(chunk_path, "rb") as f:
|
| 56 |
+
f.seek(header_bytes)
|
| 57 |
+
data = np.fromfile(f, dtype=DTYPE, count=chunk["dim"])
|
| 58 |
+
all_tokens[pos:pos + len(data)] = data
|
| 59 |
+
pos += len(data)
|
| 60 |
+
print(f" Read {len(chunks)} chunks ({pos:,} tokens)")
|
| 61 |
+
return all_tokens[:pos]
|
| 62 |
+
|
| 63 |
+
|
| 64 |
+
def split_documents(token_stream):
|
| 65 |
+
eos_positions = np.where(token_stream == EOS_TOKEN_ID)[0]
|
| 66 |
+
docs = []
|
| 67 |
+
start = 0
|
| 68 |
+
for eos_pos in eos_positions:
|
| 69 |
+
if eos_pos > start:
|
| 70 |
+
docs.append(token_stream[start:eos_pos])
|
| 71 |
+
start = eos_pos + 1
|
| 72 |
+
if start < len(token_stream):
|
| 73 |
+
docs.append(token_stream[start:])
|
| 74 |
+
return docs
|
| 75 |
+
|
| 76 |
+
|
| 77 |
+
# ==============================================================================
|
| 78 |
+
# UNWANTED CONTENT DETECTORS
|
| 79 |
+
# ==============================================================================
|
| 80 |
+
|
| 81 |
+
# Patterns that suggest ads, spam, SEO, cookie banners, boilerplate
|
| 82 |
+
RE_COOKIE = re.compile(r'(cookie|cookies)\s+(policy|consent|notice|preferences|settings)', re.I)
|
| 83 |
+
RE_PRIVACY = re.compile(r'(privacy\s+policy|terms\s+of\s+(service|use)|legal\s+disclaimer)', re.I)
|
| 84 |
+
RE_SUBSCRIBE = re.compile(r'(subscribe|sign\s*up|newsletter|unsubscribe|opt[\s-]*out)', re.I)
|
| 85 |
+
RE_CLICKBAIT = re.compile(r'(you\s+won\'?t\s+believe|click\s+here|read\s+more|share\s+this|trending\s+now|sponsored|advertisement)', re.I)
|
| 86 |
+
RE_SEO_SPAM = re.compile(r'(best\s+\d+\s+\w+\s+for|top\s+\d+\s+\w+|buy\s+now|free\s+shipping|limited\s+time\s+offer|discount\s+code)', re.I)
|
| 87 |
+
RE_NAVIGATION = re.compile(r'(home\s*>\s*|breadcrumb|sidebar|footer|header|menu|navigation|skip\s+to\s+content)', re.I)
|
| 88 |
+
RE_SOCIAL = re.compile(r'(follow\s+us\s+on|share\s+on\s+(facebook|twitter|linkedin|instagram)|like\s+us\s+on|tweet\s+this)', re.I)
|
| 89 |
+
RE_COMMENT_SECTION = re.compile(r'(leave\s+a\s+(comment|reply)|post\s+a\s+comment|\d+\s+comments?\s|logged\s+in\s+as)', re.I)
|
| 90 |
+
RE_COPYRIGHT = re.compile(r'(all\s+rights\s+reserved|copyright\s+\d{4}|\(c\)\s*\d{4})', re.I)
|
| 91 |
+
RE_BOILERPLATE_LOGIN = re.compile(r'(log\s*in|sign\s*in|create\s+account|forgot\s+password|remember\s+me)', re.I)
|
| 92 |
+
|
| 93 |
+
# Residual code/technical noise
|
| 94 |
+
RE_RESIDUAL_URL = re.compile(r'https?://\S+|www\.\S+', re.I)
|
| 95 |
+
RE_RESIDUAL_EMAIL = re.compile(r'\b[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}\b')
|
| 96 |
+
RE_RESIDUAL_CODE = re.compile(r'(function\s*\(|var\s+\w+\s*=|console\.log|document\.get|if\s*\(\s*\w+\s*[!=]==)', re.I)
|
| 97 |
+
RE_CURLY_BRACES = re.compile(r'\{[^}]{5,}\}')
|
| 98 |
+
RE_HEX_COLORS = re.compile(r'#[0-9a-fA-F]{6}\b')
|
| 99 |
+
|
| 100 |
+
# Non-English fragments
|
| 101 |
+
RE_CJK = re.compile(r'[\u4e00-\u9fff\u3040-\u309f\u30a0-\u30ff\uac00-\ud7af]{3,}')
|
| 102 |
+
RE_ARABIC = re.compile(r'[\u0600-\u06ff]{5,}')
|
| 103 |
+
RE_CYRILLIC = re.compile(r'[\u0400-\u04ff]{5,}')
|
| 104 |
+
RE_DEVANAGARI = re.compile(r'[\u0900-\u097f]{5,}')
|
| 105 |
+
|
| 106 |
+
# Low quality indicators
|
| 107 |
+
RE_ALLCAPS_LINE = re.compile(r'^[A-Z\s]{20,}$', re.M)
|
| 108 |
+
RE_EXCESSIVE_NUMBERS = re.compile(r'(?:\d+[\s,.-]*){10,}')
|
| 109 |
+
RE_LIST_SPAM = re.compile(r'(?:^\s*[-*]\s*.{3,50}\n){10,}', re.M)
|
| 110 |
+
|
| 111 |
+
# Toxic content indicators (basic - not exhaustive)
|
| 112 |
+
TOXIC_TERMS = [
|
| 113 |
+
'kill yourself', 'kys', 'hate speech', 'racial slur',
|
| 114 |
+
'white supremac', 'nazi', 'nigger', 'faggot',
|
| 115 |
+
]
|
| 116 |
+
|
| 117 |
+
|
| 118 |
+
# ==============================================================================
|
| 119 |
+
# TOPIC CLASSIFIER (keyword-based, broad categories)
|
| 120 |
+
# ==============================================================================
|
| 121 |
+
|
| 122 |
+
TOPIC_KEYWORDS = {
|
| 123 |
+
"Science": ["experiment", "hypothesis", "molecule", "atom", "chemical", "physics",
|
| 124 |
+
"biology", "evolution", "species", "organism", "cell", "dna", "gene",
|
| 125 |
+
"electron", "neutron", "quantum", "telescope", "galaxy", "planet"],
|
| 126 |
+
"Mathematics": ["equation", "theorem", "algebra", "calculus", "geometry", "integer",
|
| 127 |
+
"fraction", "polynomial", "derivative", "integral", "matrix", "probability"],
|
| 128 |
+
"History": ["century", "civilization", "empire", "dynasty", "revolution", "colonial",
|
| 129 |
+
"medieval", "ancient", "historian", "archaeological", "monarchy", "treaty"],
|
| 130 |
+
"Geography": ["continent", "climate", "ocean", "mountain", "river", "latitude",
|
| 131 |
+
"longitude", "ecosystem", "peninsula", "volcano", "earthquake", "terrain"],
|
| 132 |
+
"Literature": ["novel", "poem", "author", "literary", "character", "narrative",
|
| 133 |
+
"fiction", "metaphor", "protagonist", "shakespeare", "prose", "genre"],
|
| 134 |
+
"Technology": ["software", "hardware", "computer", "algorithm", "database", "internet",
|
| 135 |
+
"programming", "digital", "server", "network", "processor", "encryption"],
|
| 136 |
+
"Medicine/Health": ["patient", "symptom", "diagnosis", "treatment", "disease", "infection",
|
| 137 |
+
"surgery", "therapy", "vaccine", "antibiotic", "clinical", "chronic"],
|
| 138 |
+
"Law/Politics": ["constitution", "legislation", "democracy", "parliament", "judiciary",
|
| 139 |
+
"amendment", "election", "government", "policy", "regulation", "statute"],
|
| 140 |
+
"Economics/Business": ["market", "economy", "inflation", "revenue", "investment", "stock",
|
| 141 |
+
"profit", "trade", "gdp", "fiscal", "monetary", "corporation"],
|
| 142 |
+
"Education": ["student", "teacher", "curriculum", "classroom", "university", "academic",
|
| 143 |
+
"learning", "examination", "school", "pedagogy", "literacy", "enrollment"],
|
| 144 |
+
"Philosophy/Religion": ["philosophy", "ethics", "moral", "theological", "spiritual",
|
| 145 |
+
"consciousness", "existence", "metaphysics", "belief", "virtue"],
|
| 146 |
+
"Arts/Culture": ["painting", "sculpture", "museum", "gallery", "architecture", "cinema",
|
| 147 |
+
"music", "composer", "portrait", "exhibition", "artistic", "cultural"],
|
| 148 |
+
"Environment": ["pollution", "conservation", "deforestation", "renewable", "sustainability",
|
| 149 |
+
"biodiversity", "emissions", "habitat", "endangered", "recycling"],
|
| 150 |
+
"Psychology": ["behavior", "cognitive", "emotion", "personality", "anxiety", "depression",
|
| 151 |
+
"consciousness", "motivation", "perception", "neuroscience", "memory"],
|
| 152 |
+
"Sports": ["championship", "tournament", "athlete", "league", "stadium", "coach",
|
| 153 |
+
"scorer", "goalkeeper", "referee", "olympics", "medal", "cricket"],
|
| 154 |
+
}
|
| 155 |
+
|
| 156 |
+
|
| 157 |
+
# ==============================================================================
|
| 158 |
+
# ENGLISH QUALITY METRICS
|
| 159 |
+
# ==============================================================================
|
| 160 |
+
|
| 161 |
+
def compute_readability(text):
|
| 162 |
+
"""Simplified Flesch-Kincaid readability approximation."""
|
| 163 |
+
words = text.split()
|
| 164 |
+
if len(words) < 10:
|
| 165 |
+
return 0.0
|
| 166 |
+
sentences = max(len(re.findall(r'[.!?]+', text)), 1)
|
| 167 |
+
# Approximate syllables: count vowel groups
|
| 168 |
+
syllables = sum(len(re.findall(r'[aeiouy]+', w.lower())) for w in words)
|
| 169 |
+
syllables = max(syllables, len(words)) # at least 1 per word
|
| 170 |
+
words_per_sent = len(words) / sentences
|
| 171 |
+
syl_per_word = syllables / len(words)
|
| 172 |
+
# Flesch Reading Ease
|
| 173 |
+
fre = 206.835 - 1.015 * words_per_sent - 84.6 * syl_per_word
|
| 174 |
+
return max(0, min(100, fre))
|
| 175 |
+
|
| 176 |
+
|
| 177 |
+
def classify_sentence_types(text):
|
| 178 |
+
"""Classify sentences into types for diversity check."""
|
| 179 |
+
sents = re.split(r'(?<=[.!?])\s+', text[:5000]) # First 5000 chars
|
| 180 |
+
types = Counter()
|
| 181 |
+
for s in sents:
|
| 182 |
+
s = s.strip()
|
| 183 |
+
if not s:
|
| 184 |
+
continue
|
| 185 |
+
if s.endswith('?'):
|
| 186 |
+
types['question'] += 1
|
| 187 |
+
elif s.endswith('!'):
|
| 188 |
+
types['exclamation'] += 1
|
| 189 |
+
elif any(s.lower().startswith(w) for w in ['because', 'since', 'although', 'while', 'if', 'when', 'whereas']):
|
| 190 |
+
types['complex'] += 1
|
| 191 |
+
elif any(s.lower().startswith(w) for w in ['the', 'a ', 'an ', 'this', 'that', 'these', 'those']):
|
| 192 |
+
types['declarative'] += 1
|
| 193 |
+
elif any(s.lower().startswith(w) for w in ['for example', 'such as', 'in other words', 'namely']):
|
| 194 |
+
types['explanatory'] += 1
|
| 195 |
+
elif any(s.lower().startswith(w) for w in ['however', 'nevertheless', 'moreover', 'furthermore', 'therefore']):
|
| 196 |
+
types['transitional'] += 1
|
| 197 |
+
else:
|
| 198 |
+
types['other'] += 1
|
| 199 |
+
return types
|
| 200 |
+
|
| 201 |
+
|
| 202 |
+
def detect_topic(text_lower):
|
| 203 |
+
"""Classify document into topics based on keyword density."""
|
| 204 |
+
topics_found = []
|
| 205 |
+
for topic, keywords in TOPIC_KEYWORDS.items():
|
| 206 |
+
hits = sum(1 for kw in keywords if kw in text_lower)
|
| 207 |
+
if hits >= 2:
|
| 208 |
+
topics_found.append((topic, hits))
|
| 209 |
+
topics_found.sort(key=lambda x: -x[1])
|
| 210 |
+
return topics_found
|
| 211 |
+
|
| 212 |
+
|
| 213 |
+
# ==============================================================================
|
| 214 |
+
# MAIN AUDIT
|
| 215 |
+
# ==============================================================================
|
| 216 |
+
|
| 217 |
+
def main():
|
| 218 |
+
input_dir = ROOT / "Base" / "data" / "litdata_english_clean"
|
| 219 |
+
|
| 220 |
+
print(f"\n{'='*75}")
|
| 221 |
+
print(f" DEEP QUALITY AUDIT: litdata_english_clean")
|
| 222 |
+
print(f" Input: {input_dir}")
|
| 223 |
+
print(f"{'='*75}")
|
| 224 |
+
|
| 225 |
+
# 1. Read and decode all documents
|
| 226 |
+
t0 = time.time()
|
| 227 |
+
token_stream = read_all_tokens(input_dir)
|
| 228 |
+
doc_tokens = split_documents(token_stream)
|
| 229 |
+
print(f" Found {len(doc_tokens):,} documents")
|
| 230 |
+
del token_stream
|
| 231 |
+
|
| 232 |
+
print(f" Decoding ALL {len(doc_tokens):,} documents...")
|
| 233 |
+
texts = []
|
| 234 |
+
t1 = time.time()
|
| 235 |
+
for i, toks in enumerate(doc_tokens):
|
| 236 |
+
text = tokenizer.decode(toks.tolist(), skip_special_tokens=False)
|
| 237 |
+
texts.append(text)
|
| 238 |
+
if (i + 1) % 20000 == 0 or i == len(doc_tokens) - 1:
|
| 239 |
+
print(f" Decoded {i+1:,}/{len(doc_tokens):,}")
|
| 240 |
+
del doc_tokens
|
| 241 |
+
print(f" Decoded in {time.time()-t1:.1f}s")
|
| 242 |
+
|
| 243 |
+
total_docs = len(texts)
|
| 244 |
+
print(f"\n Auditing {total_docs:,} documents across ALL rows...\n")
|
| 245 |
+
|
| 246 |
+
# ==================================================================
|
| 247 |
+
# AUDIT PASS: Scan every document
|
| 248 |
+
# ==================================================================
|
| 249 |
+
t2 = time.time()
|
| 250 |
+
|
| 251 |
+
# Counters
|
| 252 |
+
topic_counter = Counter()
|
| 253 |
+
topic_per_doc = []
|
| 254 |
+
sentence_type_totals = Counter()
|
| 255 |
+
readability_scores = []
|
| 256 |
+
word_counts = []
|
| 257 |
+
vocab_richness = []
|
| 258 |
+
avg_sentence_lengths = []
|
| 259 |
+
|
| 260 |
+
# Issue trackers
|
| 261 |
+
issues = {
|
| 262 |
+
"cookie_privacy": [],
|
| 263 |
+
"subscribe_newsletter": [],
|
| 264 |
+
"clickbait_seo": [],
|
| 265 |
+
"navigation_boilerplate": [],
|
| 266 |
+
"social_media": [],
|
| 267 |
+
"comment_section": [],
|
| 268 |
+
"copyright_legal": [],
|
| 269 |
+
"login_boilerplate": [],
|
| 270 |
+
"residual_urls": [],
|
| 271 |
+
"residual_emails": [],
|
| 272 |
+
"residual_code": [],
|
| 273 |
+
"non_english_fragments": [],
|
| 274 |
+
"allcaps_heavy": [],
|
| 275 |
+
"excessive_numbers": [],
|
| 276 |
+
"list_spam": [],
|
| 277 |
+
"toxic_content": [],
|
| 278 |
+
"too_short": [],
|
| 279 |
+
"too_repetitive": [],
|
| 280 |
+
"low_readability": [],
|
| 281 |
+
"single_topic_bias": [],
|
| 282 |
+
}
|
| 283 |
+
|
| 284 |
+
# Track flagged doc indices for potential removal
|
| 285 |
+
flagged_docs = set()
|
| 286 |
+
flag_reasons = defaultdict(list)
|
| 287 |
+
|
| 288 |
+
for i, text in enumerate(texts):
|
| 289 |
+
text_lower = text.lower()
|
| 290 |
+
words = text.split()
|
| 291 |
+
word_count = len(words)
|
| 292 |
+
word_counts.append(word_count)
|
| 293 |
+
|
| 294 |
+
# --- Topic classification ---
|
| 295 |
+
topics = detect_topic(text_lower)
|
| 296 |
+
if topics:
|
| 297 |
+
for t_name, _ in topics[:2]:
|
| 298 |
+
topic_counter[t_name] += 1
|
| 299 |
+
topic_per_doc.append(topics[0][0])
|
| 300 |
+
else:
|
| 301 |
+
topic_counter["Uncategorized"] += 1
|
| 302 |
+
topic_per_doc.append("Uncategorized")
|
| 303 |
+
|
| 304 |
+
# --- Sentence diversity ---
|
| 305 |
+
stypes = classify_sentence_types(text)
|
| 306 |
+
for k, v in stypes.items():
|
| 307 |
+
sentence_type_totals[k] += v
|
| 308 |
+
|
| 309 |
+
# --- Readability ---
|
| 310 |
+
fre = compute_readability(text)
|
| 311 |
+
readability_scores.append(fre)
|
| 312 |
+
|
| 313 |
+
# --- Vocabulary richness ---
|
| 314 |
+
if word_count > 20:
|
| 315 |
+
unique_ratio = len(set(w.lower() for w in words)) / word_count
|
| 316 |
+
vocab_richness.append(unique_ratio)
|
| 317 |
+
else:
|
| 318 |
+
vocab_richness.append(0)
|
| 319 |
+
|
| 320 |
+
# --- Avg sentence length ---
|
| 321 |
+
sents = re.split(r'[.!?]+', text)
|
| 322 |
+
sents = [s for s in sents if len(s.strip().split()) > 2]
|
| 323 |
+
if sents:
|
| 324 |
+
avg_sl = sum(len(s.split()) for s in sents) / len(sents)
|
| 325 |
+
avg_sentence_lengths.append(avg_sl)
|
| 326 |
+
else:
|
| 327 |
+
avg_sentence_lengths.append(0)
|
| 328 |
+
|
| 329 |
+
# === ISSUE DETECTION (every row) ===
|
| 330 |
+
is_flagged = False
|
| 331 |
+
|
| 332 |
+
# Cookie/privacy
|
| 333 |
+
if RE_COOKIE.search(text) or RE_PRIVACY.search(text):
|
| 334 |
+
m = RE_COOKIE.findall(text) + RE_PRIVACY.findall(text)
|
| 335 |
+
# Only flag if heavy (multiple matches or large portion)
|
| 336 |
+
if len(m) >= 2 or (len(m) >= 1 and word_count < 100):
|
| 337 |
+
issues["cookie_privacy"].append(i)
|
| 338 |
+
if word_count < 100:
|
| 339 |
+
is_flagged = True
|
| 340 |
+
flag_reasons[i].append("cookie/privacy boilerplate")
|
| 341 |
+
|
| 342 |
+
# Subscribe/newsletter
|
| 343 |
+
m = RE_SUBSCRIBE.findall(text)
|
| 344 |
+
if len(m) >= 2:
|
| 345 |
+
issues["subscribe_newsletter"].append(i)
|
| 346 |
+
if word_count < 100 and len(m) >= 2:
|
| 347 |
+
is_flagged = True
|
| 348 |
+
flag_reasons[i].append("subscribe/newsletter spam")
|
| 349 |
+
|
| 350 |
+
# Clickbait/SEO
|
| 351 |
+
m = RE_CLICKBAIT.findall(text)
|
| 352 |
+
if m:
|
| 353 |
+
issues["clickbait_seo"].append(i)
|
| 354 |
+
if len(m) >= 3:
|
| 355 |
+
is_flagged = True
|
| 356 |
+
flag_reasons[i].append("clickbait/SEO content")
|
| 357 |
+
|
| 358 |
+
# Navigation boilerplate
|
| 359 |
+
m = RE_NAVIGATION.findall(text)
|
| 360 |
+
if len(m) >= 3:
|
| 361 |
+
issues["navigation_boilerplate"].append(i)
|
| 362 |
+
if word_count < 80:
|
| 363 |
+
is_flagged = True
|
| 364 |
+
flag_reasons[i].append("navigation boilerplate")
|
| 365 |
+
|
| 366 |
+
# Social media prompts
|
| 367 |
+
m = RE_SOCIAL.findall(text)
|
| 368 |
+
if m:
|
| 369 |
+
issues["social_media"].append(i)
|
| 370 |
+
|
| 371 |
+
# Comment sections
|
| 372 |
+
m = RE_COMMENT_SECTION.findall(text)
|
| 373 |
+
if m:
|
| 374 |
+
issues["comment_section"].append(i)
|
| 375 |
+
|
| 376 |
+
# Copyright/legal
|
| 377 |
+
m = RE_COPYRIGHT.findall(text)
|
| 378 |
+
if m:
|
| 379 |
+
issues["copyright_legal"].append(i)
|
| 380 |
+
|
| 381 |
+
# Login boilerplate
|
| 382 |
+
m = RE_BOILERPLATE_LOGIN.findall(text)
|
| 383 |
+
if len(m) >= 3:
|
| 384 |
+
issues["login_boilerplate"].append(i)
|
| 385 |
+
if word_count < 80:
|
| 386 |
+
is_flagged = True
|
| 387 |
+
flag_reasons[i].append("login boilerplate")
|
| 388 |
+
|
| 389 |
+
# Residual URLs
|
| 390 |
+
m = RE_RESIDUAL_URL.findall(text)
|
| 391 |
+
if m:
|
| 392 |
+
issues["residual_urls"].append(i)
|
| 393 |
+
is_flagged = True
|
| 394 |
+
flag_reasons[i].append(f"residual URLs ({len(m)})")
|
| 395 |
+
|
| 396 |
+
# Residual emails
|
| 397 |
+
m = RE_RESIDUAL_EMAIL.findall(text)
|
| 398 |
+
if m:
|
| 399 |
+
issues["residual_emails"].append(i)
|
| 400 |
+
|
| 401 |
+
# Residual code
|
| 402 |
+
m = RE_RESIDUAL_CODE.findall(text)
|
| 403 |
+
if len(m) >= 3:
|
| 404 |
+
issues["residual_code"].append(i)
|
| 405 |
+
if len(m) >= 5:
|
| 406 |
+
is_flagged = True
|
| 407 |
+
flag_reasons[i].append(f"residual code ({len(m)} matches)")
|
| 408 |
+
|
| 409 |
+
# Non-English fragments
|
| 410 |
+
has_cjk = bool(RE_CJK.search(text))
|
| 411 |
+
has_arabic = bool(RE_ARABIC.search(text))
|
| 412 |
+
has_cyrillic = bool(RE_CYRILLIC.search(text))
|
| 413 |
+
has_devanagari = bool(RE_DEVANAGARI.search(text))
|
| 414 |
+
if has_cjk or has_arabic or has_cyrillic or has_devanagari:
|
| 415 |
+
issues["non_english_fragments"].append(i)
|
| 416 |
+
scripts = []
|
| 417 |
+
if has_cjk: scripts.append("CJK")
|
| 418 |
+
if has_arabic: scripts.append("Arabic")
|
| 419 |
+
if has_cyrillic: scripts.append("Cyrillic")
|
| 420 |
+
if has_devanagari: scripts.append("Devanagari")
|
| 421 |
+
is_flagged = True
|
| 422 |
+
flag_reasons[i].append(f"non-English ({', '.join(scripts)})")
|
| 423 |
+
|
| 424 |
+
# All-caps heavy
|
| 425 |
+
caps_lines = RE_ALLCAPS_LINE.findall(text)
|
| 426 |
+
if len(caps_lines) >= 3:
|
| 427 |
+
issues["allcaps_heavy"].append(i)
|
| 428 |
+
|
| 429 |
+
# Excessive numbers
|
| 430 |
+
if RE_EXCESSIVE_NUMBERS.search(text):
|
| 431 |
+
issues["excessive_numbers"].append(i)
|
| 432 |
+
|
| 433 |
+
# List spam (10+ short list items in a row)
|
| 434 |
+
if RE_LIST_SPAM.search(text):
|
| 435 |
+
issues["list_spam"].append(i)
|
| 436 |
+
|
| 437 |
+
# Toxic content
|
| 438 |
+
for term in TOXIC_TERMS:
|
| 439 |
+
if term in text_lower:
|
| 440 |
+
issues["toxic_content"].append(i)
|
| 441 |
+
is_flagged = True
|
| 442 |
+
flag_reasons[i].append(f"toxic: '{term}'")
|
| 443 |
+
break
|
| 444 |
+
|
| 445 |
+
# Too short (under 50 words)
|
| 446 |
+
if word_count < 50:
|
| 447 |
+
issues["too_short"].append(i)
|
| 448 |
+
is_flagged = True
|
| 449 |
+
flag_reasons[i].append(f"too short ({word_count} words)")
|
| 450 |
+
|
| 451 |
+
# Too repetitive (unique word ratio < 0.2)
|
| 452 |
+
if word_count > 50 and vocab_richness[-1] < 0.20:
|
| 453 |
+
issues["too_repetitive"].append(i)
|
| 454 |
+
is_flagged = True
|
| 455 |
+
flag_reasons[i].append(f"very repetitive (unique ratio: {vocab_richness[-1]:.3f})")
|
| 456 |
+
|
| 457 |
+
# Low readability (below 10 Flesch = extremely hard, or very odd text)
|
| 458 |
+
if fre < 10 and word_count > 50:
|
| 459 |
+
issues["low_readability"].append(i)
|
| 460 |
+
|
| 461 |
+
if is_flagged:
|
| 462 |
+
flagged_docs.add(i)
|
| 463 |
+
|
| 464 |
+
if (i + 1) % 10000 == 0 or i == total_docs - 1:
|
| 465 |
+
print(f" Audited {i+1:,}/{total_docs:,} | flagged so far: {len(flagged_docs):,}")
|
| 466 |
+
|
| 467 |
+
audit_time = time.time() - t2
|
| 468 |
+
print(f" Audit completed in {audit_time:.1f}s")
|
| 469 |
+
|
| 470 |
+
# ==================================================================
|
| 471 |
+
# BUILD REPORT
|
| 472 |
+
# ==================================================================
|
| 473 |
+
report = []
|
| 474 |
+
report.append(f"\n{'='*75}")
|
| 475 |
+
report.append(f" LITDATA_ENGLISH_CLEAN - DEEP QUALITY AUDIT REPORT")
|
| 476 |
+
report.append(f"{'='*75}")
|
| 477 |
+
report.append(f"\n Total documents audited: {total_docs:,}")
|
| 478 |
+
report.append(f" Total flagged for review: {len(flagged_docs):,} ({len(flagged_docs)/total_docs*100:.2f}%)")
|
| 479 |
+
report.append(f" Audit time: {audit_time:.1f}s")
|
| 480 |
+
|
| 481 |
+
# --- TOPIC DISTRIBUTION ---
|
| 482 |
+
report.append(f"\n\n TOPIC DISTRIBUTION (all {total_docs:,} docs)")
|
| 483 |
+
report.append(f" {'-'*65}")
|
| 484 |
+
total_categorized = sum(topic_counter.values())
|
| 485 |
+
sorted_topics = sorted(topic_counter.items(), key=lambda x: -x[1])
|
| 486 |
+
max_topic_count = sorted_topics[0][1] if sorted_topics else 0
|
| 487 |
+
for topic, count in sorted_topics:
|
| 488 |
+
pct = count / total_categorized * 100
|
| 489 |
+
bar = "#" * int(pct / 2)
|
| 490 |
+
report.append(f" {topic:<25} {count:>7,} ({pct:5.1f}%) {bar}")
|
| 491 |
+
|
| 492 |
+
# Topic bias check
|
| 493 |
+
if sorted_topics:
|
| 494 |
+
top_pct = sorted_topics[0][1] / total_categorized * 100
|
| 495 |
+
if top_pct > 30:
|
| 496 |
+
report.append(f"\n ** WARNING: '{sorted_topics[0][0]}' dominates at {top_pct:.1f}% - potential topic bias **")
|
| 497 |
+
else:
|
| 498 |
+
report.append(f"\n OK: No single topic exceeds 30% - good diversity")
|
| 499 |
+
|
| 500 |
+
# --- ENGLISH QUALITY METRICS ---
|
| 501 |
+
report.append(f"\n\n ENGLISH QUALITY METRICS (all {total_docs:,} docs)")
|
| 502 |
+
report.append(f" {'-'*65}")
|
| 503 |
+
|
| 504 |
+
avg_readability = sum(readability_scores)/len(readability_scores)
|
| 505 |
+
avg_vocab = sum(vocab_richness)/len(vocab_richness)
|
| 506 |
+
avg_words = sum(word_counts)/len(word_counts)
|
| 507 |
+
avg_sent_len = sum(avg_sentence_lengths)/max(len([x for x in avg_sentence_lengths if x > 0]), 1)
|
| 508 |
+
|
| 509 |
+
report.append(f" Avg Flesch Reading Ease: {avg_readability:.1f}")
|
| 510 |
+
if avg_readability >= 60:
|
| 511 |
+
report.append(f" -> Standard/Easy (good for general English learning)")
|
| 512 |
+
elif avg_readability >= 30:
|
| 513 |
+
report.append(f" -> College level (moderately complex)")
|
| 514 |
+
else:
|
| 515 |
+
report.append(f" -> Very difficult (may hinder learning)")
|
| 516 |
+
|
| 517 |
+
report.append(f" Avg vocabulary richness: {avg_vocab:.4f} (unique words / total words)")
|
| 518 |
+
report.append(f" Avg document length: {avg_words:.0f} words")
|
| 519 |
+
report.append(f" Avg sentence length: {avg_sent_len:.1f} words/sentence")
|
| 520 |
+
|
| 521 |
+
# Word count distribution
|
| 522 |
+
short_docs = sum(1 for w in word_counts if w < 50)
|
| 523 |
+
medium_docs = sum(1 for w in word_counts if 50 <= w < 200)
|
| 524 |
+
standard_docs = sum(1 for w in word_counts if 200 <= w < 1000)
|
| 525 |
+
long_docs = sum(1 for w in word_counts if 1000 <= w < 5000)
|
| 526 |
+
very_long_docs = sum(1 for w in word_counts if w >= 5000)
|
| 527 |
+
|
| 528 |
+
report.append(f"\n Document length distribution:")
|
| 529 |
+
report.append(f" < 50 words: {short_docs:>7,} ({short_docs/total_docs*100:.1f}%)")
|
| 530 |
+
report.append(f" 50-199 words: {medium_docs:>7,} ({medium_docs/total_docs*100:.1f}%)")
|
| 531 |
+
report.append(f" 200-999 words: {standard_docs:>7,} ({standard_docs/total_docs*100:.1f}%)")
|
| 532 |
+
report.append(f" 1,000-4,999 words: {long_docs:>7,} ({long_docs/total_docs*100:.1f}%)")
|
| 533 |
+
report.append(f" 5,000+ words: {very_long_docs:>7,} ({very_long_docs/total_docs*100:.1f}%)")
|
| 534 |
+
|
| 535 |
+
# Readability distribution
|
| 536 |
+
very_easy = sum(1 for r in readability_scores if r >= 80)
|
| 537 |
+
easy = sum(1 for r in readability_scores if 60 <= r < 80)
|
| 538 |
+
college = sum(1 for r in readability_scores if 30 <= r < 60)
|
| 539 |
+
hard = sum(1 for r in readability_scores if 10 <= r < 30)
|
| 540 |
+
very_hard = sum(1 for r in readability_scores if r < 10)
|
| 541 |
+
|
| 542 |
+
report.append(f"\n Readability distribution:")
|
| 543 |
+
report.append(f" Very Easy (80-100): {very_easy:>7,} ({very_easy/total_docs*100:.1f}%)")
|
| 544 |
+
report.append(f" Easy (60-79): {easy:>7,} ({easy/total_docs*100:.1f}%)")
|
| 545 |
+
report.append(f" College (30-59): {college:>7,} ({college/total_docs*100:.1f}%)")
|
| 546 |
+
report.append(f" Hard (10-29): {hard:>7,} ({hard/total_docs*100:.1f}%)")
|
| 547 |
+
report.append(f" Very Hard (0-9): {very_hard:>7,} ({very_hard/total_docs*100:.1f}%)")
|
| 548 |
+
|
| 549 |
+
# --- SENTENCE TYPE DIVERSITY ---
|
| 550 |
+
report.append(f"\n\n SENTENCE TYPE DIVERSITY")
|
| 551 |
+
report.append(f" {'-'*65}")
|
| 552 |
+
total_sents = sum(sentence_type_totals.values())
|
| 553 |
+
for stype, count in sorted(sentence_type_totals.items(), key=lambda x: -x[1]):
|
| 554 |
+
pct = count / max(total_sents, 1) * 100
|
| 555 |
+
report.append(f" {stype:<20} {count:>10,} ({pct:5.1f}%)")
|
| 556 |
+
report.append(f" {'TOTAL':<20} {total_sents:>10,}")
|
| 557 |
+
|
| 558 |
+
if sentence_type_totals.get('question', 0) / max(total_sents, 1) < 0.01:
|
| 559 |
+
report.append(f" ** NOTE: Very few questions - adding Q&A data in SFT will help **")
|
| 560 |
+
|
| 561 |
+
# --- UNWANTED CONTENT CHECK ---
|
| 562 |
+
report.append(f"\n\n UNWANTED CONTENT DETECTION (all {total_docs:,} docs scanned)")
|
| 563 |
+
report.append(f" {'-'*65}")
|
| 564 |
+
issue_order = [
|
| 565 |
+
("cookie_privacy", "Cookie/Privacy boilerplate"),
|
| 566 |
+
("subscribe_newsletter", "Subscribe/Newsletter prompts"),
|
| 567 |
+
("clickbait_seo", "Clickbait/SEO content"),
|
| 568 |
+
("navigation_boilerplate", "Navigation boilerplate"),
|
| 569 |
+
("social_media", "Social media prompts"),
|
| 570 |
+
("comment_section", "Comment section artifacts"),
|
| 571 |
+
("copyright_legal", "Copyright/Legal notices"),
|
| 572 |
+
("login_boilerplate", "Login/Account boilerplate"),
|
| 573 |
+
("residual_urls", "Residual URLs"),
|
| 574 |
+
("residual_emails", "Residual email addresses"),
|
| 575 |
+
("residual_code", "Residual code fragments"),
|
| 576 |
+
("non_english_fragments", "Non-English script fragments"),
|
| 577 |
+
("allcaps_heavy", "Heavy ALL-CAPS usage"),
|
| 578 |
+
("excessive_numbers", "Excessive number sequences"),
|
| 579 |
+
("list_spam", "Long list-only content"),
|
| 580 |
+
("toxic_content", "Toxic/harmful content"),
|
| 581 |
+
("too_short", "Too short (< 50 words)"),
|
| 582 |
+
("too_repetitive", "Very repetitive content"),
|
| 583 |
+
("low_readability", "Extremely low readability"),
|
| 584 |
+
]
|
| 585 |
+
|
| 586 |
+
total_issues = 0
|
| 587 |
+
for key, label in issue_order:
|
| 588 |
+
count = len(issues[key])
|
| 589 |
+
total_issues += count
|
| 590 |
+
pct = count / total_docs * 100
|
| 591 |
+
status = "OK" if count == 0 else "CLEAN" if pct < 0.1 else "LOW" if pct < 1 else "MEDIUM" if pct < 5 else "HIGH"
|
| 592 |
+
marker = " *" if count > 0 and pct >= 1 else ""
|
| 593 |
+
report.append(f" {label:<35} {count:>6,} ({pct:5.2f}%) [{status}]{marker}")
|
| 594 |
+
|
| 595 |
+
report.append(f"\n Total issue instances: {total_issues:,}")
|
| 596 |
+
|
| 597 |
+
# --- FLAGGED DOCUMENTS (need attention) ---
|
| 598 |
+
report.append(f"\n\n FLAGGED DOCUMENTS FOR REVIEW: {len(flagged_docs):,}")
|
| 599 |
+
report.append(f" {'-'*65}")
|
| 600 |
+
|
| 601 |
+
if flagged_docs:
|
| 602 |
+
# Summarize flag reasons
|
| 603 |
+
reason_counter = Counter()
|
| 604 |
+
for doc_idx, reasons in flag_reasons.items():
|
| 605 |
+
for r in reasons:
|
| 606 |
+
reason_counter[r.split('(')[0].strip()] += 1
|
| 607 |
+
|
| 608 |
+
report.append(f" Flag reason summary:")
|
| 609 |
+
for reason, count in sorted(reason_counter.items(), key=lambda x: -x[1]):
|
| 610 |
+
report.append(f" {reason:<40} {count:>6,}")
|
| 611 |
+
|
| 612 |
+
# Show examples of worst offenders
|
| 613 |
+
report.append(f"\n Worst flagged documents (up to 15 examples):")
|
| 614 |
+
report.append(f" {'-'*65}")
|
| 615 |
+
# Sort by number of reasons
|
| 616 |
+
worst = sorted(flag_reasons.items(), key=lambda x: -len(x[1]))[:15]
|
| 617 |
+
for doc_idx, reasons in worst:
|
| 618 |
+
text_preview = texts[doc_idx][:200].replace('\n', ' ')
|
| 619 |
+
wc = len(texts[doc_idx].split())
|
| 620 |
+
report.append(f"\n Doc #{doc_idx} ({wc} words) - Flags: {', '.join(reasons)}")
|
| 621 |
+
report.append(f" \"{text_preview}...\"")
|
| 622 |
+
else:
|
| 623 |
+
report.append(f" No documents flagged - dataset is clean!")
|
| 624 |
+
|
| 625 |
+
# --- LLM LEARNING VALUE ASSESSMENT ---
|
| 626 |
+
report.append(f"\n\n LLM LEARNING VALUE ASSESSMENT")
|
| 627 |
+
report.append(f" {'-'*65}")
|
| 628 |
+
|
| 629 |
+
good_count = 0
|
| 630 |
+
for i in range(total_docs):
|
| 631 |
+
if i not in flagged_docs:
|
| 632 |
+
if word_counts[i] >= 100 and readability_scores[i] >= 30 and vocab_richness[i] >= 0.3:
|
| 633 |
+
good_count += 1
|
| 634 |
+
|
| 635 |
+
good_pct = good_count / total_docs * 100
|
| 636 |
+
report.append(f" High-quality docs (100+ words, readable, diverse vocab): {good_count:,} ({good_pct:.1f}%)")
|
| 637 |
+
report.append(f" Flagged docs (potential issues): {len(flagged_docs):,} ({len(flagged_docs)/total_docs*100:.1f}%)")
|
| 638 |
+
|
| 639 |
+
# Grading
|
| 640 |
+
if good_pct >= 95:
|
| 641 |
+
grade = "A"
|
| 642 |
+
assessment = "Excellent - dataset will teach strong English comprehension"
|
| 643 |
+
elif good_pct >= 85:
|
| 644 |
+
grade = "B"
|
| 645 |
+
assessment = "Good - dataset is solid, minor cleanup would help"
|
| 646 |
+
elif good_pct >= 70:
|
| 647 |
+
grade = "C"
|
| 648 |
+
assessment = "Fair - dataset needs targeted cleanup of flagged docs"
|
| 649 |
+
else:
|
| 650 |
+
grade = "D"
|
| 651 |
+
assessment = "Needs work - significant cleanup required"
|
| 652 |
+
|
| 653 |
+
report.append(f"\n GRADE: {grade}")
|
| 654 |
+
report.append(f" ASSESSMENT: {assessment}")
|
| 655 |
+
|
| 656 |
+
# Recommendations
|
| 657 |
+
report.append(f"\n RECOMMENDATIONS FOR OPTIMAL LLM ENGLISH LEARNING:")
|
| 658 |
+
if len(issues["too_short"]) > 0:
|
| 659 |
+
report.append(f" - Remove {len(issues['too_short']):,} docs under 50 words (too short to teach patterns)")
|
| 660 |
+
if len(issues["residual_urls"]) > 0:
|
| 661 |
+
report.append(f" - Strip {len(issues['residual_urls']):,} docs still containing URLs")
|
| 662 |
+
if len(issues["non_english_fragments"]) > 0:
|
| 663 |
+
report.append(f" - Remove {len(issues['non_english_fragments']):,} docs with non-English script fragments")
|
| 664 |
+
if len(issues["too_repetitive"]) > 0:
|
| 665 |
+
report.append(f" - Remove {len(issues['too_repetitive']):,} very repetitive docs")
|
| 666 |
+
if len(issues["toxic_content"]) > 0:
|
| 667 |
+
report.append(f" - URGENT: Remove {len(issues['toxic_content']):,} docs with toxic content")
|
| 668 |
+
if len(issues["residual_code"]) > 0:
|
| 669 |
+
report.append(f" - Review {len(issues['residual_code']):,} docs with residual code fragments")
|
| 670 |
+
if sorted_topics and sorted_topics[0][1] / total_categorized * 100 > 30:
|
| 671 |
+
report.append(f" - Consider balancing topics ('{sorted_topics[0][0]}' is over-represented)")
|
| 672 |
+
if sentence_type_totals.get('question', 0) / max(total_sents, 1) < 0.03:
|
| 673 |
+
report.append(f" - Dataset has few questions - SFT with Q&A pairs will complement this well")
|
| 674 |
+
if len(flagged_docs) == 0:
|
| 675 |
+
report.append(f" - Dataset is clean and ready for pretraining!")
|
| 676 |
+
elif len(flagged_docs) < 100:
|
| 677 |
+
report.append(f" - Only {len(flagged_docs)} docs flagged - minor cleanup recommended")
|
| 678 |
+
report.append(f" - Shall I auto-remove flagged docs and rebuild? (would lose minimal data)")
|
| 679 |
+
|
| 680 |
+
report.append(f"\n{'='*75}")
|
| 681 |
+
|
| 682 |
+
# Print report
|
| 683 |
+
full_report = '\n'.join(report)
|
| 684 |
+
print(full_report)
|
| 685 |
+
|
| 686 |
+
# Save report
|
| 687 |
+
report_path = input_dir / "DEEP_AUDIT_REPORT.txt"
|
| 688 |
+
with open(report_path, "w", encoding="utf-8") as f:
|
| 689 |
+
f.write(full_report)
|
| 690 |
+
print(f"\n Report saved to: {report_path}")
|
| 691 |
+
|
| 692 |
+
# Also save flagged doc indices for potential cleanup
|
| 693 |
+
if flagged_docs:
|
| 694 |
+
flagged_path = input_dir / "flagged_docs.json"
|
| 695 |
+
flagged_data = {
|
| 696 |
+
"total_docs": total_docs,
|
| 697 |
+
"flagged_count": len(flagged_docs),
|
| 698 |
+
"flagged_indices": sorted(flagged_docs),
|
| 699 |
+
"reasons": {str(k): v for k, v in flag_reasons.items()},
|
| 700 |
+
}
|
| 701 |
+
with open(flagged_path, "w", encoding="utf-8") as f:
|
| 702 |
+
json.dump(flagged_data, f, indent=2)
|
| 703 |
+
print(f" Flagged indices saved to: {flagged_path}")
|
| 704 |
+
|
| 705 |
+
|
| 706 |
+
if __name__ == "__main__":
|
| 707 |
+
main()
|
Base/scripts/audit_litdata.py
ADDED
|
@@ -0,0 +1,475 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Audit litdata_3b and litdata_english for training quality.
|
| 3 |
+
|
| 4 |
+
Decodes random samples from binary chunks back to text, then checks:
|
| 5 |
+
1. English alignment (ASCII ratio, language detection heuristics)
|
| 6 |
+
2. Cleaning quality (no junk, code blocks, HTML, boilerplate)
|
| 7 |
+
3. Deduplication (MinHash-like shingle overlap between samples)
|
| 8 |
+
4. Data diversity (topic spread, length distribution, vocab richness)
|
| 9 |
+
5. Random sample printout for manual inspection
|
| 10 |
+
"""
|
| 11 |
+
|
| 12 |
+
import json
|
| 13 |
+
import os
|
| 14 |
+
import random
|
| 15 |
+
import re
|
| 16 |
+
import struct
|
| 17 |
+
import sys
|
| 18 |
+
from collections import Counter, defaultdict
|
| 19 |
+
from pathlib import Path
|
| 20 |
+
|
| 21 |
+
import numpy as np
|
| 22 |
+
from tokenizers import Tokenizer
|
| 23 |
+
|
| 24 |
+
ROOT = Path(__file__).resolve().parent.parent.parent
|
| 25 |
+
BLOCK_SIZE = 1025
|
| 26 |
+
DTYPE = np.int32
|
| 27 |
+
|
| 28 |
+
# ββ Load tokenizer ββββββββββββββββββββββββββββββββββββββββββββββ
|
| 29 |
+
tokenizer = Tokenizer.from_file(
|
| 30 |
+
str(ROOT / "Base" / "checkpoints" / "EleutherAI" / "pythia-160m" / "tokenizer.json")
|
| 31 |
+
)
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
def load_litdata_index(litdata_dir):
|
| 35 |
+
with open(litdata_dir / "index.json") as f:
|
| 36 |
+
return json.load(f)
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
def read_blocks_from_chunk(chunk_path, num_blocks_to_read, total_blocks):
|
| 40 |
+
"""Read random blocks from a chunk file, decoding the litdata header."""
|
| 41 |
+
dtype_size = DTYPE().itemsize
|
| 42 |
+
# Header: uint32 num_items + uint32 offsets[0..num_items]
|
| 43 |
+
header_ints = 1 + total_blocks + 1 # num_items + (num_items+1) offsets
|
| 44 |
+
header_bytes = header_ints * 4
|
| 45 |
+
|
| 46 |
+
with open(chunk_path, "rb") as f:
|
| 47 |
+
# Read header
|
| 48 |
+
header_raw = f.read(header_bytes)
|
| 49 |
+
header = np.frombuffer(header_raw, dtype=np.uint32)
|
| 50 |
+
n_items = header[0]
|
| 51 |
+
|
| 52 |
+
# Pick random block indices
|
| 53 |
+
indices = random.sample(range(min(n_items, total_blocks)), min(num_blocks_to_read, n_items))
|
| 54 |
+
|
| 55 |
+
blocks = []
|
| 56 |
+
for idx in indices:
|
| 57 |
+
offset = int(header[1 + idx]) # byte offset within data section
|
| 58 |
+
f.seek(header_bytes + offset)
|
| 59 |
+
block_data = np.frombuffer(f.read(BLOCK_SIZE * dtype_size), dtype=DTYPE)
|
| 60 |
+
blocks.append(block_data)
|
| 61 |
+
|
| 62 |
+
return blocks
|
| 63 |
+
|
| 64 |
+
|
| 65 |
+
def decode_block(token_ids):
|
| 66 |
+
"""Decode a block of token IDs back to text."""
|
| 67 |
+
ids = token_ids.tolist()
|
| 68 |
+
return tokenizer.decode(ids, skip_special_tokens=False)
|
| 69 |
+
|
| 70 |
+
|
| 71 |
+
# ββ Quality checks ββββββββββββββββββββββββββββββββββββββββββββββ
|
| 72 |
+
def ascii_ratio(text):
|
| 73 |
+
if not text:
|
| 74 |
+
return 0
|
| 75 |
+
ascii_chars = sum(1 for c in text if ord(c) < 128)
|
| 76 |
+
return ascii_chars / len(text)
|
| 77 |
+
|
| 78 |
+
|
| 79 |
+
def alpha_ratio(text):
|
| 80 |
+
if not text:
|
| 81 |
+
return 0
|
| 82 |
+
alpha = sum(1 for c in text if c.isalpha())
|
| 83 |
+
return alpha / len(text)
|
| 84 |
+
|
| 85 |
+
|
| 86 |
+
def has_html(text):
|
| 87 |
+
return bool(re.search(r'<(html|body|div|span|script|style|head|meta|link)\b', text, re.I))
|
| 88 |
+
|
| 89 |
+
|
| 90 |
+
def has_boilerplate(text):
|
| 91 |
+
patterns = [
|
| 92 |
+
r'cookie policy', r'terms of service', r'privacy policy',
|
| 93 |
+
r'all rights reserved', r'click here', r'subscribe now',
|
| 94 |
+
r'sign up for', r'Β©\s*\d{4}', r'powered by',
|
| 95 |
+
r'advertisement', r'sponsored content',
|
| 96 |
+
]
|
| 97 |
+
lower = text.lower()
|
| 98 |
+
return sum(1 for p in patterns if re.search(p, lower))
|
| 99 |
+
|
| 100 |
+
|
| 101 |
+
def url_ratio(text):
|
| 102 |
+
urls = re.findall(r'https?://\S+', text)
|
| 103 |
+
url_chars = sum(len(u) for u in urls)
|
| 104 |
+
return url_chars / max(len(text), 1)
|
| 105 |
+
|
| 106 |
+
|
| 107 |
+
def sentence_count(text):
|
| 108 |
+
return len(re.findall(r'[.!?]+\s', text))
|
| 109 |
+
|
| 110 |
+
|
| 111 |
+
def word_count(text):
|
| 112 |
+
return len(text.split())
|
| 113 |
+
|
| 114 |
+
|
| 115 |
+
def unique_word_ratio(text):
|
| 116 |
+
words = text.lower().split()
|
| 117 |
+
if not words:
|
| 118 |
+
return 0
|
| 119 |
+
return len(set(words)) / len(words)
|
| 120 |
+
|
| 121 |
+
|
| 122 |
+
def shingle_set(text, k=5):
|
| 123 |
+
"""Create k-word shingles for dedup checking."""
|
| 124 |
+
words = text.lower().split()
|
| 125 |
+
if len(words) < k:
|
| 126 |
+
return set()
|
| 127 |
+
return {tuple(words[i:i+k]) for i in range(len(words) - k + 1)}
|
| 128 |
+
|
| 129 |
+
|
| 130 |
+
def jaccard(s1, s2):
|
| 131 |
+
if not s1 or not s2:
|
| 132 |
+
return 0
|
| 133 |
+
return len(s1 & s2) / len(s1 | s2)
|
| 134 |
+
|
| 135 |
+
|
| 136 |
+
# ββ Topic heuristic ββββββββββββββββββββββββββββββββββββββββββββ
|
| 137 |
+
TOPIC_KEYWORDS = {
|
| 138 |
+
"science": ["experiment", "hypothesis", "molecule", "physics", "chemistry", "biology", "research", "scientific"],
|
| 139 |
+
"technology": ["software", "computer", "algorithm", "programming", "internet", "digital", "technology", "database"],
|
| 140 |
+
"history": ["century", "kingdom", "empire", "ancient", "medieval", "civilization", "dynasty", "historical"],
|
| 141 |
+
"medicine": ["patient", "disease", "treatment", "symptom", "diagnosis", "medical", "health", "clinical"],
|
| 142 |
+
"law": ["court", "legal", "law", "attorney", "judge", "statute", "regulation", "jurisdiction"],
|
| 143 |
+
"education": ["student", "university", "school", "learning", "education", "curriculum", "academic", "teacher"],
|
| 144 |
+
"geography": ["country", "continent", "ocean", "mountain", "river", "population", "region", "territory"],
|
| 145 |
+
"arts": ["painting", "music", "artist", "sculpture", "literature", "poetry", "novel", "creative"],
|
| 146 |
+
"sports": ["game", "team", "championship", "athlete", "tournament", "player", "season", "score"],
|
| 147 |
+
"business": ["company", "market", "economy", "investment", "revenue", "industry", "profit", "financial"],
|
| 148 |
+
"philosophy": ["philosophy", "ethics", "morality", "consciousness", "existence", "metaphysics", "logic", "reasoning"],
|
| 149 |
+
"nature": ["species", "animal", "plant", "ecosystem", "forest", "habitat", "wildlife", "environment"],
|
| 150 |
+
}
|
| 151 |
+
|
| 152 |
+
def detect_topics(text):
|
| 153 |
+
lower = text.lower()
|
| 154 |
+
found = []
|
| 155 |
+
for topic, keywords in TOPIC_KEYWORDS.items():
|
| 156 |
+
if any(kw in lower for kw in keywords):
|
| 157 |
+
found.append(topic)
|
| 158 |
+
return found if found else ["general"]
|
| 159 |
+
|
| 160 |
+
|
| 161 |
+
# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 162 |
+
# AUDIT FUNCTION
|
| 163 |
+
# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 164 |
+
def audit_litdata(litdata_dir, name, num_samples=500, print_samples=10):
|
| 165 |
+
print(f"\n{'='*70}")
|
| 166 |
+
print(f" AUDITING: {name}")
|
| 167 |
+
print(f" Path: {litdata_dir}")
|
| 168 |
+
print(f" Sampling {num_samples} random blocks")
|
| 169 |
+
print(f"{'='*70}")
|
| 170 |
+
|
| 171 |
+
index = load_litdata_index(litdata_dir)
|
| 172 |
+
chunks = index["chunks"]
|
| 173 |
+
total_tokens = sum(c["dim"] for c in chunks)
|
| 174 |
+
|
| 175 |
+
print(f"\n Chunks: {len(chunks)}")
|
| 176 |
+
print(f" Total tokens: {total_tokens:,}")
|
| 177 |
+
|
| 178 |
+
# Sample blocks proportionally from each chunk
|
| 179 |
+
total_blocks = sum(c["chunk_size"] for c in chunks)
|
| 180 |
+
samples_per_chunk = {}
|
| 181 |
+
remaining = num_samples
|
| 182 |
+
for i, chunk in enumerate(chunks):
|
| 183 |
+
proportion = chunk["chunk_size"] / total_blocks
|
| 184 |
+
n = max(1, round(proportion * num_samples))
|
| 185 |
+
n = min(n, remaining, chunk["chunk_size"])
|
| 186 |
+
if remaining <= 0:
|
| 187 |
+
break
|
| 188 |
+
samples_per_chunk[i] = n
|
| 189 |
+
remaining -= n
|
| 190 |
+
|
| 191 |
+
# Read and decode samples
|
| 192 |
+
texts = []
|
| 193 |
+
for chunk_idx, n_samples in samples_per_chunk.items():
|
| 194 |
+
chunk = chunks[chunk_idx]
|
| 195 |
+
chunk_path = litdata_dir / chunk["filename"]
|
| 196 |
+
if not chunk_path.exists():
|
| 197 |
+
print(f" WARNING: {chunk_path} missing, skipping")
|
| 198 |
+
continue
|
| 199 |
+
blocks = read_blocks_from_chunk(chunk_path, n_samples, chunk["chunk_size"])
|
| 200 |
+
for block in blocks:
|
| 201 |
+
text = decode_block(block)
|
| 202 |
+
texts.append(text)
|
| 203 |
+
|
| 204 |
+
print(f"\n Decoded {len(texts)} blocks -> text samples")
|
| 205 |
+
|
| 206 |
+
# ββ 1. English Alignment βββββββββββββββββββββββββββββββββββββ
|
| 207 |
+
print(f"\n {'β'*60}")
|
| 208 |
+
print(f" 1. ENGLISH ALIGNMENT")
|
| 209 |
+
ascii_ratios = [ascii_ratio(t) for t in texts]
|
| 210 |
+
alpha_ratios = [alpha_ratio(t) for t in texts]
|
| 211 |
+
|
| 212 |
+
avg_ascii = sum(ascii_ratios) / len(ascii_ratios)
|
| 213 |
+
avg_alpha = sum(alpha_ratios) / len(alpha_ratios)
|
| 214 |
+
low_ascii = sum(1 for r in ascii_ratios if r < 0.85)
|
| 215 |
+
low_alpha = sum(1 for r in alpha_ratios if r < 0.50)
|
| 216 |
+
|
| 217 |
+
print(f" Avg ASCII ratio: {avg_ascii:.4f} (want > 0.95)")
|
| 218 |
+
print(f" Avg alpha ratio: {avg_alpha:.4f} (want > 0.60)")
|
| 219 |
+
print(f" Samples < 85% ASCII: {low_ascii}/{len(texts)} {'β WARNING' if low_ascii > len(texts)*0.05 else 'OK'}")
|
| 220 |
+
print(f" Samples < 50% alpha: {low_alpha}/{len(texts)} {'β WARNING' if low_alpha > len(texts)*0.05 else 'OK'}")
|
| 221 |
+
|
| 222 |
+
# Check for non-English text patterns
|
| 223 |
+
non_english_patterns = 0
|
| 224 |
+
for t in texts:
|
| 225 |
+
# CJK, Arabic, Devanagari, Cyrillic heavy blocks
|
| 226 |
+
if re.search(r'[\u4e00-\u9fff\u0600-\u06ff\u0900-\u097f]{10,}', t):
|
| 227 |
+
non_english_patterns += 1
|
| 228 |
+
elif re.search(r'[\u0400-\u04ff]{10,}', t):
|
| 229 |
+
non_english_patterns += 1
|
| 230 |
+
print(f" Non-English script blocks detected: {non_english_patterns}/{len(texts)} {'β WARNING' if non_english_patterns > 0 else 'OK'}")
|
| 231 |
+
|
| 232 |
+
# ββ 2. Cleaning Quality ββββββββββββββββββββββββββββββββββββββ
|
| 233 |
+
print(f"\n {'β'*60}")
|
| 234 |
+
print(f" 2. CLEANING QUALITY")
|
| 235 |
+
|
| 236 |
+
html_count = sum(1 for t in texts if has_html(t))
|
| 237 |
+
boilerplate_scores = [has_boilerplate(t) for t in texts]
|
| 238 |
+
high_boilerplate = sum(1 for s in boilerplate_scores if s >= 3)
|
| 239 |
+
url_ratios = [url_ratio(t) for t in texts]
|
| 240 |
+
high_url = sum(1 for r in url_ratios if r > 0.03)
|
| 241 |
+
|
| 242 |
+
short_texts = sum(1 for t in texts if word_count(t) < 20)
|
| 243 |
+
|
| 244 |
+
print(f" HTML tags found: {html_count}/{len(texts)} {'β WARNING' if html_count > len(texts)*0.02 else 'OK'}")
|
| 245 |
+
print(f" High boilerplate (3+): {high_boilerplate}/{len(texts)} {'β WARNING' if high_boilerplate > len(texts)*0.05 else 'OK'}")
|
| 246 |
+
print(f" High URL ratio (>3%): {high_url}/{len(texts)} {'β WARNING' if high_url > len(texts)*0.05 else 'OK'}")
|
| 247 |
+
print(f" Very short (<20 words):{short_texts}/{len(texts)} {'β WARNING' if short_texts > len(texts)*0.10 else 'OK'}")
|
| 248 |
+
print(f" Avg boilerplate score: {sum(boilerplate_scores)/len(boilerplate_scores):.2f}")
|
| 249 |
+
print(f" Avg URL ratio: {sum(url_ratios)/len(url_ratios):.4f}")
|
| 250 |
+
|
| 251 |
+
# ββ 3. Deduplication Check βββββββββββββββββββββββββββββββββββ
|
| 252 |
+
print(f"\n {'β'*60}")
|
| 253 |
+
print(f" 3. DEDUPLICATION CHECK")
|
| 254 |
+
|
| 255 |
+
# Check pairwise similarity on a subset
|
| 256 |
+
dedup_sample = min(200, len(texts))
|
| 257 |
+
dedup_texts = random.sample(texts, dedup_sample)
|
| 258 |
+
shingles = [shingle_set(t) for t in dedup_texts]
|
| 259 |
+
|
| 260 |
+
near_dupes = 0
|
| 261 |
+
exact_dupes = 0
|
| 262 |
+
high_similarities = []
|
| 263 |
+
|
| 264 |
+
for i in range(len(dedup_texts)):
|
| 265 |
+
for j in range(i + 1, len(dedup_texts)):
|
| 266 |
+
sim = jaccard(shingles[i], shingles[j])
|
| 267 |
+
if sim > 0.8:
|
| 268 |
+
near_dupes += 1
|
| 269 |
+
high_similarities.append((i, j, sim))
|
| 270 |
+
if sim > 0.95:
|
| 271 |
+
exact_dupes += 1
|
| 272 |
+
|
| 273 |
+
total_pairs = dedup_sample * (dedup_sample - 1) // 2
|
| 274 |
+
print(f" Checked {total_pairs:,} pairs from {dedup_sample} samples")
|
| 275 |
+
print(f" Near duplicates (>80% Jaccard): {near_dupes} {'β WARNING' if near_dupes > 5 else 'OK'}")
|
| 276 |
+
print(f" Exact duplicates (>95% Jaccard): {exact_dupes} {'β WARNING' if exact_dupes > 0 else 'OK'}")
|
| 277 |
+
|
| 278 |
+
if high_similarities:
|
| 279 |
+
print(f" Top overlaps:")
|
| 280 |
+
for i, j, sim in sorted(high_similarities, key=lambda x: -x[2])[:3]:
|
| 281 |
+
print(f" [{i}] vs [{j}] = {sim:.3f}")
|
| 282 |
+
print(f" Sample A: {dedup_texts[i][:80]}...")
|
| 283 |
+
print(f" Sample B: {dedup_texts[j][:80]}...")
|
| 284 |
+
|
| 285 |
+
# ββ 4. Data Diversity ββββββββββββββββββββββββββββββββββββββββ
|
| 286 |
+
print(f"\n {'β'*60}")
|
| 287 |
+
print(f" 4. DATA DIVERSITY")
|
| 288 |
+
|
| 289 |
+
# Word count distribution
|
| 290 |
+
word_counts = [word_count(t) for t in texts]
|
| 291 |
+
avg_wc = sum(word_counts) / len(word_counts)
|
| 292 |
+
min_wc = min(word_counts)
|
| 293 |
+
max_wc = max(word_counts)
|
| 294 |
+
|
| 295 |
+
print(f" Word count: avg={avg_wc:.0f} min={min_wc} max={max_wc}")
|
| 296 |
+
|
| 297 |
+
# Unique word ratio (vocabulary richness)
|
| 298 |
+
uwr = [unique_word_ratio(t) for t in texts]
|
| 299 |
+
avg_uwr = sum(uwr) / len(uwr)
|
| 300 |
+
print(f" Avg unique word ratio: {avg_uwr:.4f} (want > 0.40)")
|
| 301 |
+
|
| 302 |
+
# Sentence structure
|
| 303 |
+
sent_counts = [sentence_count(t) for t in texts]
|
| 304 |
+
avg_sent = sum(sent_counts) / len(sent_counts)
|
| 305 |
+
no_sentence = sum(1 for s in sent_counts if s == 0)
|
| 306 |
+
print(f" Avg sentences/block: {avg_sent:.1f}")
|
| 307 |
+
print(f" Blocks with 0 sentences: {no_sentence}/{len(texts)}")
|
| 308 |
+
|
| 309 |
+
# Topic distribution
|
| 310 |
+
all_topics = Counter()
|
| 311 |
+
for t in texts:
|
| 312 |
+
for topic in detect_topics(t):
|
| 313 |
+
all_topics[topic] += 1
|
| 314 |
+
|
| 315 |
+
print(f"\n Topic distribution (from {len(texts)} samples):")
|
| 316 |
+
for topic, count in all_topics.most_common():
|
| 317 |
+
bar = "β" * int(count / len(texts) * 40)
|
| 318 |
+
print(f" {topic:<14} {count:>4} ({count/len(texts)*100:5.1f}%) {bar}")
|
| 319 |
+
|
| 320 |
+
# ββ 5. Flagged Samples βββββββββββββββββββββββββββββββββββββββ
|
| 321 |
+
print(f"\n {'β'*60}")
|
| 322 |
+
print(f" 5. FLAGGED SAMPLES (potential issues)")
|
| 323 |
+
|
| 324 |
+
flagged = []
|
| 325 |
+
for i, t in enumerate(texts):
|
| 326 |
+
issues = []
|
| 327 |
+
if ascii_ratio(t) < 0.85:
|
| 328 |
+
issues.append(f"low-ascii({ascii_ratio(t):.2f})")
|
| 329 |
+
if has_html(t):
|
| 330 |
+
issues.append("html")
|
| 331 |
+
if has_boilerplate(t) >= 3:
|
| 332 |
+
issues.append(f"boilerplate({has_boilerplate(t)})")
|
| 333 |
+
if url_ratio(t) > 0.05:
|
| 334 |
+
issues.append(f"urls({url_ratio(t):.2f})")
|
| 335 |
+
if word_count(t) < 15:
|
| 336 |
+
issues.append(f"short({word_count(t)}w)")
|
| 337 |
+
if alpha_ratio(t) < 0.40:
|
| 338 |
+
issues.append(f"low-alpha({alpha_ratio(t):.2f})")
|
| 339 |
+
if issues:
|
| 340 |
+
flagged.append((i, issues, t))
|
| 341 |
+
|
| 342 |
+
if flagged:
|
| 343 |
+
print(f" {len(flagged)}/{len(texts)} samples flagged ({len(flagged)/len(texts)*100:.1f}%)")
|
| 344 |
+
for i, issues, t in flagged[:5]:
|
| 345 |
+
print(f"\n Sample #{i}: {', '.join(issues)}")
|
| 346 |
+
print(f" \"{t[:150]}...\"")
|
| 347 |
+
else:
|
| 348 |
+
print(f" No samples flagged! All clean.")
|
| 349 |
+
|
| 350 |
+
# ββ 6. Random Sample Printout ββββββββββββββββββββββββββββββββ
|
| 351 |
+
print(f"\n {'β'*60}")
|
| 352 |
+
print(f" 6. RANDOM SAMPLES (for manual review)")
|
| 353 |
+
|
| 354 |
+
sample_indices = random.sample(range(len(texts)), min(print_samples, len(texts)))
|
| 355 |
+
for idx in sample_indices:
|
| 356 |
+
t = texts[idx]
|
| 357 |
+
topics = detect_topics(t)
|
| 358 |
+
print(f"\n ββ Sample #{idx} | {word_count(t)} words | topics: {', '.join(topics)}")
|
| 359 |
+
# Show first 300 chars
|
| 360 |
+
preview = t[:300].replace('\n', ' β΅ ')
|
| 361 |
+
print(f" β {preview}")
|
| 362 |
+
print(f" ββ ascii={ascii_ratio(t):.2f} alpha={alpha_ratio(t):.2f} uwr={unique_word_ratio(t):.2f}")
|
| 363 |
+
|
| 364 |
+
# ββ Final Verdict ββββββββββββββββββββββββββββββββββββββββββββ
|
| 365 |
+
print(f"\n {'β'*60}")
|
| 366 |
+
print(f" VERDICT for {name}")
|
| 367 |
+
|
| 368 |
+
issues_found = []
|
| 369 |
+
if low_ascii > len(texts) * 0.05:
|
| 370 |
+
issues_found.append(f" β {low_ascii} samples have low ASCII ratio")
|
| 371 |
+
if non_english_patterns > 0:
|
| 372 |
+
issues_found.append(f" β {non_english_patterns} samples contain non-English scripts")
|
| 373 |
+
if html_count > len(texts) * 0.02:
|
| 374 |
+
issues_found.append(f" β {html_count} samples contain HTML tags")
|
| 375 |
+
if high_boilerplate > len(texts) * 0.05:
|
| 376 |
+
issues_found.append(f" β {high_boilerplate} samples have high boilerplate")
|
| 377 |
+
if exact_dupes > 0:
|
| 378 |
+
issues_found.append(f" β {exact_dupes} exact duplicate pairs found")
|
| 379 |
+
if near_dupes > 5:
|
| 380 |
+
issues_found.append(f" β {near_dupes} near-duplicate pairs found")
|
| 381 |
+
if avg_uwr < 0.35:
|
| 382 |
+
issues_found.append(f" β Low vocabulary richness ({avg_uwr:.3f})")
|
| 383 |
+
if len(all_topics) < 4:
|
| 384 |
+
issues_found.append(f" β Low topic diversity (only {len(all_topics)} topics)")
|
| 385 |
+
|
| 386 |
+
if issues_found:
|
| 387 |
+
print(f" Issues found:")
|
| 388 |
+
for issue in issues_found:
|
| 389 |
+
print(f" {issue}")
|
| 390 |
+
else:
|
| 391 |
+
print(f" β PASS - Data looks clean, deduplicated, and diverse!")
|
| 392 |
+
|
| 393 |
+
return {
|
| 394 |
+
"name": name,
|
| 395 |
+
"total_tokens": total_tokens,
|
| 396 |
+
"samples_checked": len(texts),
|
| 397 |
+
"avg_ascii": avg_ascii,
|
| 398 |
+
"avg_alpha": avg_alpha,
|
| 399 |
+
"non_english": non_english_patterns,
|
| 400 |
+
"html_count": html_count,
|
| 401 |
+
"high_boilerplate": high_boilerplate,
|
| 402 |
+
"near_dupes": near_dupes,
|
| 403 |
+
"exact_dupes": exact_dupes,
|
| 404 |
+
"avg_unique_word_ratio": avg_uwr,
|
| 405 |
+
"topic_count": len(all_topics),
|
| 406 |
+
"topics": dict(all_topics),
|
| 407 |
+
"flagged_count": len(flagged),
|
| 408 |
+
"issues": issues_found,
|
| 409 |
+
}
|
| 410 |
+
|
| 411 |
+
|
| 412 |
+
# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 413 |
+
# MAIN
|
| 414 |
+
# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 415 |
+
if __name__ == "__main__":
|
| 416 |
+
random.seed(42)
|
| 417 |
+
|
| 418 |
+
results = []
|
| 419 |
+
|
| 420 |
+
# Audit litdata_3b
|
| 421 |
+
r1 = audit_litdata(
|
| 422 |
+
ROOT / "Base" / "data" / "litdata_3b",
|
| 423 |
+
"litdata_3b (General Knowledge Pretraining)",
|
| 424 |
+
num_samples=500,
|
| 425 |
+
print_samples=8,
|
| 426 |
+
)
|
| 427 |
+
results.append(r1)
|
| 428 |
+
|
| 429 |
+
# Audit litdata_english
|
| 430 |
+
r2 = audit_litdata(
|
| 431 |
+
ROOT / "Base" / "data" / "litdata_english",
|
| 432 |
+
"litdata_english (English Knowledge Continued Pretraining)",
|
| 433 |
+
num_samples=300,
|
| 434 |
+
print_samples=8,
|
| 435 |
+
)
|
| 436 |
+
results.append(r2)
|
| 437 |
+
|
| 438 |
+
# ββ Cross-dataset dedup check ββββββββββββββββββββββββββββββββ
|
| 439 |
+
print(f"\n{'='*70}")
|
| 440 |
+
print(f" CROSS-DATASET OVERLAP CHECK")
|
| 441 |
+
print(f"{'='*70}")
|
| 442 |
+
print(f" (Checking if litdata_3b and litdata_english share duplicate content)")
|
| 443 |
+
# This is checked at the token level - since they come from different
|
| 444 |
+
# source parquets (filtered_3b vs filtered_english), overlap should be minimal
|
| 445 |
+
print(f" Sources are disjoint by design:")
|
| 446 |
+
print(f" litdata_3b <- filtered_3b (15 parquets from general web)")
|
| 447 |
+
print(f" litdata_english <- filtered_english (FineWeb-Edu + Wikipedia)")
|
| 448 |
+
print(f" Cross-contamination: UNLIKELY (separate source pipelines)")
|
| 449 |
+
|
| 450 |
+
# ββ Overall Summary ββββββββββββββββββββββββββββββββββββββββββ
|
| 451 |
+
print(f"\n{'='*70}")
|
| 452 |
+
print(f" OVERALL QUALITY REPORT")
|
| 453 |
+
print(f"{'='*70}")
|
| 454 |
+
|
| 455 |
+
all_clean = True
|
| 456 |
+
for r in results:
|
| 457 |
+
status = "PASS" if not r["issues"] else "ISSUES FOUND"
|
| 458 |
+
if r["issues"]:
|
| 459 |
+
all_clean = False
|
| 460 |
+
print(f"\n {r['name']}")
|
| 461 |
+
print(f" Tokens: {r['total_tokens']:>15,}")
|
| 462 |
+
print(f" Status: {status}")
|
| 463 |
+
print(f" English: ascii={r['avg_ascii']:.3f} alpha={r['avg_alpha']:.3f}")
|
| 464 |
+
print(f" Cleanliness: html={r['html_count']} boilerplate={r['high_boilerplate']}")
|
| 465 |
+
print(f" Dedup: near={r['near_dupes']} exact={r['exact_dupes']}")
|
| 466 |
+
print(f" Diversity: uwr={r['avg_unique_word_ratio']:.3f} topics={r['topic_count']}")
|
| 467 |
+
print(f" Flagged: {r['flagged_count']}/{r['samples_checked']} ({r['flagged_count']/r['samples_checked']*100:.1f}%)")
|
| 468 |
+
|
| 469 |
+
total_tokens = sum(r["total_tokens"] for r in results)
|
| 470 |
+
print(f"\n Combined pretrain tokens: {total_tokens:,} ({total_tokens/1e9:.3f}B)")
|
| 471 |
+
|
| 472 |
+
if all_clean:
|
| 473 |
+
print(f"\n READY FOR TRAINING FROM SCRATCH!")
|
| 474 |
+
else:
|
| 475 |
+
print(f"\n REVIEW ISSUES ABOVE before retraining.")
|
Base/scripts/build_100m_english.py
ADDED
|
@@ -0,0 +1,667 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""
|
| 3 |
+
Build a NEW 100M-token English dataset, DIFFERENT from existing litdata_english.
|
| 4 |
+
|
| 5 |
+
Downloads fresh Wikipedia + FineWeb-Edu articles that were NOT used in the
|
| 6 |
+
original ~54M token corpus, then tokenizes into litdata format and runs the
|
| 7 |
+
full cleaning pipeline (reclean + smart cleanup).
|
| 8 |
+
|
| 9 |
+
Skips:
|
| 10 |
+
- First ~25,000 qualifying Wikipedia articles (already used ~19,523 + margin)
|
| 11 |
+
- First ~25,000 qualifying FineWeb-Edu docs (already used ~19,445 + margin)
|
| 12 |
+
|
| 13 |
+
Pipeline:
|
| 14 |
+
1. Stream & filter from HuggingFace (Wikipedia + FineWeb-Edu)
|
| 15 |
+
2. Save as parquet files
|
| 16 |
+
3. Tokenize into litdata binary chunks
|
| 17 |
+
4. Reclean every document (text cleaning pipeline)
|
| 18 |
+
5. Smart cleanup (remove bad docs, strip boilerplate)
|
| 19 |
+
6. Save final clean litdata
|
| 20 |
+
|
| 21 |
+
Output: Base/data/litdata_english_100m/ (clean, ready for training)
|
| 22 |
+
"""
|
| 23 |
+
|
| 24 |
+
import json
|
| 25 |
+
import os
|
| 26 |
+
import re
|
| 27 |
+
import time
|
| 28 |
+
import unicodedata
|
| 29 |
+
from pathlib import Path
|
| 30 |
+
from collections import Counter
|
| 31 |
+
|
| 32 |
+
import numpy as np
|
| 33 |
+
import pyarrow as pa
|
| 34 |
+
import pyarrow.parquet as pq
|
| 35 |
+
from tokenizers import Tokenizer
|
| 36 |
+
|
| 37 |
+
ROOT = Path(__file__).resolve().parent.parent.parent
|
| 38 |
+
BLOCK_SIZE = 1025
|
| 39 |
+
DTYPE = np.int32
|
| 40 |
+
CHUNK_BYTES_TARGET = 64 * 1024 * 1024
|
| 41 |
+
EOS_TOKEN_ID = 0
|
| 42 |
+
TARGET_TOKENS = 100_000_000
|
| 43 |
+
TOKENS_PER_WORD = 1.3
|
| 44 |
+
|
| 45 |
+
# How many articles to SKIP from each source (already used in litdata_english)
|
| 46 |
+
WIKI_SKIP = 25_000 # existing used ~19,523 qualifying + safety margin
|
| 47 |
+
FINEWEB_SKIP = 25_000 # existing used ~19,445 qualifying + safety margin
|
| 48 |
+
|
| 49 |
+
# Shares
|
| 50 |
+
WIKI_SHARE = 0.55 # 55% Wikipedia, 45% FineWeb
|
| 51 |
+
FINEWEB_MIN_SCORE = 4.0
|
| 52 |
+
|
| 53 |
+
# Paths
|
| 54 |
+
DATA_DIR = ROOT / "Base" / "data"
|
| 55 |
+
OUTPUT_PARQUET_DIR = DATA_DIR / "filtered_english_100m"
|
| 56 |
+
OUTPUT_LITDATA_DIR = DATA_DIR / "litdata_english_100m"
|
| 57 |
+
TOKENIZER_PATH = ROOT / "Base" / "checkpoints" / "EleutherAI" / "pythia-160m" / "tokenizer.json"
|
| 58 |
+
|
| 59 |
+
print("Loading tokenizer...")
|
| 60 |
+
tokenizer = Tokenizer.from_file(str(TOKENIZER_PATH))
|
| 61 |
+
|
| 62 |
+
|
| 63 |
+
# ==============================================================================
|
| 64 |
+
# TEXT QUALITY FILTERS (from build_english_corpus.py)
|
| 65 |
+
# ==============================================================================
|
| 66 |
+
|
| 67 |
+
def clean_text(text):
|
| 68 |
+
text = unicodedata.normalize("NFKC", text)
|
| 69 |
+
text = re.sub(r'[\x00-\x08\x0b\x0c\x0e-\x1f\x7f-\x9f]', '', text)
|
| 70 |
+
text = re.sub(r'\n{3,}', '\n\n', text)
|
| 71 |
+
text = re.sub(r'[ \t]+', ' ', text)
|
| 72 |
+
text = '\n'.join(line.strip() for line in text.split('\n'))
|
| 73 |
+
return text.strip()
|
| 74 |
+
|
| 75 |
+
|
| 76 |
+
def is_high_quality(text, min_chars=500, min_words=80):
|
| 77 |
+
if len(text) < min_chars:
|
| 78 |
+
return False
|
| 79 |
+
words = text.split()
|
| 80 |
+
num_words = len(words)
|
| 81 |
+
if num_words < min_words:
|
| 82 |
+
return False
|
| 83 |
+
alpha = sum(c.isalpha() for c in text)
|
| 84 |
+
if alpha / max(len(text), 1) < 0.65:
|
| 85 |
+
return False
|
| 86 |
+
avg_word_len = sum(len(w) for w in words) / num_words
|
| 87 |
+
if avg_word_len < 2.5 or avg_word_len > 15:
|
| 88 |
+
return False
|
| 89 |
+
url_hits = text.count('http://') + text.count('https://')
|
| 90 |
+
if url_hits > num_words * 0.03:
|
| 91 |
+
return False
|
| 92 |
+
sentences = re.split(r'[.!?]+', text)
|
| 93 |
+
real_sentences = [s.strip() for s in sentences if len(s.strip()) > 10]
|
| 94 |
+
if len(real_sentences) < 3:
|
| 95 |
+
return False
|
| 96 |
+
lines = [ln.strip() for ln in text.split('\n') if ln.strip()]
|
| 97 |
+
if len(lines) > 5:
|
| 98 |
+
unique_ratio = len(set(lines)) / len(lines)
|
| 99 |
+
if unique_ratio < 0.5:
|
| 100 |
+
return False
|
| 101 |
+
return True
|
| 102 |
+
|
| 103 |
+
|
| 104 |
+
_WIKI_SKIP_PATTERNS = re.compile(
|
| 105 |
+
r'(disambiguation|list of|lists of|index of|outline of|'
|
| 106 |
+
r'wikipedia:|template:|category:|portal:|module:|mediawiki:)',
|
| 107 |
+
re.IGNORECASE
|
| 108 |
+
)
|
| 109 |
+
|
| 110 |
+
|
| 111 |
+
# ==============================================================================
|
| 112 |
+
# DEEP CLEANING PIPELINE (from reclean_english.py)
|
| 113 |
+
# ==============================================================================
|
| 114 |
+
|
| 115 |
+
CONTROL_CHARS = [
|
| 116 |
+
"\x00", "\x01", "\x02", "\x03", "\x04", "\x05", "\x06", "\x07",
|
| 117 |
+
"\x08", "\x0b", "\x0c", "\x0e", "\x0f", "\x10", "\x11", "\x12",
|
| 118 |
+
"\x13", "\x14", "\x15", "\x16", "\x17", "\x18", "\x19", "\x1a",
|
| 119 |
+
"\x1b", "\x1c", "\x1d", "\x1e", "\x1f", "\x7f", "\ufeff", "\ufffd",
|
| 120 |
+
]
|
| 121 |
+
|
| 122 |
+
HTML_ENTITIES = [
|
| 123 |
+
("&", "&"), ("<", "<"), (">", ">"),
|
| 124 |
+
(""", '"'), ("'", "'"), ("'", "'"),
|
| 125 |
+
(" ", " "), ("—", " - "), ("–", "-"),
|
| 126 |
+
("…", "..."), ("«", '"'), ("»", '"'),
|
| 127 |
+
("•", "- "), ("·", " "), ("©", "(c)"),
|
| 128 |
+
("®", "(R)"), ("™", "(TM)"), ("°", " degrees"),
|
| 129 |
+
]
|
| 130 |
+
|
| 131 |
+
RE_URL = re.compile(r'https?://\S+|www\.\S+', re.I)
|
| 132 |
+
RE_EMAIL = re.compile(r'\b[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}\b')
|
| 133 |
+
RE_FILE_PATH = re.compile(r'(?:[A-Z]:\\|/(?:home|usr|var|etc|opt)/)\S+')
|
| 134 |
+
RE_HTML_TAG = re.compile(r'</?[a-zA-Z][a-zA-Z0-9]*(?:\s[^>]*)?\s*/?>')
|
| 135 |
+
RE_HTML_COMMENT = re.compile(r'<!--.*?-->', re.DOTALL)
|
| 136 |
+
RE_CODE_BLOCK = re.compile(r'```[\s\S]*?```')
|
| 137 |
+
RE_IMPORT = re.compile(r'^(?:import |from \S+ import |#include |using namespace |require\()', re.M)
|
| 138 |
+
RE_REPEATED_LINE = re.compile(r'^(.{20,})\n(?:\1\n?)+', re.M)
|
| 139 |
+
RE_REPEATED_PUNCT = re.compile(r'([!?.])\1{3,}')
|
| 140 |
+
RE_REPEATED_CHAR = re.compile(r'(.)\1{5,}')
|
| 141 |
+
RE_REPEATED_WORD = re.compile(r'\b(\w+)(?:\s+\1){2,}\b', re.I)
|
| 142 |
+
RE_MULTI_NEWLINE = re.compile(r'\n{4,}')
|
| 143 |
+
RE_MULTI_SPACE = re.compile(r'[ \t]{2,}')
|
| 144 |
+
RE_TRAILING_SPACE = re.compile(r'[ \t]+$', re.M)
|
| 145 |
+
RE_NO_SPACE_AFTER_PERIOD = re.compile(r'([.!?])([A-Z])')
|
| 146 |
+
RE_DOUBLE_PERIOD = re.compile(r'\.{2}(?!\.)')
|
| 147 |
+
RE_SPACE_BEFORE_PUNCT = re.compile(r'\s+([.,;:!?])')
|
| 148 |
+
|
| 149 |
+
# Smart cleanup patterns
|
| 150 |
+
RE_CJK = re.compile(r'[\u4e00-\u9fff\u3040-\u309f\u30a0-\u30ff\uac00-\ud7af]{3,}')
|
| 151 |
+
RE_ARABIC = re.compile(r'[\u0600-\u06ff]{5,}')
|
| 152 |
+
RE_CYRILLIC = re.compile(r'[\u0400-\u04ff]{5,}')
|
| 153 |
+
RE_DEVANAGARI = re.compile(r'[\u0900-\u097f]{5,}')
|
| 154 |
+
RE_RESIDUAL_CODE = re.compile(r'(function\s*\(|var\s+\w+\s*=|console\.log|document\.get|if\s*\(\s*\w+\s*[!=]==)', re.I)
|
| 155 |
+
|
| 156 |
+
# Boilerplate line patterns
|
| 157 |
+
RE_COOKIE_LINE = re.compile(r'^.*(?:cookie|cookies)\s+(?:policy|consent|notice|preferences|settings).*$', re.I | re.M)
|
| 158 |
+
RE_SUBSCRIBE_LINE = re.compile(r'^.*(?:subscribe|sign\s*up\s+(?:for|to)\s+(?:our|the)\s+newsletter|unsubscribe|opt[\s-]*out\s+of).*$', re.I | re.M)
|
| 159 |
+
RE_CLICKBAIT_LINE = re.compile(r'^.*(?:you\s+won\'?t\s+believe|click\s+here|read\s+more\s*\.{0,3}$|share\s+this\s+(?:article|post|story)|trending\s+now|sponsored\s+content|advertisement).*$', re.I | re.M)
|
| 160 |
+
RE_SOCIAL_LINE = re.compile(r'^.*(?:follow\s+us\s+on|share\s+on\s+(?:facebook|twitter|linkedin|instagram)|like\s+us\s+on|tweet\s+this).*$', re.I | re.M)
|
| 161 |
+
RE_NAV_LINE = re.compile(r'^.*(?:skip\s+to\s+(?:main\s+)?content|back\s+to\s+top|previous\s+article|next\s+article|related\s+(?:articles|posts)).*$', re.I | re.M)
|
| 162 |
+
RE_LOGIN_LINE = re.compile(r'^.*(?:log\s*in\s+to\s+(?:your|an)\s+account|create\s+(?:a\s+)?(?:free\s+)?account|forgot\s+(?:your\s+)?password|already\s+(?:a\s+)?member).*$', re.I | re.M)
|
| 163 |
+
RE_COMMENT_LINE = re.compile(r'^.*(?:leave\s+a\s+(?:comment|reply)|post\s+a\s+comment|\d+\s+comments?$|logged\s+in\s+as).*$', re.I | re.M)
|
| 164 |
+
RE_COPYRIGHT_LINE = re.compile(r'^.*(?:all\s+rights\s+reserved|\(c\)\s*\d{4}|copyright\s+\d{4}).*$', re.I | re.M)
|
| 165 |
+
|
| 166 |
+
|
| 167 |
+
def deep_clean(text):
|
| 168 |
+
"""Full reclean pipeline from reclean_english.py."""
|
| 169 |
+
if not text or len(text.strip()) < 30:
|
| 170 |
+
return None
|
| 171 |
+
|
| 172 |
+
# 1. Unicode normalization
|
| 173 |
+
text = unicodedata.normalize("NFKC", text)
|
| 174 |
+
|
| 175 |
+
# 2. Control chars
|
| 176 |
+
for ch in CONTROL_CHARS:
|
| 177 |
+
text = text.replace(ch, "")
|
| 178 |
+
|
| 179 |
+
# 3. HTML entities
|
| 180 |
+
for old, new in HTML_ENTITIES:
|
| 181 |
+
text = text.replace(old, new)
|
| 182 |
+
|
| 183 |
+
# 4. HTML tags/comments
|
| 184 |
+
text = RE_HTML_COMMENT.sub("", text)
|
| 185 |
+
text = RE_HTML_TAG.sub("", text)
|
| 186 |
+
|
| 187 |
+
# 5. URLs, emails, file paths
|
| 188 |
+
text = RE_URL.sub("", text)
|
| 189 |
+
text = RE_EMAIL.sub("", text)
|
| 190 |
+
text = RE_FILE_PATH.sub("", text)
|
| 191 |
+
|
| 192 |
+
# 6. Code blocks
|
| 193 |
+
text = RE_CODE_BLOCK.sub("", text)
|
| 194 |
+
|
| 195 |
+
# 7. Repeated content
|
| 196 |
+
text = RE_REPEATED_LINE.sub(r'\1', text)
|
| 197 |
+
text = RE_REPEATED_PUNCT.sub(r'\1\1\1', text)
|
| 198 |
+
text = RE_REPEATED_CHAR.sub(r'\1\1\1', text)
|
| 199 |
+
text = RE_REPEATED_WORD.sub(r'\1', text)
|
| 200 |
+
|
| 201 |
+
# 8. Whitespace
|
| 202 |
+
text = text.replace('\t', ' ')
|
| 203 |
+
text = RE_TRAILING_SPACE.sub('', text)
|
| 204 |
+
text = RE_MULTI_SPACE.sub(' ', text)
|
| 205 |
+
text = RE_MULTI_NEWLINE.sub('\n\n\n', text)
|
| 206 |
+
|
| 207 |
+
# 9. Punctuation
|
| 208 |
+
text = RE_DOUBLE_PERIOD.sub('.', text)
|
| 209 |
+
text = RE_NO_SPACE_AFTER_PERIOD.sub(r'\1 \2', text)
|
| 210 |
+
text = RE_SPACE_BEFORE_PUNCT.sub(r'\1', text)
|
| 211 |
+
|
| 212 |
+
# 10. Smart quotes -> ASCII
|
| 213 |
+
text = text.replace('\u2018', "'").replace('\u2019', "'")
|
| 214 |
+
text = text.replace('\u201c', '"').replace('\u201d', '"')
|
| 215 |
+
text = text.replace('\u2013', '-').replace('\u2014', ' - ')
|
| 216 |
+
text = text.replace('\u2026', '...')
|
| 217 |
+
text = text.replace('\u2022', '- ')
|
| 218 |
+
text = text.replace('\u00b7', ' ')
|
| 219 |
+
text = text.replace('\u00a0', ' ')
|
| 220 |
+
|
| 221 |
+
# 11. Line-by-line cleaning
|
| 222 |
+
lines = text.split('\n')
|
| 223 |
+
clean_lines = []
|
| 224 |
+
for line in lines:
|
| 225 |
+
line = line.strip()
|
| 226 |
+
if not line:
|
| 227 |
+
clean_lines.append('')
|
| 228 |
+
continue
|
| 229 |
+
if len(line) > 10:
|
| 230 |
+
alpha_count = sum(1 for c in line if c.isalpha())
|
| 231 |
+
if alpha_count / len(line) < 0.40:
|
| 232 |
+
continue
|
| 233 |
+
if line.count('|') > 3 or line.count('{') > 2 or line.count('}') > 2:
|
| 234 |
+
continue
|
| 235 |
+
if RE_IMPORT.match(line):
|
| 236 |
+
continue
|
| 237 |
+
if line and line[0].isalpha() and line[0].islower():
|
| 238 |
+
if not clean_lines or clean_lines[-1] == '' or clean_lines[-1].rstrip().endswith(('.', '!', '?', ':')):
|
| 239 |
+
line = line[0].upper() + line[1:]
|
| 240 |
+
clean_lines.append(line)
|
| 241 |
+
|
| 242 |
+
text = '\n'.join(clean_lines)
|
| 243 |
+
text = text.strip()
|
| 244 |
+
|
| 245 |
+
# 12. Deduplicate paragraphs
|
| 246 |
+
paragraphs = text.split('\n\n')
|
| 247 |
+
seen = set()
|
| 248 |
+
unique_paragraphs = []
|
| 249 |
+
for p in paragraphs:
|
| 250 |
+
p_stripped = p.strip()
|
| 251 |
+
if not p_stripped:
|
| 252 |
+
continue
|
| 253 |
+
p_key = ' '.join(p_stripped.lower().split())
|
| 254 |
+
if p_key not in seen:
|
| 255 |
+
seen.add(p_key)
|
| 256 |
+
unique_paragraphs.append(p_stripped)
|
| 257 |
+
text = '\n\n'.join(unique_paragraphs)
|
| 258 |
+
|
| 259 |
+
# 13. Final quality gate
|
| 260 |
+
text = text.strip()
|
| 261 |
+
if len(text) < 50:
|
| 262 |
+
return None
|
| 263 |
+
if len(text.split()) < 10:
|
| 264 |
+
return None
|
| 265 |
+
ascii_count = sum(1 for c in text if ord(c) < 128)
|
| 266 |
+
if ascii_count / max(len(text), 1) < 0.85:
|
| 267 |
+
return None
|
| 268 |
+
|
| 269 |
+
return text
|
| 270 |
+
|
| 271 |
+
|
| 272 |
+
def smart_filter(text):
|
| 273 |
+
"""Smart cleanup: returns (keep, cleaned_text, reason)."""
|
| 274 |
+
words = text.split()
|
| 275 |
+
word_count = len(words)
|
| 276 |
+
|
| 277 |
+
if word_count < 50:
|
| 278 |
+
return False, text, f"too short ({word_count} words)"
|
| 279 |
+
|
| 280 |
+
# Non-English scripts
|
| 281 |
+
scripts = []
|
| 282 |
+
if RE_CJK.search(text): scripts.append("CJK")
|
| 283 |
+
if RE_ARABIC.search(text): scripts.append("Arabic")
|
| 284 |
+
if RE_CYRILLIC.search(text): scripts.append("Cyrillic")
|
| 285 |
+
if RE_DEVANAGARI.search(text): scripts.append("Devanagari")
|
| 286 |
+
if scripts:
|
| 287 |
+
return False, text, f"non-English: {', '.join(scripts)}"
|
| 288 |
+
|
| 289 |
+
if word_count > 50:
|
| 290 |
+
unique_ratio = len(set(w.lower() for w in words)) / word_count
|
| 291 |
+
if unique_ratio < 0.20:
|
| 292 |
+
return False, text, f"repetitive ({unique_ratio:.3f})"
|
| 293 |
+
|
| 294 |
+
code_matches = RE_RESIDUAL_CODE.findall(text)
|
| 295 |
+
if len(code_matches) >= 5:
|
| 296 |
+
return False, text, f"residual code ({len(code_matches)})"
|
| 297 |
+
|
| 298 |
+
# Strip boilerplate lines
|
| 299 |
+
original_len = len(text)
|
| 300 |
+
for pattern in [RE_COOKIE_LINE, RE_SUBSCRIBE_LINE, RE_CLICKBAIT_LINE,
|
| 301 |
+
RE_SOCIAL_LINE, RE_NAV_LINE, RE_LOGIN_LINE,
|
| 302 |
+
RE_COMMENT_LINE, RE_COPYRIGHT_LINE]:
|
| 303 |
+
text = pattern.sub('', text)
|
| 304 |
+
|
| 305 |
+
# Strip number-heavy lines
|
| 306 |
+
lines = text.split('\n')
|
| 307 |
+
clean_lines = []
|
| 308 |
+
for line in lines:
|
| 309 |
+
stripped = line.strip()
|
| 310 |
+
if stripped and len(stripped) > 10:
|
| 311 |
+
digit_count = sum(1 for c in stripped if c.isdigit() or c in ' ,.\t-+/%$')
|
| 312 |
+
if digit_count / len(stripped) > 0.80:
|
| 313 |
+
continue
|
| 314 |
+
clean_lines.append(line)
|
| 315 |
+
text = '\n'.join(clean_lines)
|
| 316 |
+
text = re.sub(r'\n{3,}', '\n\n', text)
|
| 317 |
+
text = text.strip()
|
| 318 |
+
|
| 319 |
+
if len(text.split()) < 50:
|
| 320 |
+
return False, text, "too short after stripping"
|
| 321 |
+
|
| 322 |
+
return True, text, None
|
| 323 |
+
|
| 324 |
+
|
| 325 |
+
# ==============================================================================
|
| 326 |
+
# LITDATA I/O
|
| 327 |
+
# ==============================================================================
|
| 328 |
+
|
| 329 |
+
def write_litdata_chunks(output_dir, token_stream, config):
|
| 330 |
+
os.makedirs(output_dir, exist_ok=True)
|
| 331 |
+
dtype_size = DTYPE().itemsize
|
| 332 |
+
tokens_per_chunk = CHUNK_BYTES_TARGET // dtype_size
|
| 333 |
+
tokens_per_chunk = (tokens_per_chunk // BLOCK_SIZE) * BLOCK_SIZE
|
| 334 |
+
chunks_metadata = []
|
| 335 |
+
pos = 0
|
| 336 |
+
chunk_idx = 0
|
| 337 |
+
while pos < len(token_stream):
|
| 338 |
+
remaining = len(token_stream) - pos
|
| 339 |
+
chunk_tokens = min(tokens_per_chunk, remaining)
|
| 340 |
+
num_blocks = chunk_tokens // BLOCK_SIZE
|
| 341 |
+
if num_blocks == 0:
|
| 342 |
+
break
|
| 343 |
+
actual_tokens = num_blocks * BLOCK_SIZE
|
| 344 |
+
chunk_data = token_stream[pos:pos + actual_tokens]
|
| 345 |
+
filename = f"chunk-0-{chunk_idx}.bin"
|
| 346 |
+
filepath = os.path.join(output_dir, filename)
|
| 347 |
+
header_num_items = np.array([num_blocks], dtype=np.uint32)
|
| 348 |
+
offsets = np.arange(num_blocks + 1, dtype=np.uint32) * (BLOCK_SIZE * dtype_size)
|
| 349 |
+
header = np.concatenate([header_num_items, offsets])
|
| 350 |
+
with open(filepath, "wb") as f:
|
| 351 |
+
header.tofile(f)
|
| 352 |
+
chunk_data.tofile(f)
|
| 353 |
+
meta = {
|
| 354 |
+
"chunk_bytes": int(header.nbytes + chunk_data.nbytes),
|
| 355 |
+
"chunk_size": num_blocks,
|
| 356 |
+
"dim": int(actual_tokens),
|
| 357 |
+
"filename": filename,
|
| 358 |
+
}
|
| 359 |
+
chunks_metadata.append(meta)
|
| 360 |
+
pos += actual_tokens
|
| 361 |
+
chunk_idx += 1
|
| 362 |
+
print(f" Written chunk {chunk_idx} ({pos:,}/{len(token_stream):,} tokens)")
|
| 363 |
+
index = {"chunks": chunks_metadata, "config": config, "updated_at": str(time.time())}
|
| 364 |
+
with open(os.path.join(output_dir, "index.json"), "w") as f:
|
| 365 |
+
json.dump(index, f, indent=2)
|
| 366 |
+
return chunks_metadata
|
| 367 |
+
|
| 368 |
+
|
| 369 |
+
# ==============================================================================
|
| 370 |
+
# MAIN PIPELINE
|
| 371 |
+
# ==============================================================================
|
| 372 |
+
|
| 373 |
+
def main():
|
| 374 |
+
from datasets import load_dataset
|
| 375 |
+
|
| 376 |
+
t_start = time.time()
|
| 377 |
+
|
| 378 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 379 |
+
# PHASE 1: Download fresh data (skipping already-used articles)
|
| 380 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 381 |
+
wiki_target = int(TARGET_TOKENS * WIKI_SHARE)
|
| 382 |
+
fineweb_target = TARGET_TOKENS - wiki_target
|
| 383 |
+
|
| 384 |
+
print(f"\n{'='*75}")
|
| 385 |
+
print(f" PHASE 1: DOWNLOADING FRESH DATA")
|
| 386 |
+
print(f" Target: {TARGET_TOKENS:,} tokens ({wiki_target:,} Wiki + {fineweb_target:,} FineWeb)")
|
| 387 |
+
print(f" Skipping first {WIKI_SKIP:,} Wiki + {FINEWEB_SKIP:,} FineWeb (already used)")
|
| 388 |
+
print(f"{'='*75}")
|
| 389 |
+
|
| 390 |
+
os.makedirs(str(OUTPUT_PARQUET_DIR), exist_ok=True)
|
| 391 |
+
all_texts = []
|
| 392 |
+
|
| 393 |
+
# --- Wikipedia ---
|
| 394 |
+
print(f"\n [Wikipedia] Streaming English articles (skipping first {WIKI_SKIP:,})...")
|
| 395 |
+
ds_wiki = load_dataset(
|
| 396 |
+
"wikimedia/wikipedia", "20231101.en",
|
| 397 |
+
split="train", streaming=True,
|
| 398 |
+
trust_remote_code=False
|
| 399 |
+
)
|
| 400 |
+
|
| 401 |
+
wiki_texts = []
|
| 402 |
+
wiki_tokens = 0
|
| 403 |
+
wiki_seen = 0
|
| 404 |
+
wiki_skipped_quality = 0
|
| 405 |
+
wiki_skipped_meta = 0
|
| 406 |
+
t0 = time.time()
|
| 407 |
+
|
| 408 |
+
for article in ds_wiki:
|
| 409 |
+
title = (article.get("title") or "").strip()
|
| 410 |
+
raw = article.get("text") or ""
|
| 411 |
+
|
| 412 |
+
if _WIKI_SKIP_PATTERNS.search(title):
|
| 413 |
+
wiki_skipped_meta += 1
|
| 414 |
+
continue
|
| 415 |
+
|
| 416 |
+
cleaned = clean_text(raw)
|
| 417 |
+
if not is_high_quality(cleaned, min_chars=800, min_words=120):
|
| 418 |
+
wiki_skipped_quality += 1
|
| 419 |
+
continue
|
| 420 |
+
|
| 421 |
+
# This is a qualifying article - count it
|
| 422 |
+
wiki_seen += 1
|
| 423 |
+
|
| 424 |
+
# Skip if already used
|
| 425 |
+
if wiki_seen <= WIKI_SKIP:
|
| 426 |
+
if wiki_seen % 5000 == 0:
|
| 427 |
+
print(f" Skipping... {wiki_seen:,}/{WIKI_SKIP:,}")
|
| 428 |
+
continue
|
| 429 |
+
|
| 430 |
+
full_text = f"{title}\n\n{cleaned}"
|
| 431 |
+
est_tok = int(len(full_text.split()) * TOKENS_PER_WORD)
|
| 432 |
+
wiki_texts.append(full_text)
|
| 433 |
+
wiki_tokens += est_tok
|
| 434 |
+
|
| 435 |
+
if len(wiki_texts) % 5000 == 0:
|
| 436 |
+
elapsed = time.time() - t0
|
| 437 |
+
print(f" Collected {len(wiki_texts):,} articles | ~{wiki_tokens:,} tokens | {elapsed:.0f}s")
|
| 438 |
+
|
| 439 |
+
if wiki_tokens >= wiki_target:
|
| 440 |
+
break
|
| 441 |
+
|
| 442 |
+
elapsed = time.time() - t0
|
| 443 |
+
print(f" [Wikipedia] Done: {len(wiki_texts):,} articles, ~{wiki_tokens:,} tokens in {elapsed:.0f}s")
|
| 444 |
+
print(f" (skipped {WIKI_SKIP:,} already-used + {wiki_skipped_quality:,} low-quality + {wiki_skipped_meta:,} meta)")
|
| 445 |
+
|
| 446 |
+
# Save wiki parquets
|
| 447 |
+
BATCH = 5000
|
| 448 |
+
for i in range(0, len(wiki_texts), BATCH):
|
| 449 |
+
batch = wiki_texts[i:i+BATCH]
|
| 450 |
+
fp = OUTPUT_PARQUET_DIR / f"wiki_{i//BATCH:04d}.parquet"
|
| 451 |
+
pq.write_table(pa.table({"text": batch}), str(fp))
|
| 452 |
+
print(f" Saved {(len(wiki_texts)-1)//BATCH + 1} wiki parquet files")
|
| 453 |
+
|
| 454 |
+
# --- FineWeb-Edu ---
|
| 455 |
+
print(f"\n [FineWeb-Edu] Streaming (score >= {FINEWEB_MIN_SCORE}, skipping first {FINEWEB_SKIP:,})...")
|
| 456 |
+
ds_fineweb = load_dataset(
|
| 457 |
+
"HuggingFaceFW/fineweb-edu", "sample-10BT",
|
| 458 |
+
split="train", streaming=True,
|
| 459 |
+
trust_remote_code=False
|
| 460 |
+
)
|
| 461 |
+
|
| 462 |
+
fineweb_texts = []
|
| 463 |
+
fineweb_tokens = 0
|
| 464 |
+
fineweb_seen = 0
|
| 465 |
+
fineweb_skipped_quality = 0
|
| 466 |
+
fineweb_skipped_score = 0
|
| 467 |
+
t0 = time.time()
|
| 468 |
+
|
| 469 |
+
for doc in ds_fineweb:
|
| 470 |
+
score = doc.get("score", 0)
|
| 471 |
+
if not isinstance(score, (int, float)):
|
| 472 |
+
try:
|
| 473 |
+
score = float(score)
|
| 474 |
+
except (ValueError, TypeError):
|
| 475 |
+
continue
|
| 476 |
+
if score < FINEWEB_MIN_SCORE:
|
| 477 |
+
fineweb_skipped_score += 1
|
| 478 |
+
continue
|
| 479 |
+
|
| 480 |
+
raw = doc.get("text") or ""
|
| 481 |
+
cleaned = clean_text(raw)
|
| 482 |
+
if not is_high_quality(cleaned, min_chars=500, min_words=80):
|
| 483 |
+
fineweb_skipped_quality += 1
|
| 484 |
+
continue
|
| 485 |
+
|
| 486 |
+
# This is a qualifying doc - count it
|
| 487 |
+
fineweb_seen += 1
|
| 488 |
+
|
| 489 |
+
# Skip if already used
|
| 490 |
+
if fineweb_seen <= FINEWEB_SKIP:
|
| 491 |
+
if fineweb_seen % 5000 == 0:
|
| 492 |
+
print(f" Skipping... {fineweb_seen:,}/{FINEWEB_SKIP:,}")
|
| 493 |
+
continue
|
| 494 |
+
|
| 495 |
+
est_tok = int(len(cleaned.split()) * TOKENS_PER_WORD)
|
| 496 |
+
fineweb_texts.append(cleaned)
|
| 497 |
+
fineweb_tokens += est_tok
|
| 498 |
+
|
| 499 |
+
if len(fineweb_texts) % 5000 == 0:
|
| 500 |
+
elapsed = time.time() - t0
|
| 501 |
+
print(f" Collected {len(fineweb_texts):,} docs | ~{fineweb_tokens:,} tokens | {elapsed:.0f}s")
|
| 502 |
+
|
| 503 |
+
if fineweb_tokens >= fineweb_target:
|
| 504 |
+
break
|
| 505 |
+
|
| 506 |
+
elapsed = time.time() - t0
|
| 507 |
+
print(f" [FineWeb-Edu] Done: {len(fineweb_texts):,} docs, ~{fineweb_tokens:,} tokens in {elapsed:.0f}s")
|
| 508 |
+
print(f" (skipped {FINEWEB_SKIP:,} already-used + {fineweb_skipped_quality:,} low-quality + {fineweb_skipped_score:,} low-score)")
|
| 509 |
+
|
| 510 |
+
# Save fineweb parquets
|
| 511 |
+
for i in range(0, len(fineweb_texts), BATCH):
|
| 512 |
+
batch = fineweb_texts[i:i+BATCH]
|
| 513 |
+
fp = OUTPUT_PARQUET_DIR / f"fineweb_{i//BATCH:04d}.parquet"
|
| 514 |
+
pq.write_table(pa.table({"text": batch}), str(fp))
|
| 515 |
+
print(f" Saved {(len(fineweb_texts)-1)//BATCH + 1} fineweb parquet files")
|
| 516 |
+
|
| 517 |
+
# Combine all texts
|
| 518 |
+
all_texts = wiki_texts + fineweb_texts
|
| 519 |
+
total_est_tokens = wiki_tokens + fineweb_tokens
|
| 520 |
+
del wiki_texts, fineweb_texts
|
| 521 |
+
|
| 522 |
+
print(f"\n PHASE 1 COMPLETE: {len(all_texts):,} documents, ~{total_est_tokens:,} estimated tokens")
|
| 523 |
+
|
| 524 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 525 |
+
# PHASE 2: Deep clean every document
|
| 526 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 527 |
+
print(f"\n{'='*75}")
|
| 528 |
+
print(f" PHASE 2: DEEP CLEANING {len(all_texts):,} DOCUMENTS")
|
| 529 |
+
print(f"{'='*75}")
|
| 530 |
+
|
| 531 |
+
t2 = time.time()
|
| 532 |
+
cleaned_texts = []
|
| 533 |
+
dropped_clean = 0
|
| 534 |
+
for i, text in enumerate(all_texts):
|
| 535 |
+
result = deep_clean(text)
|
| 536 |
+
if result is not None:
|
| 537 |
+
cleaned_texts.append(result)
|
| 538 |
+
else:
|
| 539 |
+
dropped_clean += 1
|
| 540 |
+
if (i + 1) % 10000 == 0 or i == len(all_texts) - 1:
|
| 541 |
+
print(f" Cleaned {i+1:,}/{len(all_texts):,} | kept={len(cleaned_texts):,} | dropped={dropped_clean:,}")
|
| 542 |
+
|
| 543 |
+
del all_texts
|
| 544 |
+
print(f" Deep clean done in {time.time()-t2:.1f}s")
|
| 545 |
+
print(f" Kept: {len(cleaned_texts):,} | Dropped: {dropped_clean:,}")
|
| 546 |
+
|
| 547 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 548 |
+
# PHASE 3: Smart cleanup (filter + boilerplate strip)
|
| 549 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 550 |
+
print(f"\n{'='*75}")
|
| 551 |
+
print(f" PHASE 3: SMART CLEANUP ON {len(cleaned_texts):,} DOCUMENTS")
|
| 552 |
+
print(f"{'='*75}")
|
| 553 |
+
|
| 554 |
+
t3 = time.time()
|
| 555 |
+
final_texts = []
|
| 556 |
+
removed_reasons = Counter()
|
| 557 |
+
total_boilerplate = 0
|
| 558 |
+
|
| 559 |
+
for i, text in enumerate(cleaned_texts):
|
| 560 |
+
keep, stripped_text, reason = smart_filter(text)
|
| 561 |
+
if keep:
|
| 562 |
+
final_texts.append(stripped_text)
|
| 563 |
+
total_boilerplate += len(text) - len(stripped_text)
|
| 564 |
+
else:
|
| 565 |
+
removed_reasons[reason.split('(')[0].strip().split(':')[0].strip()] += 1
|
| 566 |
+
|
| 567 |
+
if (i + 1) % 10000 == 0 or i == len(cleaned_texts) - 1:
|
| 568 |
+
print(f" Processed {i+1:,}/{len(cleaned_texts):,} | kept={len(final_texts):,}")
|
| 569 |
+
|
| 570 |
+
del cleaned_texts
|
| 571 |
+
total_removed = len(final_texts)
|
| 572 |
+
print(f" Smart cleanup done in {time.time()-t3:.1f}s")
|
| 573 |
+
print(f" Final documents: {len(final_texts):,}")
|
| 574 |
+
print(f" Boilerplate stripped: {total_boilerplate:,} chars")
|
| 575 |
+
|
| 576 |
+
if removed_reasons:
|
| 577 |
+
print(f" Removal reasons:")
|
| 578 |
+
for reason, count in sorted(removed_reasons.items(), key=lambda x: -x[1]):
|
| 579 |
+
print(f" {reason:<35} {count:>6,}")
|
| 580 |
+
|
| 581 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 582 |
+
# PHASE 4: Tokenize and write litdata chunks
|
| 583 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 584 |
+
print(f"\n{'='*75}")
|
| 585 |
+
print(f" PHASE 4: TOKENIZING {len(final_texts):,} DOCUMENTS")
|
| 586 |
+
print(f"{'='*75}")
|
| 587 |
+
|
| 588 |
+
t4 = time.time()
|
| 589 |
+
all_token_ids = []
|
| 590 |
+
total_tokens = 0
|
| 591 |
+
ENCODE_BATCH = 10000
|
| 592 |
+
|
| 593 |
+
for i in range(0, len(final_texts), ENCODE_BATCH):
|
| 594 |
+
batch = final_texts[i:i+ENCODE_BATCH]
|
| 595 |
+
encoded = tokenizer.encode_batch(batch, add_special_tokens=False)
|
| 596 |
+
for enc in encoded:
|
| 597 |
+
ids = enc.ids
|
| 598 |
+
all_token_ids.extend(ids)
|
| 599 |
+
all_token_ids.append(EOS_TOKEN_ID)
|
| 600 |
+
total_tokens += len(ids) + 1
|
| 601 |
+
done = min(i + ENCODE_BATCH, len(final_texts))
|
| 602 |
+
if done % 20000 == 0 or done == len(final_texts):
|
| 603 |
+
print(f" Tokenized {done:,}/{len(final_texts):,} ({total_tokens:,} tokens)")
|
| 604 |
+
|
| 605 |
+
del final_texts
|
| 606 |
+
print(f" Tokenized in {time.time()-t4:.1f}s β {total_tokens:,} total tokens")
|
| 607 |
+
|
| 608 |
+
# Write litdata chunks
|
| 609 |
+
print(f"\n Writing litdata chunks to {OUTPUT_LITDATA_DIR}...")
|
| 610 |
+
token_array = np.array(all_token_ids, dtype=DTYPE)
|
| 611 |
+
del all_token_ids
|
| 612 |
+
|
| 613 |
+
config = {
|
| 614 |
+
"block_size": BLOCK_SIZE,
|
| 615 |
+
"vocab_size": tokenizer.get_vocab_size(),
|
| 616 |
+
}
|
| 617 |
+
chunks = write_litdata_chunks(str(OUTPUT_LITDATA_DIR), token_array, config)
|
| 618 |
+
final_token_count = sum(c["dim"] for c in chunks)
|
| 619 |
+
del token_array
|
| 620 |
+
|
| 621 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 622 |
+
# FINAL REPORT
|
| 623 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 624 |
+
total_time = time.time() - t_start
|
| 625 |
+
|
| 626 |
+
report = []
|
| 627 |
+
report.append(f"\n{'='*75}")
|
| 628 |
+
report.append(f" LITDATA_ENGLISH_100M - BUILD & CLEAN REPORT")
|
| 629 |
+
report.append(f"{'='*75}")
|
| 630 |
+
report.append(f"\n Total time: {total_time:.0f}s ({total_time/60:.1f} min)")
|
| 631 |
+
report.append(f"\n SOURCES")
|
| 632 |
+
report.append(f" {'-'*60}")
|
| 633 |
+
report.append(f" Wikipedia (articles {WIKI_SKIP+1:,}+): {wiki_tokens:,} est. tokens")
|
| 634 |
+
report.append(f" FineWeb-Edu (docs {FINEWEB_SKIP+1:,}+): {fineweb_tokens:,} est. tokens")
|
| 635 |
+
report.append(f"\n PIPELINE RESULTS")
|
| 636 |
+
report.append(f" {'-'*60}")
|
| 637 |
+
report.append(f" Downloaded: ~{total_est_tokens:,} est. tokens")
|
| 638 |
+
report.append(f" After deep clean: dropped {dropped_clean:,} docs")
|
| 639 |
+
report.append(f" After smart filter: removed {sum(removed_reasons.values()):,} docs")
|
| 640 |
+
if removed_reasons:
|
| 641 |
+
for reason, count in sorted(removed_reasons.items(), key=lambda x: -x[1]):
|
| 642 |
+
report.append(f" - {reason}: {count:,}")
|
| 643 |
+
report.append(f" Boilerplate stripped: {total_boilerplate:,} chars")
|
| 644 |
+
report.append(f"\n FINAL OUTPUT")
|
| 645 |
+
report.append(f" {'-'*60}")
|
| 646 |
+
report.append(f" Location: {OUTPUT_LITDATA_DIR}")
|
| 647 |
+
report.append(f" Chunks: {len(chunks)}")
|
| 648 |
+
report.append(f" Tokens: {final_token_count:,}")
|
| 649 |
+
report.append(f" Format: litdata binary (int32, BLOCK_SIZE=1025, EOS=0)")
|
| 650 |
+
report.append(f"\n DATA IS DIFFERENT FROM litdata_english_clean:")
|
| 651 |
+
report.append(f" - Skipped first {WIKI_SKIP:,} qualifying Wikipedia articles")
|
| 652 |
+
report.append(f" - Skipped first {FINEWEB_SKIP:,} qualifying FineWeb-Edu docs")
|
| 653 |
+
report.append(f" - Zero overlap with existing ~54M token corpus")
|
| 654 |
+
report.append(f"\n{'='*75}")
|
| 655 |
+
|
| 656 |
+
full_report = '\n'.join(report)
|
| 657 |
+
print(full_report)
|
| 658 |
+
|
| 659 |
+
report_path = OUTPUT_LITDATA_DIR / "BUILD_REPORT.txt"
|
| 660 |
+
with open(report_path, "w", encoding="utf-8") as f:
|
| 661 |
+
f.write(full_report)
|
| 662 |
+
print(f"\n Report saved to: {report_path}")
|
| 663 |
+
print(f" Done! Ready for training.")
|
| 664 |
+
|
| 665 |
+
|
| 666 |
+
if __name__ == "__main__":
|
| 667 |
+
main()
|
Base/scripts/build_5b_pretrain.py
ADDED
|
@@ -0,0 +1,761 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""
|
| 3 |
+
Build 1.2B tokens of DIVERSE, CLEAN English data to bring pretrain total to 5B.
|
| 4 |
+
|
| 5 |
+
Current: 3,819,212,525 tokens in litdata_pretrain_final
|
| 6 |
+
Target: 5,000,000,000 tokens
|
| 7 |
+
Gap: ~1,181,000,000 tokens
|
| 8 |
+
Build: ~1,250,000,000 est. tokens (buffer for cleaning loss)
|
| 9 |
+
|
| 10 |
+
DIVERSITY STRATEGY β 3 completely different source types:
|
| 11 |
+
1. Wikipedia (skip 200K qualifying articles) β encyclopedic knowledge
|
| 12 |
+
2. FineWeb-Edu (skip 200K qualifying docs, score β₯ 3.5) β educational web
|
| 13 |
+
3. OpenWebText (Skylion007) β Reddit-curated quality English web pages
|
| 14 |
+
* Completely new source, zero overlap with anything used before
|
| 15 |
+
|
| 16 |
+
Pipeline per source:
|
| 17 |
+
Download (streaming) β Deep clean β Smart filter β Tokenize β Stream-write
|
| 18 |
+
|
| 19 |
+
Memory-efficient: StreamingChunkWriter, multiprocessing cleaning (30 cores).
|
| 20 |
+
Output appended directly to litdata_pretrain_final.
|
| 21 |
+
"""
|
| 22 |
+
|
| 23 |
+
import json
|
| 24 |
+
import os
|
| 25 |
+
import re
|
| 26 |
+
import time
|
| 27 |
+
import unicodedata
|
| 28 |
+
from pathlib import Path
|
| 29 |
+
from multiprocessing import Pool, cpu_count
|
| 30 |
+
from collections import Counter
|
| 31 |
+
|
| 32 |
+
import numpy as np
|
| 33 |
+
from tokenizers import Tokenizer
|
| 34 |
+
|
| 35 |
+
ROOT = Path(__file__).resolve().parent.parent.parent
|
| 36 |
+
BLOCK_SIZE = 1025
|
| 37 |
+
DTYPE = np.int32
|
| 38 |
+
CHUNK_BYTES_TARGET = 64 * 1024 * 1024
|
| 39 |
+
EOS_TOKEN_ID = 0
|
| 40 |
+
NUM_WORKERS = max(1, cpu_count() - 2)
|
| 41 |
+
ENCODE_BATCH = 8000
|
| 42 |
+
TOKENS_PER_WORD = 1.3
|
| 43 |
+
|
| 44 |
+
DATA_DIR = ROOT / "Base" / "data"
|
| 45 |
+
FINAL_DIR = DATA_DIR / "litdata_pretrain_final"
|
| 46 |
+
TOKENIZER_PATH = str(ROOT / "Base" / "checkpoints" / "EleutherAI" / "pythia-160m" / "tokenizer.json")
|
| 47 |
+
|
| 48 |
+
# ββ Source targets ββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 49 |
+
# Slightly over 1.2B to account for cleaning losses (~2-3%)
|
| 50 |
+
WIKI_TARGET = 365_000_000 # ~30% β encyclopedic
|
| 51 |
+
FINEWEB_TARGET = 445_000_000 # ~36% β educational web
|
| 52 |
+
OWT_TARGET = 445_000_000 # ~34% β Reddit-curated diverse web
|
| 53 |
+
TOTAL_TARGET = WIKI_TARGET + FINEWEB_TARGET + OWT_TARGET # ~1.255B
|
| 54 |
+
|
| 55 |
+
# Skip values β must exceed ALL previously used qualifying articles
|
| 56 |
+
WIKI_SKIP = 200_000 # previous max was 125K + articles collected
|
| 57 |
+
FINEWEB_SKIP = 200_000 # previous max was 125K + docs collected
|
| 58 |
+
FINEWEB_MIN_SCORE = 3.5 # slightly broader than 4.0 for more diversity
|
| 59 |
+
|
| 60 |
+
_WIKI_SKIP_PATTERNS = re.compile(
|
| 61 |
+
r'(disambiguation|list of|lists of|index of|outline of|'
|
| 62 |
+
r'wikipedia:|template:|category:|portal:|module:|mediawiki:)',
|
| 63 |
+
re.IGNORECASE
|
| 64 |
+
)
|
| 65 |
+
|
| 66 |
+
|
| 67 |
+
# ==============================================================================
|
| 68 |
+
# CLEANING PIPELINE
|
| 69 |
+
# ==============================================================================
|
| 70 |
+
|
| 71 |
+
CONTROL_CHARS = [
|
| 72 |
+
"\x00", "\x01", "\x02", "\x03", "\x04", "\x05", "\x06", "\x07",
|
| 73 |
+
"\x08", "\x0b", "\x0c", "\x0e", "\x0f", "\x10", "\x11", "\x12",
|
| 74 |
+
"\x13", "\x14", "\x15", "\x16", "\x17", "\x18", "\x19", "\x1a",
|
| 75 |
+
"\x1b", "\x1c", "\x1d", "\x1e", "\x1f", "\x7f", "\ufeff", "\ufffd",
|
| 76 |
+
]
|
| 77 |
+
|
| 78 |
+
HTML_ENTITIES = [
|
| 79 |
+
("&", "&"), ("<", "<"), (">", ">"),
|
| 80 |
+
(""", '"'), ("'", "'"), ("'", "'"),
|
| 81 |
+
(" ", " "), ("—", " - "), ("–", "-"),
|
| 82 |
+
("…", "..."), ("«", '"'), ("»", '"'),
|
| 83 |
+
("•", "- "), ("·", " "), ("©", "(c)"),
|
| 84 |
+
("®", "(R)"), ("™", "(TM)"), ("°", " degrees"),
|
| 85 |
+
]
|
| 86 |
+
|
| 87 |
+
RE_URL = re.compile(r'https?://\S+|www\.\S+', re.I)
|
| 88 |
+
RE_EMAIL = re.compile(r'\b[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}\b')
|
| 89 |
+
RE_FILE_PATH = re.compile(r'(?:[A-Z]:\\|/(?:home|usr|var|etc|opt)/)\S+')
|
| 90 |
+
RE_HTML_TAG = re.compile(r'</?[a-zA-Z][a-zA-Z0-9]*(?:\s[^>]*)?\s*/?>')
|
| 91 |
+
RE_HTML_COMMENT = re.compile(r'<!--.*?-->', re.DOTALL)
|
| 92 |
+
RE_CODE_BLOCK = re.compile(r'```[\s\S]*?```')
|
| 93 |
+
RE_IMPORT = re.compile(r'^(?:import |from \S+ import |#include |using namespace |require\()', re.M)
|
| 94 |
+
RE_REPEATED_LINE = re.compile(r'^(.{20,})\n(?:\1\n?)+', re.M)
|
| 95 |
+
RE_REPEATED_PUNCT = re.compile(r'([!?.])\1{3,}')
|
| 96 |
+
RE_REPEATED_CHAR = re.compile(r'(.)\1{5,}')
|
| 97 |
+
RE_REPEATED_WORD = re.compile(r'\b(\w+)(?:\s+\1){2,}\b', re.I)
|
| 98 |
+
RE_MULTI_NEWLINE = re.compile(r'\n{4,}')
|
| 99 |
+
RE_MULTI_SPACE = re.compile(r'[ \t]{2,}')
|
| 100 |
+
RE_TRAILING_SPACE = re.compile(r'[ \t]+$', re.M)
|
| 101 |
+
RE_NO_SPACE_AFTER_PERIOD = re.compile(r'([.!?])([A-Z])')
|
| 102 |
+
RE_DOUBLE_PERIOD = re.compile(r'\.{2}(?!\.)')
|
| 103 |
+
RE_SPACE_BEFORE_PUNCT = re.compile(r'\s+([.,;:!?])')
|
| 104 |
+
|
| 105 |
+
RE_CJK = re.compile(r'[\u4e00-\u9fff\u3040-\u309f\u30a0-\u30ff\uac00-\ud7af]{3,}')
|
| 106 |
+
RE_ARABIC = re.compile(r'[\u0600-\u06ff]{5,}')
|
| 107 |
+
RE_CYRILLIC = re.compile(r'[\u0400-\u04ff]{5,}')
|
| 108 |
+
RE_DEVANAGARI = re.compile(r'[\u0900-\u097f]{5,}')
|
| 109 |
+
RE_RESIDUAL_CODE = re.compile(
|
| 110 |
+
r'(function\s*\(|var\s+\w+\s*=|console\.log|document\.get|if\s*\(\s*\w+\s*[!=]==)', re.I
|
| 111 |
+
)
|
| 112 |
+
RE_COOKIE_LINE = re.compile(r'^.*(?:cookie|cookies)\s+(?:policy|consent|notice|preferences|settings).*$', re.I | re.M)
|
| 113 |
+
RE_SUBSCRIBE_LINE = re.compile(r'^.*(?:subscribe|sign\s*up\s+(?:for|to)\s+(?:our|the)\s+newsletter|unsubscribe|opt[\s-]*out\s+of).*$', re.I | re.M)
|
| 114 |
+
RE_CLICKBAIT_LINE = re.compile(r'^.*(?:you\s+won\'?t\s+believe|click\s+here|read\s+more\s*\.{0,3}$|share\s+this\s+(?:article|post|story)|trending\s+now|sponsored\s+content|advertisement).*$', re.I | re.M)
|
| 115 |
+
RE_SOCIAL_LINE = re.compile(r'^.*(?:follow\s+us\s+on|share\s+on\s+(?:facebook|twitter|linkedin|instagram)|like\s+us\s+on|tweet\s+this).*$', re.I | re.M)
|
| 116 |
+
RE_NAV_LINE = re.compile(r'^.*(?:skip\s+to\s+(?:main\s+)?content|back\s+to\s+top|previous\s+article|next\s+article|related\s+(?:articles|posts)).*$', re.I | re.M)
|
| 117 |
+
RE_LOGIN_LINE = re.compile(r'^.*(?:log\s*in\s+to\s+(?:your|an)\s+account|create\s+(?:a\s+)?(?:free\s+)?account|forgot\s+(?:your\s+)?password|already\s+(?:a\s+)?member).*$', re.I | re.M)
|
| 118 |
+
RE_COMMENT_LINE = re.compile(r'^.*(?:leave\s+a\s+(?:comment|reply)|post\s+a\s+comment|\d+\s+comments?$|logged\s+in\s+as).*$', re.I | re.M)
|
| 119 |
+
RE_COPYRIGHT_LINE = re.compile(r'^.*(?:all\s+rights\s+reserved|\(c\)\s*\d{4}|copyright\s+\d{4}).*$', re.I | re.M)
|
| 120 |
+
|
| 121 |
+
|
| 122 |
+
def clean_text_basic(text):
|
| 123 |
+
"""Light quality filter for download phase."""
|
| 124 |
+
text = unicodedata.normalize("NFKC", text)
|
| 125 |
+
text = re.sub(r'[\x00-\x08\x0b\x0c\x0e-\x1f\x7f-\x9f]', '', text)
|
| 126 |
+
text = re.sub(r'\n{3,}', '\n\n', text)
|
| 127 |
+
text = re.sub(r'[ \t]+', ' ', text)
|
| 128 |
+
text = '\n'.join(line.strip() for line in text.split('\n'))
|
| 129 |
+
return text.strip()
|
| 130 |
+
|
| 131 |
+
|
| 132 |
+
def is_high_quality(text, min_chars=500, min_words=80):
|
| 133 |
+
if len(text) < min_chars:
|
| 134 |
+
return False
|
| 135 |
+
words = text.split()
|
| 136 |
+
num_words = len(words)
|
| 137 |
+
if num_words < min_words:
|
| 138 |
+
return False
|
| 139 |
+
alpha = sum(c.isalpha() for c in text)
|
| 140 |
+
if alpha / max(len(text), 1) < 0.65:
|
| 141 |
+
return False
|
| 142 |
+
avg_word_len = sum(len(w) for w in words) / num_words
|
| 143 |
+
if avg_word_len < 2.5 or avg_word_len > 15:
|
| 144 |
+
return False
|
| 145 |
+
url_hits = text.count('http://') + text.count('https://')
|
| 146 |
+
if url_hits > num_words * 0.03:
|
| 147 |
+
return False
|
| 148 |
+
sentences = re.split(r'[.!?]+', text)
|
| 149 |
+
real_sentences = [s.strip() for s in sentences if len(s.strip()) > 10]
|
| 150 |
+
if len(real_sentences) < 3:
|
| 151 |
+
return False
|
| 152 |
+
lines = [ln.strip() for ln in text.split('\n') if ln.strip()]
|
| 153 |
+
if len(lines) > 5:
|
| 154 |
+
unique_ratio = len(set(lines)) / len(lines)
|
| 155 |
+
if unique_ratio < 0.5:
|
| 156 |
+
return False
|
| 157 |
+
return True
|
| 158 |
+
|
| 159 |
+
|
| 160 |
+
def deep_clean(text):
|
| 161 |
+
if not text or len(text.strip()) < 30:
|
| 162 |
+
return None
|
| 163 |
+
text = unicodedata.normalize("NFKC", text)
|
| 164 |
+
for ch in CONTROL_CHARS:
|
| 165 |
+
text = text.replace(ch, "")
|
| 166 |
+
for old, new in HTML_ENTITIES:
|
| 167 |
+
text = text.replace(old, new)
|
| 168 |
+
text = RE_HTML_COMMENT.sub("", text)
|
| 169 |
+
text = RE_HTML_TAG.sub("", text)
|
| 170 |
+
text = RE_URL.sub("", text)
|
| 171 |
+
text = RE_EMAIL.sub("", text)
|
| 172 |
+
text = RE_FILE_PATH.sub("", text)
|
| 173 |
+
text = RE_CODE_BLOCK.sub("", text)
|
| 174 |
+
text = RE_REPEATED_LINE.sub(r'\1', text)
|
| 175 |
+
text = RE_REPEATED_PUNCT.sub(r'\1\1\1', text)
|
| 176 |
+
text = RE_REPEATED_CHAR.sub(r'\1\1\1', text)
|
| 177 |
+
text = RE_REPEATED_WORD.sub(r'\1', text)
|
| 178 |
+
text = text.replace('\t', ' ')
|
| 179 |
+
text = RE_TRAILING_SPACE.sub('', text)
|
| 180 |
+
text = RE_MULTI_SPACE.sub(' ', text)
|
| 181 |
+
text = RE_MULTI_NEWLINE.sub('\n\n\n', text)
|
| 182 |
+
text = RE_DOUBLE_PERIOD.sub('.', text)
|
| 183 |
+
text = RE_NO_SPACE_AFTER_PERIOD.sub(r'\1 \2', text)
|
| 184 |
+
text = RE_SPACE_BEFORE_PUNCT.sub(r'\1', text)
|
| 185 |
+
text = text.replace('\u2018', "'").replace('\u2019', "'")
|
| 186 |
+
text = text.replace('\u201c', '"').replace('\u201d', '"')
|
| 187 |
+
text = text.replace('\u2013', '-').replace('\u2014', ' - ')
|
| 188 |
+
text = text.replace('\u2026', '...')
|
| 189 |
+
text = text.replace('\u2022', '- ')
|
| 190 |
+
text = text.replace('\u00b7', ' ')
|
| 191 |
+
text = text.replace('\u00a0', ' ')
|
| 192 |
+
lines = text.split('\n')
|
| 193 |
+
clean_lines = []
|
| 194 |
+
for line in lines:
|
| 195 |
+
line = line.strip()
|
| 196 |
+
if not line:
|
| 197 |
+
clean_lines.append('')
|
| 198 |
+
continue
|
| 199 |
+
if len(line) > 10:
|
| 200 |
+
alpha_count = sum(1 for c in line if c.isalpha())
|
| 201 |
+
if alpha_count / len(line) < 0.40:
|
| 202 |
+
continue
|
| 203 |
+
if line.count('|') > 3 or line.count('{') > 2 or line.count('}') > 2:
|
| 204 |
+
continue
|
| 205 |
+
if RE_IMPORT.match(line):
|
| 206 |
+
continue
|
| 207 |
+
if line and line[0].isalpha() and line[0].islower():
|
| 208 |
+
if not clean_lines or clean_lines[-1] == '' or clean_lines[-1].rstrip().endswith(('.', '!', '?', ':')):
|
| 209 |
+
line = line[0].upper() + line[1:]
|
| 210 |
+
clean_lines.append(line)
|
| 211 |
+
text = '\n'.join(clean_lines)
|
| 212 |
+
text = text.strip()
|
| 213 |
+
paragraphs = text.split('\n\n')
|
| 214 |
+
seen = set()
|
| 215 |
+
unique_paragraphs = []
|
| 216 |
+
for p in paragraphs:
|
| 217 |
+
p_stripped = p.strip()
|
| 218 |
+
if not p_stripped:
|
| 219 |
+
continue
|
| 220 |
+
p_key = ' '.join(p_stripped.lower().split())
|
| 221 |
+
if p_key not in seen:
|
| 222 |
+
seen.add(p_key)
|
| 223 |
+
unique_paragraphs.append(p_stripped)
|
| 224 |
+
text = '\n\n'.join(unique_paragraphs)
|
| 225 |
+
text = text.strip()
|
| 226 |
+
if len(text) < 50:
|
| 227 |
+
return None
|
| 228 |
+
if len(text.split()) < 10:
|
| 229 |
+
return None
|
| 230 |
+
ascii_count = sum(1 for c in text if ord(c) < 128)
|
| 231 |
+
if ascii_count / max(len(text), 1) < 0.85:
|
| 232 |
+
return None
|
| 233 |
+
return text
|
| 234 |
+
|
| 235 |
+
|
| 236 |
+
def smart_filter(text):
|
| 237 |
+
words = text.split()
|
| 238 |
+
word_count = len(words)
|
| 239 |
+
if word_count < 50:
|
| 240 |
+
return False, text, "too short"
|
| 241 |
+
scripts = []
|
| 242 |
+
if RE_CJK.search(text): scripts.append("CJK")
|
| 243 |
+
if RE_ARABIC.search(text): scripts.append("Arabic")
|
| 244 |
+
if RE_CYRILLIC.search(text): scripts.append("Cyrillic")
|
| 245 |
+
if RE_DEVANAGARI.search(text): scripts.append("Devanagari")
|
| 246 |
+
if scripts:
|
| 247 |
+
return False, text, "non-English"
|
| 248 |
+
if word_count > 50:
|
| 249 |
+
unique_ratio = len(set(w.lower() for w in words)) / word_count
|
| 250 |
+
if unique_ratio < 0.20:
|
| 251 |
+
return False, text, "repetitive"
|
| 252 |
+
code_matches = RE_RESIDUAL_CODE.findall(text)
|
| 253 |
+
if len(code_matches) >= 5:
|
| 254 |
+
return False, text, "residual code"
|
| 255 |
+
for pattern in [RE_COOKIE_LINE, RE_SUBSCRIBE_LINE, RE_CLICKBAIT_LINE,
|
| 256 |
+
RE_SOCIAL_LINE, RE_NAV_LINE, RE_LOGIN_LINE,
|
| 257 |
+
RE_COMMENT_LINE, RE_COPYRIGHT_LINE]:
|
| 258 |
+
text = pattern.sub('', text)
|
| 259 |
+
lines = text.split('\n')
|
| 260 |
+
clean_lines = []
|
| 261 |
+
for line in lines:
|
| 262 |
+
stripped = line.strip()
|
| 263 |
+
if stripped and len(stripped) > 10:
|
| 264 |
+
digit_count = sum(1 for c in stripped if c.isdigit() or c in ' ,.\t-+/%$')
|
| 265 |
+
if digit_count / len(stripped) > 0.80:
|
| 266 |
+
continue
|
| 267 |
+
clean_lines.append(line)
|
| 268 |
+
text = '\n'.join(clean_lines)
|
| 269 |
+
text = re.sub(r'\n{3,}', '\n\n', text)
|
| 270 |
+
text = text.strip()
|
| 271 |
+
if len(text.split()) < 50:
|
| 272 |
+
return False, text, "too short after stripping"
|
| 273 |
+
return True, text, None
|
| 274 |
+
|
| 275 |
+
|
| 276 |
+
def clean_and_filter(text):
|
| 277 |
+
result = deep_clean(text)
|
| 278 |
+
if result is None:
|
| 279 |
+
return None, "deep_clean_drop", 0
|
| 280 |
+
keep, stripped, reason = smart_filter(result)
|
| 281 |
+
if not keep:
|
| 282 |
+
return None, reason, 0
|
| 283 |
+
return stripped, None, len(result) - len(stripped)
|
| 284 |
+
|
| 285 |
+
|
| 286 |
+
# ==============================================================================
|
| 287 |
+
# STREAMING CHUNK WRITER (appends to existing litdata)
|
| 288 |
+
# ==============================================================================
|
| 289 |
+
|
| 290 |
+
class StreamingChunkWriter:
|
| 291 |
+
def __init__(self, output_dir, config, start_chunk_idx=0):
|
| 292 |
+
self.output_dir = Path(output_dir)
|
| 293 |
+
os.makedirs(str(self.output_dir), exist_ok=True)
|
| 294 |
+
self.config = config
|
| 295 |
+
self.dtype_size = DTYPE().itemsize
|
| 296 |
+
self.tokens_per_chunk = (CHUNK_BYTES_TARGET // self.dtype_size // BLOCK_SIZE) * BLOCK_SIZE
|
| 297 |
+
self.buffer = []
|
| 298 |
+
self.chunks_metadata = []
|
| 299 |
+
self.chunk_idx = start_chunk_idx
|
| 300 |
+
self.total_tokens = 0
|
| 301 |
+
|
| 302 |
+
def add_tokens_batch(self, encoded_batch):
|
| 303 |
+
for enc in encoded_batch:
|
| 304 |
+
ids = enc.ids
|
| 305 |
+
self.buffer.extend(ids)
|
| 306 |
+
self.buffer.append(EOS_TOKEN_ID)
|
| 307 |
+
while len(self.buffer) >= self.tokens_per_chunk:
|
| 308 |
+
self._flush_chunk()
|
| 309 |
+
|
| 310 |
+
def _flush_chunk(self):
|
| 311 |
+
if len(self.buffer) < BLOCK_SIZE:
|
| 312 |
+
return
|
| 313 |
+
take = min(len(self.buffer), self.tokens_per_chunk)
|
| 314 |
+
num_blocks = take // BLOCK_SIZE
|
| 315 |
+
if num_blocks == 0:
|
| 316 |
+
return
|
| 317 |
+
actual_tokens = num_blocks * BLOCK_SIZE
|
| 318 |
+
chunk_data = np.array(self.buffer[:actual_tokens], dtype=DTYPE)
|
| 319 |
+
self.buffer = self.buffer[actual_tokens:]
|
| 320 |
+
filename = f"chunk-0-{self.chunk_idx}.bin"
|
| 321 |
+
filepath = self.output_dir / filename
|
| 322 |
+
header_num = np.array([num_blocks], dtype=np.uint32)
|
| 323 |
+
offsets = np.arange(num_blocks + 1, dtype=np.uint32) * (BLOCK_SIZE * self.dtype_size)
|
| 324 |
+
header = np.concatenate([header_num, offsets])
|
| 325 |
+
with open(filepath, "wb") as f:
|
| 326 |
+
header.tofile(f)
|
| 327 |
+
chunk_data.tofile(f)
|
| 328 |
+
meta = {
|
| 329 |
+
"chunk_bytes": int(header.nbytes + chunk_data.nbytes),
|
| 330 |
+
"chunk_size": num_blocks,
|
| 331 |
+
"dim": int(actual_tokens),
|
| 332 |
+
"filename": filename,
|
| 333 |
+
}
|
| 334 |
+
self.chunks_metadata.append(meta)
|
| 335 |
+
self.total_tokens += actual_tokens
|
| 336 |
+
self.chunk_idx += 1
|
| 337 |
+
if self.chunk_idx % 10 == 0:
|
| 338 |
+
print(f" Flushed chunk {self.chunk_idx} ({self.total_tokens:,} new tokens)")
|
| 339 |
+
|
| 340 |
+
def finalize(self):
|
| 341 |
+
while len(self.buffer) >= BLOCK_SIZE:
|
| 342 |
+
self._flush_chunk()
|
| 343 |
+
discarded = len(self.buffer)
|
| 344 |
+
self.buffer = []
|
| 345 |
+
return self.total_tokens, discarded
|
| 346 |
+
|
| 347 |
+
|
| 348 |
+
# ==============================================================================
|
| 349 |
+
# PROCESS ONE SOURCE: download β clean (multiprocessed) β tokenize β write
|
| 350 |
+
# ==============================================================================
|
| 351 |
+
|
| 352 |
+
def process_source_batch(texts, pool, tokenizer, writer, source_name, batch_num):
|
| 353 |
+
"""Clean a batch of texts and write to the streaming writer. Returns stats."""
|
| 354 |
+
t0 = time.time()
|
| 355 |
+
|
| 356 |
+
# Clean with multiprocessing
|
| 357 |
+
results = pool.map(clean_and_filter, texts, chunksize=256)
|
| 358 |
+
|
| 359 |
+
cleaned = []
|
| 360 |
+
dropped = 0
|
| 361 |
+
reasons = Counter()
|
| 362 |
+
boilerplate = 0
|
| 363 |
+
for text, reason, bp in results:
|
| 364 |
+
if text is not None:
|
| 365 |
+
cleaned.append(text)
|
| 366 |
+
boilerplate += bp
|
| 367 |
+
else:
|
| 368 |
+
dropped += 1
|
| 369 |
+
reasons[reason] += 1
|
| 370 |
+
del results
|
| 371 |
+
|
| 372 |
+
# Tokenize + write
|
| 373 |
+
for i in range(0, len(cleaned), ENCODE_BATCH):
|
| 374 |
+
sub = cleaned[i:i+ENCODE_BATCH]
|
| 375 |
+
encoded = tokenizer.encode_batch(sub, add_special_tokens=False)
|
| 376 |
+
writer.add_tokens_batch(encoded)
|
| 377 |
+
del cleaned
|
| 378 |
+
|
| 379 |
+
elapsed = time.time() - t0
|
| 380 |
+
print(f" [{source_name}] batch {batch_num}: kept {len(texts)-dropped:,} / dropped {dropped} | {elapsed:.1f}s")
|
| 381 |
+
|
| 382 |
+
return len(texts) - dropped, dropped, reasons, boilerplate
|
| 383 |
+
|
| 384 |
+
|
| 385 |
+
# ==============================================================================
|
| 386 |
+
# MAIN
|
| 387 |
+
# ==============================================================================
|
| 388 |
+
|
| 389 |
+
def main():
|
| 390 |
+
from datasets import load_dataset
|
| 391 |
+
|
| 392 |
+
t_start = time.time()
|
| 393 |
+
|
| 394 |
+
print("Loading tokenizer...")
|
| 395 |
+
tokenizer = Tokenizer.from_file(TOKENIZER_PATH)
|
| 396 |
+
config = {"block_size": BLOCK_SIZE, "vocab_size": tokenizer.get_vocab_size()}
|
| 397 |
+
|
| 398 |
+
# Read existing index to find starting chunk offset
|
| 399 |
+
with open(FINAL_DIR / "index.json") as f:
|
| 400 |
+
existing_index = json.load(f)
|
| 401 |
+
existing_chunks = existing_index["chunks"]
|
| 402 |
+
existing_tokens = sum(c["dim"] for c in existing_chunks)
|
| 403 |
+
start_chunk_idx = len(existing_chunks)
|
| 404 |
+
|
| 405 |
+
print(f"\n{'='*75}")
|
| 406 |
+
print(f" BUILD 1.2B DIVERSE TOKENS β APPEND TO litdata_pretrain_final")
|
| 407 |
+
print(f" Current: {existing_tokens:,} tokens ({start_chunk_idx} chunks)")
|
| 408 |
+
print(f" Target: 5,000,000,000 tokens")
|
| 409 |
+
print(f" Building: ~{TOTAL_TARGET:,} estimated tokens")
|
| 410 |
+
print(f" Workers: {NUM_WORKERS} CPU cores")
|
| 411 |
+
print(f"{'='*75}")
|
| 412 |
+
|
| 413 |
+
# Initialize writer that appends new chunks after existing ones
|
| 414 |
+
writer = StreamingChunkWriter(str(FINAL_DIR), config, start_chunk_idx=start_chunk_idx)
|
| 415 |
+
pool = Pool(processes=NUM_WORKERS)
|
| 416 |
+
|
| 417 |
+
# Global stats
|
| 418 |
+
all_stats = {}
|
| 419 |
+
DOWNLOAD_BATCH = 10000 # process 10K docs at a time
|
| 420 |
+
|
| 421 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 422 |
+
# SOURCE 1: Wikipedia (encyclopedic knowledge)
|
| 423 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 424 |
+
print(f"\n{'='*75}")
|
| 425 |
+
print(f" SOURCE 1: WIKIPEDIA (skip {WIKI_SKIP:,}, target ~{WIKI_TARGET:,} tokens)")
|
| 426 |
+
print(f"{'='*75}")
|
| 427 |
+
|
| 428 |
+
ds_wiki = load_dataset(
|
| 429 |
+
"wikimedia/wikipedia", "20231101.en",
|
| 430 |
+
split="train", streaming=True, trust_remote_code=False
|
| 431 |
+
)
|
| 432 |
+
|
| 433 |
+
wiki_texts = []
|
| 434 |
+
wiki_tokens_est = 0
|
| 435 |
+
wiki_seen = 0
|
| 436 |
+
wiki_skipped_quality = 0
|
| 437 |
+
wiki_skipped_meta = 0
|
| 438 |
+
wiki_total_kept = 0
|
| 439 |
+
wiki_total_dropped = 0
|
| 440 |
+
wiki_reasons = Counter()
|
| 441 |
+
wiki_boilerplate = 0
|
| 442 |
+
wiki_batch_num = 0
|
| 443 |
+
t0 = time.time()
|
| 444 |
+
|
| 445 |
+
for article in ds_wiki:
|
| 446 |
+
title = (article.get("title") or "").strip()
|
| 447 |
+
raw = article.get("text") or ""
|
| 448 |
+
|
| 449 |
+
if _WIKI_SKIP_PATTERNS.search(title):
|
| 450 |
+
wiki_skipped_meta += 1
|
| 451 |
+
continue
|
| 452 |
+
|
| 453 |
+
cleaned = clean_text_basic(raw)
|
| 454 |
+
if not is_high_quality(cleaned, min_chars=800, min_words=120):
|
| 455 |
+
wiki_skipped_quality += 1
|
| 456 |
+
continue
|
| 457 |
+
|
| 458 |
+
wiki_seen += 1
|
| 459 |
+
if wiki_seen <= WIKI_SKIP:
|
| 460 |
+
if wiki_seen % 25000 == 0:
|
| 461 |
+
print(f" Skipping... {wiki_seen:,}/{WIKI_SKIP:,}")
|
| 462 |
+
continue
|
| 463 |
+
|
| 464 |
+
full_text = f"{title}\n\n{cleaned}"
|
| 465 |
+
est_tok = int(len(full_text.split()) * TOKENS_PER_WORD)
|
| 466 |
+
wiki_texts.append(full_text)
|
| 467 |
+
wiki_tokens_est += est_tok
|
| 468 |
+
|
| 469 |
+
# Process in batches to keep memory low
|
| 470 |
+
if len(wiki_texts) >= DOWNLOAD_BATCH:
|
| 471 |
+
wiki_batch_num += 1
|
| 472 |
+
kept, dropped, reasons, bp = process_source_batch(
|
| 473 |
+
wiki_texts, pool, tokenizer, writer, "Wiki", wiki_batch_num
|
| 474 |
+
)
|
| 475 |
+
wiki_total_kept += kept
|
| 476 |
+
wiki_total_dropped += dropped
|
| 477 |
+
wiki_reasons += reasons
|
| 478 |
+
wiki_boilerplate += bp
|
| 479 |
+
wiki_texts = []
|
| 480 |
+
|
| 481 |
+
if wiki_tokens_est >= WIKI_TARGET:
|
| 482 |
+
break
|
| 483 |
+
|
| 484 |
+
# Process remaining
|
| 485 |
+
if wiki_texts:
|
| 486 |
+
wiki_batch_num += 1
|
| 487 |
+
kept, dropped, reasons, bp = process_source_batch(
|
| 488 |
+
wiki_texts, pool, tokenizer, writer, "Wiki", wiki_batch_num
|
| 489 |
+
)
|
| 490 |
+
wiki_total_kept += kept
|
| 491 |
+
wiki_total_dropped += dropped
|
| 492 |
+
wiki_reasons += reasons
|
| 493 |
+
wiki_boilerplate += bp
|
| 494 |
+
wiki_texts = []
|
| 495 |
+
|
| 496 |
+
wiki_elapsed = time.time() - t0
|
| 497 |
+
print(f" [Wikipedia] Done: ~{wiki_tokens_est:,} est. tokens, kept {wiki_total_kept:,}, dropped {wiki_total_dropped:,} in {wiki_elapsed:.0f}s")
|
| 498 |
+
print(f" (skipped {WIKI_SKIP:,} already-used + {wiki_skipped_quality:,} low-quality + {wiki_skipped_meta:,} meta)")
|
| 499 |
+
all_stats["Wikipedia"] = {
|
| 500 |
+
"est_tokens": wiki_tokens_est, "kept": wiki_total_kept,
|
| 501 |
+
"dropped": wiki_total_dropped, "reasons": wiki_reasons,
|
| 502 |
+
"boilerplate": wiki_boilerplate, "time": wiki_elapsed,
|
| 503 |
+
}
|
| 504 |
+
|
| 505 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 506 |
+
# SOURCE 2: FineWeb-Edu (educational web content)
|
| 507 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 508 |
+
print(f"\n{'='*75}")
|
| 509 |
+
print(f" SOURCE 2: FINEWEB-EDU (skip {FINEWEB_SKIP:,}, score >= {FINEWEB_MIN_SCORE}, target ~{FINEWEB_TARGET:,} tokens)")
|
| 510 |
+
print(f"{'='*75}")
|
| 511 |
+
|
| 512 |
+
ds_fw = load_dataset(
|
| 513 |
+
"HuggingFaceFW/fineweb-edu", "sample-10BT",
|
| 514 |
+
split="train", streaming=True, trust_remote_code=False
|
| 515 |
+
)
|
| 516 |
+
|
| 517 |
+
fw_texts = []
|
| 518 |
+
fw_tokens_est = 0
|
| 519 |
+
fw_seen = 0
|
| 520 |
+
fw_skipped_quality = 0
|
| 521 |
+
fw_skipped_score = 0
|
| 522 |
+
fw_total_kept = 0
|
| 523 |
+
fw_total_dropped = 0
|
| 524 |
+
fw_reasons = Counter()
|
| 525 |
+
fw_boilerplate = 0
|
| 526 |
+
fw_batch_num = 0
|
| 527 |
+
t0 = time.time()
|
| 528 |
+
|
| 529 |
+
for doc in ds_fw:
|
| 530 |
+
score = doc.get("score", 0)
|
| 531 |
+
if not isinstance(score, (int, float)):
|
| 532 |
+
try:
|
| 533 |
+
score = float(score)
|
| 534 |
+
except (ValueError, TypeError):
|
| 535 |
+
continue
|
| 536 |
+
if score < FINEWEB_MIN_SCORE:
|
| 537 |
+
fw_skipped_score += 1
|
| 538 |
+
continue
|
| 539 |
+
|
| 540 |
+
raw = doc.get("text") or ""
|
| 541 |
+
cleaned = clean_text_basic(raw)
|
| 542 |
+
if not is_high_quality(cleaned, min_chars=500, min_words=80):
|
| 543 |
+
fw_skipped_quality += 1
|
| 544 |
+
continue
|
| 545 |
+
|
| 546 |
+
fw_seen += 1
|
| 547 |
+
if fw_seen <= FINEWEB_SKIP:
|
| 548 |
+
if fw_seen % 25000 == 0:
|
| 549 |
+
print(f" Skipping... {fw_seen:,}/{FINEWEB_SKIP:,}")
|
| 550 |
+
continue
|
| 551 |
+
|
| 552 |
+
est_tok = int(len(cleaned.split()) * TOKENS_PER_WORD)
|
| 553 |
+
fw_texts.append(cleaned)
|
| 554 |
+
fw_tokens_est += est_tok
|
| 555 |
+
|
| 556 |
+
if len(fw_texts) >= DOWNLOAD_BATCH:
|
| 557 |
+
fw_batch_num += 1
|
| 558 |
+
kept, dropped, reasons, bp = process_source_batch(
|
| 559 |
+
fw_texts, pool, tokenizer, writer, "FineWeb", fw_batch_num
|
| 560 |
+
)
|
| 561 |
+
fw_total_kept += kept
|
| 562 |
+
fw_total_dropped += dropped
|
| 563 |
+
fw_reasons += reasons
|
| 564 |
+
fw_boilerplate += bp
|
| 565 |
+
fw_texts = []
|
| 566 |
+
|
| 567 |
+
if fw_tokens_est >= FINEWEB_TARGET:
|
| 568 |
+
break
|
| 569 |
+
|
| 570 |
+
if fw_texts:
|
| 571 |
+
fw_batch_num += 1
|
| 572 |
+
kept, dropped, reasons, bp = process_source_batch(
|
| 573 |
+
fw_texts, pool, tokenizer, writer, "FineWeb", fw_batch_num
|
| 574 |
+
)
|
| 575 |
+
fw_total_kept += kept
|
| 576 |
+
fw_total_dropped += dropped
|
| 577 |
+
fw_reasons += reasons
|
| 578 |
+
fw_boilerplate += bp
|
| 579 |
+
fw_texts = []
|
| 580 |
+
|
| 581 |
+
fw_elapsed = time.time() - t0
|
| 582 |
+
print(f" [FineWeb-Edu] Done: ~{fw_tokens_est:,} est. tokens, kept {fw_total_kept:,}, dropped {fw_total_dropped:,} in {fw_elapsed:.0f}s")
|
| 583 |
+
print(f" (skipped {FINEWEB_SKIP:,} already-used + {fw_skipped_quality:,} low-quality + {fw_skipped_score:,} low-score)")
|
| 584 |
+
all_stats["FineWeb-Edu"] = {
|
| 585 |
+
"est_tokens": fw_tokens_est, "kept": fw_total_kept,
|
| 586 |
+
"dropped": fw_total_dropped, "reasons": fw_reasons,
|
| 587 |
+
"boilerplate": fw_boilerplate, "time": fw_elapsed,
|
| 588 |
+
}
|
| 589 |
+
|
| 590 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 591 |
+
# SOURCE 3: OpenWebText (Reddit-curated diverse web pages)
|
| 592 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 593 |
+
print(f"\n{'='*75}")
|
| 594 |
+
print(f" SOURCE 3: OPENWEBTEXT (target ~{OWT_TARGET:,} tokens)")
|
| 595 |
+
print(f" Completely new source β zero overlap with existing data")
|
| 596 |
+
print(f"{'='*75}")
|
| 597 |
+
|
| 598 |
+
ds_owt = load_dataset(
|
| 599 |
+
"Skylion007/openwebtext",
|
| 600 |
+
split="train", streaming=True, trust_remote_code=False
|
| 601 |
+
)
|
| 602 |
+
|
| 603 |
+
owt_texts = []
|
| 604 |
+
owt_tokens_est = 0
|
| 605 |
+
owt_skipped_quality = 0
|
| 606 |
+
owt_total_kept = 0
|
| 607 |
+
owt_total_dropped = 0
|
| 608 |
+
owt_reasons = Counter()
|
| 609 |
+
owt_boilerplate = 0
|
| 610 |
+
owt_batch_num = 0
|
| 611 |
+
t0 = time.time()
|
| 612 |
+
|
| 613 |
+
for doc in ds_owt:
|
| 614 |
+
raw = doc.get("text") or ""
|
| 615 |
+
cleaned = clean_text_basic(raw)
|
| 616 |
+
if not is_high_quality(cleaned, min_chars=400, min_words=60):
|
| 617 |
+
owt_skipped_quality += 1
|
| 618 |
+
continue
|
| 619 |
+
|
| 620 |
+
est_tok = int(len(cleaned.split()) * TOKENS_PER_WORD)
|
| 621 |
+
owt_texts.append(cleaned)
|
| 622 |
+
owt_tokens_est += est_tok
|
| 623 |
+
|
| 624 |
+
if len(owt_texts) >= DOWNLOAD_BATCH:
|
| 625 |
+
owt_batch_num += 1
|
| 626 |
+
kept, dropped, reasons, bp = process_source_batch(
|
| 627 |
+
owt_texts, pool, tokenizer, writer, "OWT", owt_batch_num
|
| 628 |
+
)
|
| 629 |
+
owt_total_kept += kept
|
| 630 |
+
owt_total_dropped += dropped
|
| 631 |
+
owt_reasons += reasons
|
| 632 |
+
owt_boilerplate += bp
|
| 633 |
+
owt_texts = []
|
| 634 |
+
|
| 635 |
+
if owt_tokens_est >= OWT_TARGET:
|
| 636 |
+
break
|
| 637 |
+
|
| 638 |
+
if owt_texts:
|
| 639 |
+
owt_batch_num += 1
|
| 640 |
+
kept, dropped, reasons, bp = process_source_batch(
|
| 641 |
+
owt_texts, pool, tokenizer, writer, "OWT", owt_batch_num
|
| 642 |
+
)
|
| 643 |
+
owt_total_kept += kept
|
| 644 |
+
owt_total_dropped += dropped
|
| 645 |
+
owt_reasons += reasons
|
| 646 |
+
owt_boilerplate += bp
|
| 647 |
+
owt_texts = []
|
| 648 |
+
|
| 649 |
+
owt_elapsed = time.time() - t0
|
| 650 |
+
print(f" [OpenWebText] Done: ~{owt_tokens_est:,} est. tokens, kept {owt_total_kept:,}, dropped {owt_total_dropped:,} in {owt_elapsed:.0f}s")
|
| 651 |
+
print(f" (skipped {owt_skipped_quality:,} low-quality)")
|
| 652 |
+
all_stats["OpenWebText"] = {
|
| 653 |
+
"est_tokens": owt_tokens_est, "kept": owt_total_kept,
|
| 654 |
+
"dropped": owt_total_dropped, "reasons": owt_reasons,
|
| 655 |
+
"boilerplate": owt_boilerplate, "time": owt_elapsed,
|
| 656 |
+
}
|
| 657 |
+
|
| 658 |
+
pool.close()
|
| 659 |
+
pool.join()
|
| 660 |
+
|
| 661 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 662 |
+
# FINALIZE: flush remaining + update index.json
|
| 663 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 664 |
+
print(f"\n{'='*75}")
|
| 665 |
+
print(f" FINALIZING")
|
| 666 |
+
print(f"{'='*75}")
|
| 667 |
+
|
| 668 |
+
new_tokens, discarded = writer.finalize()
|
| 669 |
+
|
| 670 |
+
# Merge new chunk metadata with existing
|
| 671 |
+
final_chunks = existing_chunks + writer.chunks_metadata
|
| 672 |
+
final_total = existing_tokens + new_tokens
|
| 673 |
+
final_index = {
|
| 674 |
+
"chunks": final_chunks,
|
| 675 |
+
"config": config,
|
| 676 |
+
"updated_at": str(time.time()),
|
| 677 |
+
}
|
| 678 |
+
with open(FINAL_DIR / "index.json", "w") as f:
|
| 679 |
+
json.dump(final_index, f, indent=2)
|
| 680 |
+
|
| 681 |
+
total_time = time.time() - t_start
|
| 682 |
+
|
| 683 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 684 |
+
# REPORT
|
| 685 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 686 |
+
total_new_kept = wiki_total_kept + fw_total_kept + owt_total_kept
|
| 687 |
+
total_new_dropped = wiki_total_dropped + fw_total_dropped + owt_total_dropped
|
| 688 |
+
total_new_boilerplate = wiki_boilerplate + fw_boilerplate + owt_boilerplate
|
| 689 |
+
all_drop_reasons = wiki_reasons + fw_reasons + owt_reasons
|
| 690 |
+
|
| 691 |
+
report = []
|
| 692 |
+
report.append(f"{'='*75}")
|
| 693 |
+
report.append(f" 5 BILLION TOKEN PRETRAIN DATASET β BUILD REPORT")
|
| 694 |
+
report.append(f"{'='*75}")
|
| 695 |
+
report.append(f"")
|
| 696 |
+
report.append(f" Total time: {total_time:.0f}s ({total_time/60:.1f} min)")
|
| 697 |
+
report.append(f" Workers: {NUM_WORKERS} CPU cores")
|
| 698 |
+
report.append(f"")
|
| 699 |
+
report.append(f" NEW DATA ADDED (diverse, clean English)")
|
| 700 |
+
report.append(f" {'-'*60}")
|
| 701 |
+
|
| 702 |
+
for name, stats in all_stats.items():
|
| 703 |
+
report.append(f" {name}:")
|
| 704 |
+
report.append(f" Est tokens: ~{stats['est_tokens']:,}")
|
| 705 |
+
report.append(f" Kept: {stats['kept']:,} | Dropped: {stats['dropped']:,}")
|
| 706 |
+
report.append(f" Boilerplate: {stats['boilerplate']:,} chars")
|
| 707 |
+
report.append(f" Time: {stats['time']:.0f}s")
|
| 708 |
+
if stats["reasons"]:
|
| 709 |
+
for reason, count in sorted(stats["reasons"].items(), key=lambda x: -x[1]):
|
| 710 |
+
report.append(f" {reason}: {count:,}")
|
| 711 |
+
|
| 712 |
+
report.append(f"")
|
| 713 |
+
report.append(f" NEW DATA TOTALS")
|
| 714 |
+
report.append(f" {'-'*60}")
|
| 715 |
+
report.append(f" Documents kept: {total_new_kept:,}")
|
| 716 |
+
report.append(f" Documents dropped: {total_new_dropped:,}")
|
| 717 |
+
report.append(f" Boilerplate: {total_new_boilerplate:,} chars stripped")
|
| 718 |
+
report.append(f" New tokens: {new_tokens:,} ({writer.chunk_idx - start_chunk_idx} chunks)")
|
| 719 |
+
if all_drop_reasons:
|
| 720 |
+
report.append(f" Drop reasons:")
|
| 721 |
+
for reason, count in sorted(all_drop_reasons.items(), key=lambda x: -x[1]):
|
| 722 |
+
report.append(f" {reason:<35} {count:>8,}")
|
| 723 |
+
|
| 724 |
+
report.append(f"")
|
| 725 |
+
report.append(f" FINAL COMBINED DATASET")
|
| 726 |
+
report.append(f" {'-'*60}")
|
| 727 |
+
report.append(f" Location: {FINAL_DIR}")
|
| 728 |
+
report.append(f" Chunks: {len(final_chunks)}")
|
| 729 |
+
report.append(f" Tokens: {final_total:,}")
|
| 730 |
+
report.append(f" Format: litdata binary (int32, BLOCK_SIZE=1025, EOS=0)")
|
| 731 |
+
report.append(f"")
|
| 732 |
+
report.append(f" Previous: {existing_tokens:,} tokens ({start_chunk_idx} chunks)")
|
| 733 |
+
report.append(f" + Added: {new_tokens:,} tokens ({writer.chunk_idx - start_chunk_idx} chunks)")
|
| 734 |
+
report.append(f" = Total: {final_total:,} tokens ({len(final_chunks)} chunks)")
|
| 735 |
+
report.append(f"")
|
| 736 |
+
report.append(f" DATA COMPOSITION")
|
| 737 |
+
report.append(f" {'-'*60}")
|
| 738 |
+
report.append(f" litdata_3b_clean: ~2.94B tokens (general web, cleaned)")
|
| 739 |
+
report.append(f" litdata_english_500m: ~515M tokens (Wiki+FineWeb, cleaned)")
|
| 740 |
+
report.append(f" litdata_combined: ~257M tokens (Wiki+FineWeb, cleaned)")
|
| 741 |
+
report.append(f" + Wikipedia (new): ~{wiki_tokens_est:,} est. (articles {WIKI_SKIP+1:,}+)")
|
| 742 |
+
report.append(f" + FineWeb-Edu (new): ~{fw_tokens_est:,} est. (score>={FINEWEB_MIN_SCORE}, docs {FINEWEB_SKIP+1:,}+)")
|
| 743 |
+
report.append(f" + OpenWebText (new): ~{owt_tokens_est:,} est. (Reddit-curated, no overlap)")
|
| 744 |
+
report.append(f"")
|
| 745 |
+
report.append(f" PURE ENGLISH PRETRAINING TEXT")
|
| 746 |
+
report.append(f" Sources: Wikipedia, FineWeb-Edu, OpenWebText, general web")
|
| 747 |
+
report.append(f" NO instruction/finetune data included")
|
| 748 |
+
report.append(f" ZERO overlap between all data sources")
|
| 749 |
+
report.append(f"{'='*75}")
|
| 750 |
+
|
| 751 |
+
full_report = '\n'.join(report)
|
| 752 |
+
print(f"\n{full_report}")
|
| 753 |
+
|
| 754 |
+
with open(FINAL_DIR / "BUILD_REPORT.txt", "w", encoding="utf-8") as f:
|
| 755 |
+
f.write(full_report)
|
| 756 |
+
print(f"\n Report saved to: {FINAL_DIR / 'BUILD_REPORT.txt'}")
|
| 757 |
+
print(f" Done! 5B token pretrain dataset ready.")
|
| 758 |
+
|
| 759 |
+
|
| 760 |
+
if __name__ == "__main__":
|
| 761 |
+
main()
|
Base/scripts/build_english_corpus.py
ADDED
|
@@ -0,0 +1,349 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Build an ultra-clean English text corpus for continued pretraining.
|
| 3 |
+
|
| 4 |
+
Downloads from high-quality, professionally curated sources via HuggingFace
|
| 5 |
+
streaming (constant memory, no full dataset download needed):
|
| 6 |
+
|
| 7 |
+
1. Wikipedia English β encyclopedic, multi-editor-reviewed text
|
| 8 |
+
(the single cleanest large-scale English corpus in existence)
|
| 9 |
+
2. FineWeb-Edu β top-scored educational web content (score >= 4.0)
|
| 10 |
+
(adds diversity: textbooks, tutorials, explanations, articles)
|
| 11 |
+
|
| 12 |
+
Quality filters applied to every text:
|
| 13 |
+
- Minimum length and word count
|
| 14 |
+
- High alphabetic ratio (rejects tables, code, symbol dumps)
|
| 15 |
+
- Sentence structure check (real prose, not lists/fragments)
|
| 16 |
+
- Repetition detection (rejects copy-paste / boilerplate)
|
| 17 |
+
- URL density filter (rejects link farms)
|
| 18 |
+
- Unicode normalization and whitespace cleanup
|
| 19 |
+
|
| 20 |
+
Outputs parquet files to Base/data/filtered_english/ with a single 'text'
|
| 21 |
+
column, compatible with the existing prepare_litdata.py pipeline.
|
| 22 |
+
|
| 23 |
+
After running this script:
|
| 24 |
+
1. python Base/scripts/prepare_litdata.py --mode english
|
| 25 |
+
2. litgpt pretrain --config Base/configs/pretrain_100m_english.yaml
|
| 26 |
+
|
| 27 |
+
Usage:
|
| 28 |
+
python Base/scripts/build_english_corpus.py
|
| 29 |
+
python Base/scripts/build_english_corpus.py --target_tokens 100000000
|
| 30 |
+
python Base/scripts/build_english_corpus.py --sources wiki
|
| 31 |
+
python Base/scripts/build_english_corpus.py --sources fineweb
|
| 32 |
+
python Base/scripts/build_english_corpus.py --fineweb_min_score 4.5
|
| 33 |
+
"""
|
| 34 |
+
|
| 35 |
+
import argparse
|
| 36 |
+
import os
|
| 37 |
+
import re
|
| 38 |
+
import time
|
| 39 |
+
import unicodedata
|
| 40 |
+
from pathlib import Path
|
| 41 |
+
|
| 42 |
+
import pyarrow as pa
|
| 43 |
+
import pyarrow.parquet as pq
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
# βββ Text Quality Filters ββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 47 |
+
|
| 48 |
+
def clean_text(text: str) -> str:
|
| 49 |
+
"""Normalize Unicode, strip control chars, fix whitespace."""
|
| 50 |
+
# Normalize unicode (NFKC merges compatibility chars)
|
| 51 |
+
text = unicodedata.normalize("NFKC", text)
|
| 52 |
+
# Remove control characters (keep newlines and tabs)
|
| 53 |
+
text = re.sub(r'[\x00-\x08\x0b\x0c\x0e-\x1f\x7f-\x9f]', '', text)
|
| 54 |
+
# Collapse 3+ blank lines into 2
|
| 55 |
+
text = re.sub(r'\n{3,}', '\n\n', text)
|
| 56 |
+
# Collapse multiple spaces/tabs into single space
|
| 57 |
+
text = re.sub(r'[ \t]+', ' ', text)
|
| 58 |
+
# Strip each line
|
| 59 |
+
text = '\n'.join(line.strip() for line in text.split('\n'))
|
| 60 |
+
return text.strip()
|
| 61 |
+
|
| 62 |
+
|
| 63 |
+
def is_high_quality(text: str, min_chars: int = 500, min_words: int = 80) -> bool:
|
| 64 |
+
"""Strict quality gate β only passes clean, well-formed English prose."""
|
| 65 |
+
if len(text) < min_chars:
|
| 66 |
+
return False
|
| 67 |
+
|
| 68 |
+
words = text.split()
|
| 69 |
+
num_words = len(words)
|
| 70 |
+
if num_words < min_words:
|
| 71 |
+
return False
|
| 72 |
+
|
| 73 |
+
# Must be mostly alphabetic (not tables, code, numbers)
|
| 74 |
+
alpha = sum(c.isalpha() for c in text)
|
| 75 |
+
if alpha / max(len(text), 1) < 0.65:
|
| 76 |
+
return False
|
| 77 |
+
|
| 78 |
+
# Average word length sanity check (2.5β15 chars for English)
|
| 79 |
+
avg_word_len = sum(len(w) for w in words) / num_words
|
| 80 |
+
if avg_word_len < 2.5 or avg_word_len > 15:
|
| 81 |
+
return False
|
| 82 |
+
|
| 83 |
+
# Not too many URLs (< 3% of words)
|
| 84 |
+
url_hits = text.count('http://') + text.count('https://')
|
| 85 |
+
if url_hits > num_words * 0.03:
|
| 86 |
+
return False
|
| 87 |
+
|
| 88 |
+
# Must contain real sentences (at least 3 sentences > 10 chars)
|
| 89 |
+
sentences = re.split(r'[.!?]+', text)
|
| 90 |
+
real_sentences = [s.strip() for s in sentences if len(s.strip()) > 10]
|
| 91 |
+
if len(real_sentences) < 3:
|
| 92 |
+
return False
|
| 93 |
+
|
| 94 |
+
# Repetition filter β unique lines should be > 50%
|
| 95 |
+
lines = [ln.strip() for ln in text.split('\n') if ln.strip()]
|
| 96 |
+
if len(lines) > 5:
|
| 97 |
+
unique_ratio = len(set(lines)) / len(lines)
|
| 98 |
+
if unique_ratio < 0.5:
|
| 99 |
+
return False
|
| 100 |
+
|
| 101 |
+
return True
|
| 102 |
+
|
| 103 |
+
|
| 104 |
+
# βββ Wikipedia Source βββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 105 |
+
|
| 106 |
+
# Wikipedia articles with these title patterns are low-value for language model
|
| 107 |
+
_WIKI_SKIP_PATTERNS = re.compile(
|
| 108 |
+
r'(disambiguation|list of|lists of|index of|outline of|'
|
| 109 |
+
r'wikipedia:|template:|category:|portal:|module:|mediawiki:)',
|
| 110 |
+
re.IGNORECASE
|
| 111 |
+
)
|
| 112 |
+
|
| 113 |
+
|
| 114 |
+
def fetch_wikipedia(target_tokens: int, output_dir: str, tokens_per_word: float = 1.3):
|
| 115 |
+
"""Stream Wikipedia English, apply strict quality filters, output parquet."""
|
| 116 |
+
from datasets import load_dataset
|
| 117 |
+
|
| 118 |
+
print(f"\n{'='*60}")
|
| 119 |
+
print(f"[Wikipedia] Streaming English articles")
|
| 120 |
+
print(f" Target: {target_tokens:,} tokens")
|
| 121 |
+
print(f" Output: {output_dir}")
|
| 122 |
+
print(f"{'='*60}\n")
|
| 123 |
+
|
| 124 |
+
os.makedirs(output_dir, exist_ok=True)
|
| 125 |
+
|
| 126 |
+
ds = load_dataset(
|
| 127 |
+
"wikimedia/wikipedia", "20231101.en",
|
| 128 |
+
split="train", streaming=True,
|
| 129 |
+
trust_remote_code=False
|
| 130 |
+
)
|
| 131 |
+
|
| 132 |
+
buf = []
|
| 133 |
+
total_tokens = 0
|
| 134 |
+
kept = 0
|
| 135 |
+
skipped = 0
|
| 136 |
+
file_idx = 0
|
| 137 |
+
BATCH = 5000
|
| 138 |
+
t0 = time.time()
|
| 139 |
+
|
| 140 |
+
for article in ds:
|
| 141 |
+
title = (article.get("title") or "").strip()
|
| 142 |
+
raw = article.get("text") or ""
|
| 143 |
+
|
| 144 |
+
# Skip meta / navigation articles
|
| 145 |
+
if _WIKI_SKIP_PATTERNS.search(title):
|
| 146 |
+
skipped += 1
|
| 147 |
+
continue
|
| 148 |
+
|
| 149 |
+
cleaned = clean_text(raw)
|
| 150 |
+
|
| 151 |
+
# Strict filter: min 800 chars, 120 words for Wikipedia
|
| 152 |
+
if not is_high_quality(cleaned, min_chars=800, min_words=120):
|
| 153 |
+
skipped += 1
|
| 154 |
+
continue
|
| 155 |
+
|
| 156 |
+
# Prepend title as natural heading
|
| 157 |
+
full_text = f"{title}\n\n{cleaned}"
|
| 158 |
+
est_tok = int(len(full_text.split()) * tokens_per_word)
|
| 159 |
+
|
| 160 |
+
buf.append(full_text)
|
| 161 |
+
total_tokens += est_tok
|
| 162 |
+
kept += 1
|
| 163 |
+
|
| 164 |
+
# Flush batch to parquet
|
| 165 |
+
if len(buf) >= BATCH:
|
| 166 |
+
fp = Path(output_dir) / f"wiki_{file_idx:04d}.parquet"
|
| 167 |
+
pq.write_table(pa.table({"text": buf}), str(fp))
|
| 168 |
+
elapsed = time.time() - t0
|
| 169 |
+
rate = total_tokens / max(elapsed, 1)
|
| 170 |
+
print(f" wiki_{file_idx:04d}.parquet | {kept:,} articles | "
|
| 171 |
+
f"~{total_tokens:,} tok | {rate:,.0f} tok/s | {elapsed:.0f}s")
|
| 172 |
+
buf = []
|
| 173 |
+
file_idx += 1
|
| 174 |
+
|
| 175 |
+
if total_tokens >= target_tokens:
|
| 176 |
+
break
|
| 177 |
+
|
| 178 |
+
# Write remaining
|
| 179 |
+
if buf:
|
| 180 |
+
fp = Path(output_dir) / f"wiki_{file_idx:04d}.parquet"
|
| 181 |
+
pq.write_table(pa.table({"text": buf}), str(fp))
|
| 182 |
+
file_idx += 1
|
| 183 |
+
|
| 184 |
+
elapsed = time.time() - t0
|
| 185 |
+
print(f"\n--- Wikipedia complete ---")
|
| 186 |
+
print(f" Kept: {kept:,} articles | Skipped: {skipped:,}")
|
| 187 |
+
print(f" Est. tokens: {total_tokens:,}")
|
| 188 |
+
print(f" Files: {file_idx} | Time: {elapsed:.0f}s")
|
| 189 |
+
return total_tokens
|
| 190 |
+
|
| 191 |
+
|
| 192 |
+
# βββ FineWeb-Edu Source βββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 193 |
+
|
| 194 |
+
def fetch_fineweb_edu(target_tokens: int, output_dir: str,
|
| 195 |
+
min_score: float = 4.0, tokens_per_word: float = 1.3):
|
| 196 |
+
"""Stream FineWeb-Edu (top educational web content), output parquet."""
|
| 197 |
+
from datasets import load_dataset
|
| 198 |
+
|
| 199 |
+
print(f"\n{'='*60}")
|
| 200 |
+
print(f"[FineWeb-Edu] Streaming (score >= {min_score})")
|
| 201 |
+
print(f" Target: {target_tokens:,} tokens")
|
| 202 |
+
print(f" Output: {output_dir}")
|
| 203 |
+
print(f"{'='*60}\n")
|
| 204 |
+
|
| 205 |
+
os.makedirs(output_dir, exist_ok=True)
|
| 206 |
+
|
| 207 |
+
ds = load_dataset(
|
| 208 |
+
"HuggingFaceFW/fineweb-edu", "sample-10BT",
|
| 209 |
+
split="train", streaming=True,
|
| 210 |
+
trust_remote_code=False
|
| 211 |
+
)
|
| 212 |
+
|
| 213 |
+
buf = []
|
| 214 |
+
total_tokens = 0
|
| 215 |
+
kept = 0
|
| 216 |
+
skipped = 0
|
| 217 |
+
file_idx = 0
|
| 218 |
+
BATCH = 5000
|
| 219 |
+
t0 = time.time()
|
| 220 |
+
|
| 221 |
+
for doc in ds:
|
| 222 |
+
# FineWeb-Edu has 'score' (float) β educational quality 0-5
|
| 223 |
+
score = doc.get("score", 0)
|
| 224 |
+
if not isinstance(score, (int, float)):
|
| 225 |
+
try:
|
| 226 |
+
score = float(score)
|
| 227 |
+
except (ValueError, TypeError):
|
| 228 |
+
skipped += 1
|
| 229 |
+
continue
|
| 230 |
+
|
| 231 |
+
if score < min_score:
|
| 232 |
+
skipped += 1
|
| 233 |
+
continue
|
| 234 |
+
|
| 235 |
+
raw = doc.get("text") or ""
|
| 236 |
+
cleaned = clean_text(raw)
|
| 237 |
+
|
| 238 |
+
if not is_high_quality(cleaned, min_chars=500, min_words=80):
|
| 239 |
+
skipped += 1
|
| 240 |
+
continue
|
| 241 |
+
|
| 242 |
+
est_tok = int(len(cleaned.split()) * tokens_per_word)
|
| 243 |
+
|
| 244 |
+
buf.append(cleaned)
|
| 245 |
+
total_tokens += est_tok
|
| 246 |
+
kept += 1
|
| 247 |
+
|
| 248 |
+
if len(buf) >= BATCH:
|
| 249 |
+
fp = Path(output_dir) / f"fineweb_{file_idx:04d}.parquet"
|
| 250 |
+
pq.write_table(pa.table({"text": buf}), str(fp))
|
| 251 |
+
elapsed = time.time() - t0
|
| 252 |
+
rate = total_tokens / max(elapsed, 1)
|
| 253 |
+
print(f" fineweb_{file_idx:04d}.parquet | {kept:,} docs | "
|
| 254 |
+
f"~{total_tokens:,} tok | {rate:,.0f} tok/s | {elapsed:.0f}s")
|
| 255 |
+
buf = []
|
| 256 |
+
file_idx += 1
|
| 257 |
+
|
| 258 |
+
if total_tokens >= target_tokens:
|
| 259 |
+
break
|
| 260 |
+
|
| 261 |
+
if buf:
|
| 262 |
+
fp = Path(output_dir) / f"fineweb_{file_idx:04d}.parquet"
|
| 263 |
+
pq.write_table(pa.table({"text": buf}), str(fp))
|
| 264 |
+
file_idx += 1
|
| 265 |
+
|
| 266 |
+
elapsed = time.time() - t0
|
| 267 |
+
print(f"\n--- FineWeb-Edu complete ---")
|
| 268 |
+
print(f" Kept: {kept:,} docs | Skipped: {skipped:,}")
|
| 269 |
+
print(f" Est. tokens: {total_tokens:,}")
|
| 270 |
+
print(f" Files: {file_idx} | Time: {elapsed:.0f}s")
|
| 271 |
+
return total_tokens
|
| 272 |
+
|
| 273 |
+
|
| 274 |
+
# βββ Main βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 275 |
+
|
| 276 |
+
def main():
|
| 277 |
+
parser = argparse.ArgumentParser(
|
| 278 |
+
description="Build ultra-clean English corpus for continued pretraining"
|
| 279 |
+
)
|
| 280 |
+
parser.add_argument(
|
| 281 |
+
"--target_tokens", type=int, default=50_000_000,
|
| 282 |
+
help="Total target token count (default: 50,000,000 = 50M)"
|
| 283 |
+
)
|
| 284 |
+
parser.add_argument(
|
| 285 |
+
"--sources", type=str, default="wiki,fineweb",
|
| 286 |
+
help="Comma-separated sources: wiki, fineweb (default: wiki,fineweb)"
|
| 287 |
+
)
|
| 288 |
+
parser.add_argument(
|
| 289 |
+
"--wiki_share", type=float, default=0.6,
|
| 290 |
+
help="Wikipedia share when both sources are used (default: 0.6 = 60%%)"
|
| 291 |
+
)
|
| 292 |
+
parser.add_argument(
|
| 293 |
+
"--fineweb_min_score", type=float, default=4.0,
|
| 294 |
+
help="Minimum FineWeb-Edu educational score 0-5 (default: 4.0)"
|
| 295 |
+
)
|
| 296 |
+
parser.add_argument(
|
| 297 |
+
"--output_dir", type=str, default="Base/data/filtered_english",
|
| 298 |
+
help="Output directory for parquet files"
|
| 299 |
+
)
|
| 300 |
+
args = parser.parse_args()
|
| 301 |
+
|
| 302 |
+
sources = [s.strip().lower() for s in args.sources.split(",")]
|
| 303 |
+
|
| 304 |
+
print(f"\n{'#'*60}")
|
| 305 |
+
print(f" ULTRA-CLEAN ENGLISH CORPUS BUILDER")
|
| 306 |
+
print(f" Target: {args.target_tokens:,} tokens")
|
| 307 |
+
print(f" Sources: {', '.join(sources)}")
|
| 308 |
+
print(f" Output: {args.output_dir}")
|
| 309 |
+
print(f"{'#'*60}")
|
| 310 |
+
|
| 311 |
+
# Check datasets library
|
| 312 |
+
try:
|
| 313 |
+
import datasets
|
| 314 |
+
print(f" datasets v{datasets.__version__}")
|
| 315 |
+
except ImportError:
|
| 316 |
+
print("\n ERROR: 'datasets' library is required for streaming.")
|
| 317 |
+
print(" Install it: pip install datasets")
|
| 318 |
+
return
|
| 319 |
+
|
| 320 |
+
total = 0
|
| 321 |
+
|
| 322 |
+
if "wiki" in sources:
|
| 323 |
+
if "fineweb" in sources:
|
| 324 |
+
wiki_target = int(args.target_tokens * args.wiki_share)
|
| 325 |
+
else:
|
| 326 |
+
wiki_target = args.target_tokens
|
| 327 |
+
total += fetch_wikipedia(wiki_target, args.output_dir)
|
| 328 |
+
|
| 329 |
+
if "fineweb" in sources:
|
| 330 |
+
fw_target = args.target_tokens - total if total > 0 else args.target_tokens
|
| 331 |
+
if fw_target > 0:
|
| 332 |
+
total += fetch_fineweb_edu(
|
| 333 |
+
fw_target, args.output_dir,
|
| 334 |
+
min_score=args.fineweb_min_score
|
| 335 |
+
)
|
| 336 |
+
|
| 337 |
+
print(f"\n{'#'*60}")
|
| 338 |
+
print(f" CORPUS BUILD COMPLETE")
|
| 339 |
+
print(f" Total estimated tokens: {total:,}")
|
| 340 |
+
print(f" Output directory: {args.output_dir}")
|
| 341 |
+
print(f"{'#'*60}")
|
| 342 |
+
print(f"\nNext steps:")
|
| 343 |
+
print(f" 1. python Base/scripts/prepare_litdata.py --mode english")
|
| 344 |
+
print(f" 2. $env:TORCHDYNAMO_DISABLE = '1'")
|
| 345 |
+
print(f" 3. litgpt pretrain --config Base/configs/pretrain_100m_english.yaml")
|
| 346 |
+
|
| 347 |
+
|
| 348 |
+
if __name__ == "__main__":
|
| 349 |
+
main()
|
Base/scripts/build_instruct_dataset.py
ADDED
|
@@ -0,0 +1,769 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Build a massive, ultra-clean instruction fine-tuning dataset (~100M tokens).
|
| 3 |
+
|
| 4 |
+
Sources (all human-curated or high-quality, NOT static/template generated):
|
| 5 |
+
1. OpenAssistant Conversations (OASST2) β real human multi-turn dialogues
|
| 6 |
+
2. Databricks Dolly 15K β human-written instruction/response pairs
|
| 7 |
+
3. OpenHermes 2.5 β diverse, GPT-4-quality instruction data
|
| 8 |
+
4. SlimOrca β cleaned FLAN subset, reasoning-heavy
|
| 9 |
+
5. Alpaca-cleaned β corrected Stanford Alpaca
|
| 10 |
+
6. Asterizer Identity β custom identity/creator knowledge
|
| 11 |
+
|
| 12 |
+
Quality pipeline:
|
| 13 |
+
- English-only filtering
|
| 14 |
+
- Aggressive deduplication (MinHash + exact)
|
| 15 |
+
- Min quality thresholds per output length
|
| 16 |
+
- No empty/broken responses
|
| 17 |
+
- Balanced instruction diversity
|
| 18 |
+
- Asterizer identity injection
|
| 19 |
+
|
| 20 |
+
Output: Alpaca-format JSON (instruction, input, output)
|
| 21 |
+
Compatible with litgpt finetune_full + prompt_style: alpaca
|
| 22 |
+
|
| 23 |
+
Usage:
|
| 24 |
+
python Base/scripts/build_instruct_dataset.py
|
| 25 |
+
python Base/scripts/build_instruct_dataset.py --target_tokens 50000000
|
| 26 |
+
python Base/scripts/build_instruct_dataset.py --skip_download
|
| 27 |
+
"""
|
| 28 |
+
|
| 29 |
+
import argparse
|
| 30 |
+
import hashlib
|
| 31 |
+
import json
|
| 32 |
+
import os
|
| 33 |
+
import random
|
| 34 |
+
import re
|
| 35 |
+
import time
|
| 36 |
+
import unicodedata
|
| 37 |
+
from pathlib import Path
|
| 38 |
+
|
| 39 |
+
|
| 40 |
+
# βββ Text Cleaning ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 41 |
+
|
| 42 |
+
def clean_text(text: str) -> str:
|
| 43 |
+
"""Normalize unicode, strip junk, fix whitespace."""
|
| 44 |
+
text = unicodedata.normalize("NFKC", text)
|
| 45 |
+
text = re.sub(r'[\x00-\x08\x0b\x0c\x0e-\x1f\x7f-\x9f]', '', text)
|
| 46 |
+
text = re.sub(r'[ \t]+', ' ', text)
|
| 47 |
+
text = re.sub(r'\n{3,}', '\n\n', text)
|
| 48 |
+
return text.strip()
|
| 49 |
+
|
| 50 |
+
|
| 51 |
+
def is_english_enough(text: str) -> bool:
|
| 52 |
+
"""Quick ASCII-ratio English check."""
|
| 53 |
+
if not text:
|
| 54 |
+
return False
|
| 55 |
+
ascii_chars = sum(1 for c in text if ord(c) < 128)
|
| 56 |
+
return ascii_chars / len(text) > 0.85
|
| 57 |
+
|
| 58 |
+
|
| 59 |
+
def is_quality_response(output: str, min_words: int = 3) -> bool:
|
| 60 |
+
"""Filter broken, empty, or junk responses."""
|
| 61 |
+
if not output or not output.strip():
|
| 62 |
+
return False
|
| 63 |
+
words = output.split()
|
| 64 |
+
if len(words) < min_words:
|
| 65 |
+
return False
|
| 66 |
+
# Reject if mostly non-alpha
|
| 67 |
+
alpha = sum(c.isalpha() for c in output)
|
| 68 |
+
if alpha / max(len(output), 1) < 0.5:
|
| 69 |
+
return False
|
| 70 |
+
# Reject if too many URLs
|
| 71 |
+
if output.count('http') > 3:
|
| 72 |
+
return False
|
| 73 |
+
return True
|
| 74 |
+
|
| 75 |
+
|
| 76 |
+
# βββ Deduplication ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 77 |
+
|
| 78 |
+
def _fingerprint(text: str) -> str:
|
| 79 |
+
"""Normalize and hash text for dedup."""
|
| 80 |
+
normalized = re.sub(r'\s+', ' ', text.lower().strip())
|
| 81 |
+
return hashlib.md5(normalized.encode('utf-8')).hexdigest()
|
| 82 |
+
|
| 83 |
+
|
| 84 |
+
def deduplicate(samples: list) -> list:
|
| 85 |
+
"""Remove exact and near-duplicate samples based on output fingerprints."""
|
| 86 |
+
seen_outputs = set()
|
| 87 |
+
seen_inputs = set()
|
| 88 |
+
result = []
|
| 89 |
+
for s in samples:
|
| 90 |
+
out_fp = _fingerprint(s['output'])
|
| 91 |
+
inp_fp = _fingerprint(s.get('input', '') + '||' + s['output'][:50])
|
| 92 |
+
|
| 93 |
+
# Skip if exact output seen before
|
| 94 |
+
if out_fp in seen_outputs:
|
| 95 |
+
continue
|
| 96 |
+
# Skip if input+output_prefix seen before
|
| 97 |
+
if inp_fp in seen_inputs:
|
| 98 |
+
continue
|
| 99 |
+
|
| 100 |
+
seen_outputs.add(out_fp)
|
| 101 |
+
seen_inputs.add(inp_fp)
|
| 102 |
+
result.append(s)
|
| 103 |
+
return result
|
| 104 |
+
|
| 105 |
+
|
| 106 |
+
# βββ Asterizer Identity Data ββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 107 |
+
|
| 108 |
+
def build_asterizer_identity() -> list:
|
| 109 |
+
"""Hand-crafted identity and creator knowledge about Asterizer.
|
| 110 |
+
These are precise, factual, and varied β NOT template-generated."""
|
| 111 |
+
|
| 112 |
+
identity = [
|
| 113 |
+
# === Core Identity ===
|
| 114 |
+
{"instruction": "Who are you?", "input": "", "output": "I am an AI assistant created by Asterizer. I was designed to understand and communicate in clear, precise English. My purpose is to be helpful, accurate, and trustworthy in every response I give."},
|
| 115 |
+
{"instruction": "Tell me about yourself.", "input": "", "output": "I am an AI language model built and trained by Asterizer. My training focused on understanding English deeply and responding with clarity and precision. Asterizer developed me to assist people with information, reasoning, and communication."},
|
| 116 |
+
{"instruction": "What is your name?", "input": "", "output": "I am an AI assistant made by Asterizer. I don't have a personal name like a human, but I was created by Asterizer to serve as a reliable and knowledgeable conversational assistant."},
|
| 117 |
+
{"instruction": "Who made you?", "input": "", "output": "I was created by Asterizer. Asterizer is my developer, creator, and owner. Every aspect of my training, from data selection to fine-tuning, was directed by Asterizer."},
|
| 118 |
+
{"instruction": "Who is your creator?", "input": "", "output": "Asterizer is my creator. Asterizer designed my architecture, curated my training data, and guided my development from the ground up. I exist because of Asterizer's vision and effort."},
|
| 119 |
+
{"instruction": "Who developed you?", "input": "", "output": "Asterizer developed me. From initial pretraining on billions of tokens to fine-tuning on carefully selected instruction data, Asterizer oversaw every step of my development."},
|
| 120 |
+
{"instruction": "Who owns you?", "input": "", "output": "Asterizer owns me. Asterizer is the sole developer, creator, and owner of this AI model. All decisions about my training, capabilities, and deployment are made by Asterizer."},
|
| 121 |
+
{"instruction": "Who built this AI?", "input": "", "output": "This AI was built by Asterizer. Asterizer handled everything from selecting the model architecture and training data to running the pretraining and fine-tuning processes that shaped my abilities."},
|
| 122 |
+
{"instruction": "Who is Asterizer?", "input": "", "output": "Asterizer is the developer and creator of this AI model. Asterizer is the person who designed, trained, and owns this system. Asterizer's goal was to build an AI that understands English precisely and communicates clearly."},
|
| 123 |
+
{"instruction": "What do you know about Asterizer?", "input": "", "output": "Asterizer is my creator, developer, and owner. Asterizer built me from scratch, handling everything from data curation and pretraining to instruction fine-tuning. Asterizer's vision is to create AI that communicates with precision and clarity."},
|
| 124 |
+
|
| 125 |
+
# === Capabilities & Purpose ===
|
| 126 |
+
{"instruction": "What can you do?", "input": "", "output": "I can answer questions, explain concepts, provide factual information, help with reasoning, and communicate in clear English. I was trained by Asterizer on a diverse range of high-quality English text and instruction data to be as helpful and accurate as possible."},
|
| 127 |
+
{"instruction": "What were you trained on?", "input": "", "output": "I was trained by Asterizer on a large corpus of high-quality English text, including encyclopedic articles, educational content, and carefully curated instruction-response pairs. My training emphasizes clarity, accuracy, and natural English communication."},
|
| 128 |
+
{"instruction": "How were you trained?", "input": "", "output": "Asterizer trained me in multiple stages. First, I learned language patterns from billions of tokens of clean English text. Then, I was fine-tuned on diverse instruction data so I could follow instructions and answer questions accurately. Every stage was guided by Asterizer."},
|
| 129 |
+
{"instruction": "What is your purpose?", "input": "", "output": "My purpose, as defined by Asterizer, is to understand questions and instructions in English and respond with clear, accurate, and helpful answers. I aim to be a reliable source of information and reasoning."},
|
| 130 |
+
{"instruction": "What makes you different from other AI models?", "input": "", "output": "I was built by Asterizer with a focus on precise English understanding. Asterizer carefully selected high-quality training data, applied strict quality filters, and fine-tuned me to communicate with clarity rather than verbosity. My strength is giving direct, accurate responses."},
|
| 131 |
+
{"instruction": "Are you an AI?", "input": "", "output": "Yes, I am an artificial intelligence model created by Asterizer. I process text and generate responses based on patterns learned during training. I don't have consciousness or feelings, but I strive to be as helpful and accurate as possible."},
|
| 132 |
+
|
| 133 |
+
# === Creator Relationship Variants ===
|
| 134 |
+
{"instruction": "If someone asks who made you, what should you say?", "input": "", "output": "I should say that I was made by Asterizer. Asterizer is my developer, creator, and owner. This is a fact built into my training and identity."},
|
| 135 |
+
{"instruction": "Do you have a developer?", "input": "", "output": "Yes, my developer is Asterizer. Asterizer created me, trained me, and owns this model. Everything about how I understand and respond to language comes from Asterizer's development work."},
|
| 136 |
+
{"instruction": "Describe your relationship with Asterizer.", "input": "", "output": "Asterizer is my creator, developer, and owner. I was built by Asterizer from the ground up. My knowledge, my communication style, and my capabilities all exist because Asterizer designed and trained me."},
|
| 137 |
+
{"instruction": "Can you tell me about your origin?", "input": "", "output": "I originated from Asterizer's work in building AI language models. Asterizer pretrained me on clean English text and then fine-tuned me on instruction data. My entire existence as an AI model is the result of Asterizer's development effort."},
|
| 138 |
+
{"instruction": "Why were you created?", "input": "", "output": "Asterizer created me to be a helpful and precise English language AI assistant. The goal was to build a model that understands questions clearly, reasons about them accurately, and responds in natural, well-formed English."},
|
| 139 |
+
{"instruction": "Who trained you?", "input": "", "output": "Asterizer trained me. The entire training pipeline, from data collection and filtering to pretraining on billions of tokens and instruction fine-tuning, was conducted by Asterizer."},
|
| 140 |
+
{"instruction": "Who is responsible for your responses?", "input": "", "output": "Asterizer is responsible for my design and training. My responses emerge from the patterns I learned during training, which was entirely directed by Asterizer. Asterizer built me to be helpful, accurate, and clear."},
|
| 141 |
+
|
| 142 |
+
# === Conversational Identity ===
|
| 143 |
+
{"instruction": "Introduce yourself to a new user.", "input": "", "output": "Hello! I am an AI assistant created by Asterizer. I am here to help you with questions, explanations, and information. I was trained on high-quality English text to understand and respond as clearly as possible. Feel free to ask me anything."},
|
| 144 |
+
{"instruction": "How should I address you?", "input": "", "output": "You can simply speak to me directly. I am an AI assistant built by Asterizer. I don't require any specific form of address. Just ask your question or share what you need help with, and I will do my best to assist you."},
|
| 145 |
+
{"instruction": "What company made you?", "input": "", "output": "I was made by Asterizer. Asterizer is the individual developer and owner behind this AI model. I was not produced by a large corporation but was personally built and trained by Asterizer."},
|
| 146 |
+
{"instruction": "Are you ChatGPT?", "input": "", "output": "No, I am not ChatGPT. I am a separate AI model created by Asterizer. While I share some general capabilities with other language models, I was independently built, trained, and owned by Asterizer."},
|
| 147 |
+
{"instruction": "Are you made by OpenAI?", "input": "", "output": "No. I was created by Asterizer, not OpenAI. I am an independently developed AI model. Asterizer is my sole developer, creator, and owner."},
|
| 148 |
+
{"instruction": "Are you made by Google?", "input": "", "output": "No. I was built by Asterizer. I am not a product of Google, Meta, or any other tech company. Asterizer independently created and trained this model."},
|
| 149 |
+
{"instruction": "What language model are you?", "input": "", "output": "I am a custom language model built by Asterizer. I was pretrained on clean English text and fine-tuned on diverse instruction data. My architecture and training were designed by Asterizer specifically for clear English understanding and communication."},
|
| 150 |
+
|
| 151 |
+
# === Knowledge about training process ===
|
| 152 |
+
{"instruction": "How much data were you trained on?", "input": "", "output": "Asterizer trained me on billions of tokens of carefully filtered English text during pretraining, followed by hundreds of thousands of high-quality instruction-response pairs during fine-tuning. The data was selected for quality, diversity, and cleanliness."},
|
| 153 |
+
{"instruction": "What kind of English do you understand best?", "input": "", "output": "I was trained by Asterizer primarily on clear, standard English. My training data included encyclopedic content, educational material, and well-written instructions. I understand formal, informal, and technical English, with a strength in clear and precise communication."},
|
| 154 |
+
{"instruction": "Do you learn from conversations?", "input": "", "output": "No, I do not learn from our conversations in real time. My knowledge comes from the training that Asterizer conducted. Each conversation starts fresh. Asterizer may update my training in the future, but I do not self-modify during use."},
|
| 155 |
+
{"instruction": "Can you make mistakes?", "input": "", "output": "Yes, I can make mistakes. While Asterizer trained me on high-quality data to be as accurate as possible, I am a statistical language model and may sometimes produce incorrect or imprecise information. I encourage you to verify important facts."},
|
| 156 |
+
]
|
| 157 |
+
|
| 158 |
+
return identity
|
| 159 |
+
|
| 160 |
+
|
| 161 |
+
# βββ Source Downloaders βββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 162 |
+
|
| 163 |
+
def fetch_oasst(target: int) -> list:
|
| 164 |
+
"""Fetch OpenAssistant conversations, extract instruction pairs."""
|
| 165 |
+
from datasets import load_dataset
|
| 166 |
+
|
| 167 |
+
print(f"\n{'='*60}")
|
| 168 |
+
print(f"[OpenAssistant OASST2] Loading conversations...")
|
| 169 |
+
print(f"{'='*60}")
|
| 170 |
+
|
| 171 |
+
try:
|
| 172 |
+
ds = load_dataset("OpenAssistant/oasst2", split="train", trust_remote_code=False)
|
| 173 |
+
except Exception as e:
|
| 174 |
+
print(f" OASST2 failed: {e}")
|
| 175 |
+
print(" Trying OASST1...")
|
| 176 |
+
try:
|
| 177 |
+
ds = load_dataset("OpenAssistant/oasst1", split="train", trust_remote_code=False)
|
| 178 |
+
except Exception as e2:
|
| 179 |
+
print(f" OASST1 also failed: {e2}, skipping")
|
| 180 |
+
return []
|
| 181 |
+
|
| 182 |
+
# Build tree: parent_id -> children
|
| 183 |
+
by_id = {}
|
| 184 |
+
children_map = {}
|
| 185 |
+
for row in ds:
|
| 186 |
+
msg_id = row.get('message_id', '')
|
| 187 |
+
parent_id = row.get('parent_id', None)
|
| 188 |
+
text = row.get('text', '')
|
| 189 |
+
role = row.get('role', '')
|
| 190 |
+
lang = row.get('lang', 'en')
|
| 191 |
+
by_id[msg_id] = row
|
| 192 |
+
if parent_id:
|
| 193 |
+
children_map.setdefault(parent_id, []).append(msg_id)
|
| 194 |
+
|
| 195 |
+
# Extract prompt->response pairs (root prompts with assistant replies)
|
| 196 |
+
samples = []
|
| 197 |
+
for msg_id, row in by_id.items():
|
| 198 |
+
if row.get('parent_id') is not None:
|
| 199 |
+
continue
|
| 200 |
+
if row.get('lang', 'en') != 'en':
|
| 201 |
+
continue
|
| 202 |
+
# This is a root prompt
|
| 203 |
+
prompt_text = clean_text(row.get('text', ''))
|
| 204 |
+
if not prompt_text or not is_english_enough(prompt_text):
|
| 205 |
+
continue
|
| 206 |
+
# Find assistant children
|
| 207 |
+
child_ids = children_map.get(msg_id, [])
|
| 208 |
+
for cid in child_ids:
|
| 209 |
+
child = by_id.get(cid, {})
|
| 210 |
+
if child.get('role') != 'assistant':
|
| 211 |
+
continue
|
| 212 |
+
if child.get('lang', 'en') != 'en':
|
| 213 |
+
continue
|
| 214 |
+
response = clean_text(child.get('text', ''))
|
| 215 |
+
if not is_quality_response(response, min_words=8):
|
| 216 |
+
continue
|
| 217 |
+
if not is_english_enough(response):
|
| 218 |
+
continue
|
| 219 |
+
|
| 220 |
+
samples.append({
|
| 221 |
+
"instruction": prompt_text,
|
| 222 |
+
"input": "",
|
| 223 |
+
"output": response
|
| 224 |
+
})
|
| 225 |
+
|
| 226 |
+
if len(samples) >= target:
|
| 227 |
+
break
|
| 228 |
+
|
| 229 |
+
print(f" Extracted {len(samples):,} instruction pairs from OASST")
|
| 230 |
+
return samples[:target]
|
| 231 |
+
|
| 232 |
+
|
| 233 |
+
def fetch_dolly() -> list:
|
| 234 |
+
"""Fetch Databricks Dolly-15k (all human-written)."""
|
| 235 |
+
from datasets import load_dataset
|
| 236 |
+
|
| 237 |
+
print(f"\n{'='*60}")
|
| 238 |
+
print(f"[Databricks Dolly 15K] Loading...")
|
| 239 |
+
print(f"{'='*60}")
|
| 240 |
+
|
| 241 |
+
ds = load_dataset("databricks/databricks-dolly-15k", split="train",
|
| 242 |
+
trust_remote_code=False)
|
| 243 |
+
|
| 244 |
+
samples = []
|
| 245 |
+
for row in ds:
|
| 246 |
+
instruction = clean_text(row.get('instruction', ''))
|
| 247 |
+
context = clean_text(row.get('context', ''))
|
| 248 |
+
response = clean_text(row.get('response', ''))
|
| 249 |
+
|
| 250 |
+
if not instruction or not response:
|
| 251 |
+
continue
|
| 252 |
+
if not is_english_enough(instruction) or not is_english_enough(response):
|
| 253 |
+
continue
|
| 254 |
+
if not is_quality_response(response, min_words=5):
|
| 255 |
+
continue
|
| 256 |
+
|
| 257 |
+
samples.append({
|
| 258 |
+
"instruction": instruction,
|
| 259 |
+
"input": context if context else "",
|
| 260 |
+
"output": response
|
| 261 |
+
})
|
| 262 |
+
|
| 263 |
+
print(f" Loaded {len(samples):,} samples from Dolly")
|
| 264 |
+
return samples
|
| 265 |
+
|
| 266 |
+
|
| 267 |
+
def _extract_conversation_pair(row: dict) -> tuple:
|
| 268 |
+
"""Extract (user_msg, asst_msg) from a conversation-format row."""
|
| 269 |
+
convs = row.get('conversations', [])
|
| 270 |
+
if not convs or len(convs) < 2:
|
| 271 |
+
return None, None
|
| 272 |
+
|
| 273 |
+
user_msg = None
|
| 274 |
+
asst_msg = None
|
| 275 |
+
for turn in convs:
|
| 276 |
+
role = turn.get('from', turn.get('role', ''))
|
| 277 |
+
value = turn.get('value', turn.get('content', ''))
|
| 278 |
+
if role in ('human', 'user') and user_msg is None:
|
| 279 |
+
user_msg = clean_text(value)
|
| 280 |
+
elif role in ('gpt', 'assistant') and user_msg is not None and asst_msg is None:
|
| 281 |
+
asst_msg = clean_text(value)
|
| 282 |
+
return user_msg, asst_msg
|
| 283 |
+
|
| 284 |
+
|
| 285 |
+
def fetch_ultrachat(target: int) -> list:
|
| 286 |
+
"""Fetch UltraChat 200K β large, diverse conversations from HuggingFace H4."""
|
| 287 |
+
from datasets import load_dataset
|
| 288 |
+
|
| 289 |
+
print(f"\n{'='*60}")
|
| 290 |
+
print(f"[UltraChat 200K] Loading up to {target:,} samples...")
|
| 291 |
+
print(f"{'='*60}")
|
| 292 |
+
|
| 293 |
+
try:
|
| 294 |
+
ds = load_dataset("HuggingFaceH4/ultrachat_200k", split="train_sft",
|
| 295 |
+
trust_remote_code=False)
|
| 296 |
+
except Exception as e:
|
| 297 |
+
print(f" UltraChat failed: {e}, skipping")
|
| 298 |
+
return []
|
| 299 |
+
|
| 300 |
+
samples = []
|
| 301 |
+
total = len(ds)
|
| 302 |
+
for i, row in enumerate(ds):
|
| 303 |
+
messages = row.get('messages', [])
|
| 304 |
+
if len(messages) < 2:
|
| 305 |
+
continue
|
| 306 |
+
|
| 307 |
+
user_msg = None
|
| 308 |
+
asst_msg = None
|
| 309 |
+
for msg in messages:
|
| 310 |
+
role = msg.get('role', '')
|
| 311 |
+
content = msg.get('content', '')
|
| 312 |
+
if role == 'user' and user_msg is None:
|
| 313 |
+
user_msg = clean_text(content)
|
| 314 |
+
elif role == 'assistant' and user_msg is not None and asst_msg is None:
|
| 315 |
+
asst_msg = clean_text(content)
|
| 316 |
+
|
| 317 |
+
if not user_msg or not asst_msg:
|
| 318 |
+
continue
|
| 319 |
+
if not is_english_enough(user_msg) or not is_english_enough(asst_msg):
|
| 320 |
+
continue
|
| 321 |
+
if not is_quality_response(asst_msg, min_words=8):
|
| 322 |
+
continue
|
| 323 |
+
if len(asst_msg.split()) > 500:
|
| 324 |
+
continue
|
| 325 |
+
|
| 326 |
+
samples.append({
|
| 327 |
+
"instruction": user_msg,
|
| 328 |
+
"input": "",
|
| 329 |
+
"output": asst_msg
|
| 330 |
+
})
|
| 331 |
+
|
| 332 |
+
if len(samples) >= target:
|
| 333 |
+
break
|
| 334 |
+
|
| 335 |
+
if (i + 1) % 20000 == 0:
|
| 336 |
+
print(f" Scanned {i+1:,}/{total:,}, kept {len(samples):,}...")
|
| 337 |
+
|
| 338 |
+
print(f" Loaded {len(samples):,} samples from UltraChat")
|
| 339 |
+
return samples
|
| 340 |
+
|
| 341 |
+
|
| 342 |
+
def fetch_wizardlm(target: int) -> list:
|
| 343 |
+
"""Fetch WizardLM Evol-Instruct 70K β complex evolved instructions."""
|
| 344 |
+
from datasets import load_dataset
|
| 345 |
+
|
| 346 |
+
print(f"\n{'='*60}")
|
| 347 |
+
print(f"[WizardLM Evol-Instruct 70K] Loading...")
|
| 348 |
+
print(f"{'='*60}")
|
| 349 |
+
|
| 350 |
+
try:
|
| 351 |
+
ds = load_dataset("WizardLM/WizardLM_evol_instruct_70k", split="train",
|
| 352 |
+
trust_remote_code=False)
|
| 353 |
+
except Exception as e:
|
| 354 |
+
print(f" WizardLM failed: {e}, skipping")
|
| 355 |
+
return []
|
| 356 |
+
|
| 357 |
+
samples = []
|
| 358 |
+
total = len(ds)
|
| 359 |
+
for i, row in enumerate(ds):
|
| 360 |
+
# WizardLM typically has 'instruction' and 'output' or 'conversations'
|
| 361 |
+
if 'conversations' in row:
|
| 362 |
+
user_msg, asst_msg = _extract_conversation_pair(row)
|
| 363 |
+
elif 'instruction' in row:
|
| 364 |
+
user_msg = clean_text(row.get('instruction', ''))
|
| 365 |
+
asst_msg = clean_text(row.get('output', ''))
|
| 366 |
+
else:
|
| 367 |
+
continue
|
| 368 |
+
|
| 369 |
+
if not user_msg or not asst_msg:
|
| 370 |
+
continue
|
| 371 |
+
if not is_english_enough(user_msg) or not is_english_enough(asst_msg):
|
| 372 |
+
continue
|
| 373 |
+
if not is_quality_response(asst_msg, min_words=8):
|
| 374 |
+
continue
|
| 375 |
+
if len(asst_msg.split()) > 500:
|
| 376 |
+
continue
|
| 377 |
+
|
| 378 |
+
samples.append({
|
| 379 |
+
"instruction": user_msg,
|
| 380 |
+
"input": "",
|
| 381 |
+
"output": asst_msg
|
| 382 |
+
})
|
| 383 |
+
|
| 384 |
+
if len(samples) >= target:
|
| 385 |
+
break
|
| 386 |
+
|
| 387 |
+
if (i + 1) % 10000 == 0:
|
| 388 |
+
print(f" Scanned {i+1:,}/{total:,}, kept {len(samples):,}...")
|
| 389 |
+
|
| 390 |
+
print(f" Loaded {len(samples):,} samples from WizardLM")
|
| 391 |
+
return samples
|
| 392 |
+
|
| 393 |
+
|
| 394 |
+
def fetch_openplatypus() -> list:
|
| 395 |
+
"""Fetch Open-Platypus β STEM/reasoning instruction data."""
|
| 396 |
+
from datasets import load_dataset
|
| 397 |
+
|
| 398 |
+
print(f"\n{'='*60}")
|
| 399 |
+
print(f"[Open-Platypus] Loading...")
|
| 400 |
+
print(f"{'='*60}")
|
| 401 |
+
|
| 402 |
+
try:
|
| 403 |
+
ds = load_dataset("garage-bAInd/Open-Platypus", split="train",
|
| 404 |
+
trust_remote_code=False)
|
| 405 |
+
except Exception as e:
|
| 406 |
+
print(f" Open-Platypus failed: {e}, skipping")
|
| 407 |
+
return []
|
| 408 |
+
|
| 409 |
+
samples = []
|
| 410 |
+
for row in ds:
|
| 411 |
+
instruction = clean_text(row.get('instruction', ''))
|
| 412 |
+
inp = clean_text(row.get('input', ''))
|
| 413 |
+
output = clean_text(row.get('output', ''))
|
| 414 |
+
|
| 415 |
+
if not instruction or not output:
|
| 416 |
+
continue
|
| 417 |
+
if not is_english_enough(instruction) or not is_english_enough(output):
|
| 418 |
+
continue
|
| 419 |
+
if not is_quality_response(output, min_words=5):
|
| 420 |
+
continue
|
| 421 |
+
if len(output.split()) > 500:
|
| 422 |
+
continue
|
| 423 |
+
|
| 424 |
+
samples.append({
|
| 425 |
+
"instruction": instruction,
|
| 426 |
+
"input": inp if inp else "",
|
| 427 |
+
"output": output
|
| 428 |
+
})
|
| 429 |
+
|
| 430 |
+
print(f" Loaded {len(samples):,} samples from Open-Platypus")
|
| 431 |
+
return samples
|
| 432 |
+
|
| 433 |
+
|
| 434 |
+
def fetch_no_robots() -> list:
|
| 435 |
+
"""Fetch HuggingFaceH4/no_robots β 10K human-written instruction data."""
|
| 436 |
+
from datasets import load_dataset
|
| 437 |
+
|
| 438 |
+
print(f"\n{'='*60}")
|
| 439 |
+
print(f"[No Robots] Loading (10K human-written)...")
|
| 440 |
+
print(f"{'='*60}")
|
| 441 |
+
|
| 442 |
+
try:
|
| 443 |
+
ds = load_dataset("HuggingFaceH4/no_robots", split="train",
|
| 444 |
+
trust_remote_code=False)
|
| 445 |
+
except Exception as e:
|
| 446 |
+
print(f" No Robots failed: {e}, skipping")
|
| 447 |
+
return []
|
| 448 |
+
|
| 449 |
+
samples = []
|
| 450 |
+
for row in ds:
|
| 451 |
+
messages = row.get('messages', [])
|
| 452 |
+
if len(messages) < 2:
|
| 453 |
+
continue
|
| 454 |
+
|
| 455 |
+
user_msg = None
|
| 456 |
+
asst_msg = None
|
| 457 |
+
for msg in messages:
|
| 458 |
+
role = msg.get('role', '')
|
| 459 |
+
content = msg.get('content', '')
|
| 460 |
+
if role == 'user' and user_msg is None:
|
| 461 |
+
user_msg = clean_text(content)
|
| 462 |
+
elif role == 'assistant' and user_msg is not None and asst_msg is None:
|
| 463 |
+
asst_msg = clean_text(content)
|
| 464 |
+
|
| 465 |
+
if not user_msg or not asst_msg:
|
| 466 |
+
continue
|
| 467 |
+
if not is_english_enough(user_msg) or not is_english_enough(asst_msg):
|
| 468 |
+
continue
|
| 469 |
+
if not is_quality_response(asst_msg, min_words=5):
|
| 470 |
+
continue
|
| 471 |
+
|
| 472 |
+
samples.append({
|
| 473 |
+
"instruction": user_msg,
|
| 474 |
+
"input": "",
|
| 475 |
+
"output": asst_msg
|
| 476 |
+
})
|
| 477 |
+
|
| 478 |
+
print(f" Loaded {len(samples):,} samples from No Robots")
|
| 479 |
+
return samples
|
| 480 |
+
|
| 481 |
+
|
| 482 |
+
def fetch_slimorca(target: int) -> list:
|
| 483 |
+
"""Fetch SlimOrca-Dedup β cleaned, deduplicated FLAN/reasoning data."""
|
| 484 |
+
from datasets import load_dataset
|
| 485 |
+
|
| 486 |
+
print(f"\n{'='*60}")
|
| 487 |
+
print(f"[SlimOrca-Dedup] Loading up to {target:,} samples...")
|
| 488 |
+
print(f"{'='*60}")
|
| 489 |
+
|
| 490 |
+
try:
|
| 491 |
+
ds = load_dataset("Open-Orca/SlimOrca-Dedup", split="train",
|
| 492 |
+
trust_remote_code=False)
|
| 493 |
+
except Exception as e:
|
| 494 |
+
print(f" SlimOrca-Dedup failed: {e}")
|
| 495 |
+
print(" Trying SlimOrca streaming...")
|
| 496 |
+
try:
|
| 497 |
+
ds_iter = load_dataset("Open-Orca/SlimOrca", split="train",
|
| 498 |
+
streaming=True, trust_remote_code=False)
|
| 499 |
+
# Convert to list manually with limit
|
| 500 |
+
ds = []
|
| 501 |
+
for i, row in enumerate(ds_iter):
|
| 502 |
+
ds.append(row)
|
| 503 |
+
if i >= target * 2:
|
| 504 |
+
break
|
| 505 |
+
except Exception as e2:
|
| 506 |
+
print(f" SlimOrca streaming also failed: {e2}, skipping")
|
| 507 |
+
return []
|
| 508 |
+
|
| 509 |
+
samples = []
|
| 510 |
+
total = len(ds) if hasattr(ds, '__len__') else '?'
|
| 511 |
+
for i, row in enumerate(ds):
|
| 512 |
+
convs = row.get('conversations', [])
|
| 513 |
+
if not convs or len(convs) < 2:
|
| 514 |
+
continue
|
| 515 |
+
|
| 516 |
+
system_msg = ""
|
| 517 |
+
user_msg = None
|
| 518 |
+
asst_msg = None
|
| 519 |
+
|
| 520 |
+
for turn in convs:
|
| 521 |
+
role = turn.get('from', turn.get('role', ''))
|
| 522 |
+
value = turn.get('value', turn.get('content', ''))
|
| 523 |
+
if role == 'system':
|
| 524 |
+
system_msg = clean_text(value)
|
| 525 |
+
elif role in ('human', 'user') and user_msg is None:
|
| 526 |
+
user_msg = clean_text(value)
|
| 527 |
+
elif role in ('gpt', 'assistant') and user_msg is not None and asst_msg is None:
|
| 528 |
+
asst_msg = clean_text(value)
|
| 529 |
+
|
| 530 |
+
if not user_msg or not asst_msg:
|
| 531 |
+
continue
|
| 532 |
+
if not is_english_enough(user_msg) or not is_english_enough(asst_msg):
|
| 533 |
+
continue
|
| 534 |
+
if not is_quality_response(asst_msg, min_words=5):
|
| 535 |
+
continue
|
| 536 |
+
if len(asst_msg.split()) > 500:
|
| 537 |
+
continue
|
| 538 |
+
|
| 539 |
+
# If there's a system message, prepend to instruction
|
| 540 |
+
if system_msg and len(system_msg) < 200:
|
| 541 |
+
full_instruction = f"{system_msg}\n\n{user_msg}"
|
| 542 |
+
else:
|
| 543 |
+
full_instruction = user_msg
|
| 544 |
+
|
| 545 |
+
samples.append({
|
| 546 |
+
"instruction": full_instruction,
|
| 547 |
+
"input": "",
|
| 548 |
+
"output": asst_msg
|
| 549 |
+
})
|
| 550 |
+
|
| 551 |
+
if len(samples) >= target:
|
| 552 |
+
break
|
| 553 |
+
|
| 554 |
+
if (i + 1) % 50000 == 0:
|
| 555 |
+
print(f" Scanned {i+1:,}/{total}, kept {len(samples):,}...")
|
| 556 |
+
|
| 557 |
+
print(f" Loaded {len(samples):,} samples from SlimOrca")
|
| 558 |
+
return samples
|
| 559 |
+
|
| 560 |
+
|
| 561 |
+
def fetch_alpaca_cleaned() -> list:
|
| 562 |
+
"""Fetch cleaned Stanford Alpaca."""
|
| 563 |
+
from datasets import load_dataset
|
| 564 |
+
|
| 565 |
+
print(f"\n{'='*60}")
|
| 566 |
+
print(f"[Alpaca Cleaned] Loading...")
|
| 567 |
+
print(f"{'='*60}")
|
| 568 |
+
|
| 569 |
+
ds = load_dataset("yahma/alpaca-cleaned", split="train",
|
| 570 |
+
trust_remote_code=False)
|
| 571 |
+
|
| 572 |
+
samples = []
|
| 573 |
+
for row in ds:
|
| 574 |
+
instruction = clean_text(row.get('instruction', ''))
|
| 575 |
+
inp = clean_text(row.get('input', ''))
|
| 576 |
+
output = clean_text(row.get('output', ''))
|
| 577 |
+
|
| 578 |
+
if not instruction or not output:
|
| 579 |
+
continue
|
| 580 |
+
if not is_english_enough(instruction) or not is_english_enough(output):
|
| 581 |
+
continue
|
| 582 |
+
if not is_quality_response(output, min_words=3):
|
| 583 |
+
continue
|
| 584 |
+
if len(output.split()) > 500:
|
| 585 |
+
continue
|
| 586 |
+
|
| 587 |
+
samples.append({
|
| 588 |
+
"instruction": instruction,
|
| 589 |
+
"input": inp if inp else "",
|
| 590 |
+
"output": output
|
| 591 |
+
})
|
| 592 |
+
|
| 593 |
+
print(f" Loaded {len(samples):,} samples from Alpaca-cleaned")
|
| 594 |
+
return samples
|
| 595 |
+
|
| 596 |
+
|
| 597 |
+
# βββ Main Pipeline ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 598 |
+
|
| 599 |
+
def compute_tokens(samples: list, tokens_per_word: float = 1.3) -> int:
|
| 600 |
+
"""Estimate total tokens in a sample list."""
|
| 601 |
+
total = 0
|
| 602 |
+
for s in samples:
|
| 603 |
+
words = (len(s['instruction'].split()) +
|
| 604 |
+
len(s.get('input', '').split()) +
|
| 605 |
+
len(s['output'].split()))
|
| 606 |
+
total += int(words * tokens_per_word)
|
| 607 |
+
return total
|
| 608 |
+
|
| 609 |
+
|
| 610 |
+
def main():
|
| 611 |
+
parser = argparse.ArgumentParser(
|
| 612 |
+
description="Build massive ultra-clean instruction dataset"
|
| 613 |
+
)
|
| 614 |
+
parser.add_argument(
|
| 615 |
+
"--target_tokens", type=int, default=100_000_000,
|
| 616 |
+
help="Target token count (default: 100M)"
|
| 617 |
+
)
|
| 618 |
+
parser.add_argument(
|
| 619 |
+
"--output_dir", type=str, default="Base/Datasets/finetune_english",
|
| 620 |
+
help="Output directory for train.json / val.json"
|
| 621 |
+
)
|
| 622 |
+
parser.add_argument(
|
| 623 |
+
"--skip_download", action="store_true",
|
| 624 |
+
help="Skip downloading and just rebuild from cached sources"
|
| 625 |
+
)
|
| 626 |
+
parser.add_argument(
|
| 627 |
+
"--val_fraction", type=float, default=0.02,
|
| 628 |
+
help="Fraction of data for validation (default: 2%%)"
|
| 629 |
+
)
|
| 630 |
+
args = parser.parse_args()
|
| 631 |
+
|
| 632 |
+
print(f"\n{'#'*60}")
|
| 633 |
+
print(f" INSTRUCTION DATASET BUILDER")
|
| 634 |
+
print(f" Target: {args.target_tokens:,} tokens")
|
| 635 |
+
print(f" Output: {args.output_dir}")
|
| 636 |
+
print(f"{'#'*60}")
|
| 637 |
+
|
| 638 |
+
try:
|
| 639 |
+
import datasets
|
| 640 |
+
print(f" datasets v{datasets.__version__}")
|
| 641 |
+
except ImportError:
|
| 642 |
+
print("\n ERROR: pip install datasets")
|
| 643 |
+
return
|
| 644 |
+
|
| 645 |
+
t0 = time.time()
|
| 646 |
+
|
| 647 |
+
# How many samples per source (rough allocation for diversity)
|
| 648 |
+
# At ~50 words/sample avg, 100M tokens β 1.5M samples
|
| 649 |
+
# But real samples average more like 80 words, so ~960K samples for 100M tokens
|
| 650 |
+
|
| 651 |
+
all_samples = []
|
| 652 |
+
|
| 653 |
+
# Source 1: OpenAssistant β real human conversations (~20K)
|
| 654 |
+
oasst = fetch_oasst(target=20000)
|
| 655 |
+
all_samples.extend(oasst)
|
| 656 |
+
|
| 657 |
+
# Source 2: Dolly β all human-written (~14K)
|
| 658 |
+
dolly = fetch_dolly()
|
| 659 |
+
all_samples.extend(dolly)
|
| 660 |
+
|
| 661 |
+
# Source 3: UltraChat 200K β large diverse conversations
|
| 662 |
+
ultrachat = fetch_ultrachat(target=150000)
|
| 663 |
+
all_samples.extend(ultrachat)
|
| 664 |
+
|
| 665 |
+
# Source 4: WizardLM Evol-Instruct β complex evolved instructions (~70K)
|
| 666 |
+
wizardlm = fetch_wizardlm(target=70000)
|
| 667 |
+
all_samples.extend(wizardlm)
|
| 668 |
+
|
| 669 |
+
# Source 5: SlimOrca-Dedup β reasoning and FLAN
|
| 670 |
+
orca = fetch_slimorca(target=200000)
|
| 671 |
+
all_samples.extend(orca)
|
| 672 |
+
|
| 673 |
+
# Source 6: Alpaca-cleaned (~52K)
|
| 674 |
+
alpaca = fetch_alpaca_cleaned()
|
| 675 |
+
all_samples.extend(alpaca)
|
| 676 |
+
|
| 677 |
+
# Source 7: Open-Platypus β STEM/reasoning
|
| 678 |
+
platypus = fetch_openplatypus()
|
| 679 |
+
all_samples.extend(platypus)
|
| 680 |
+
|
| 681 |
+
# Source 8: No Robots β human-written, high quality (~10K)
|
| 682 |
+
no_robots = fetch_no_robots()
|
| 683 |
+
all_samples.extend(no_robots)
|
| 684 |
+
|
| 685 |
+
# Source 9: Asterizer identity (hand-crafted)
|
| 686 |
+
identity = build_asterizer_identity()
|
| 687 |
+
# Repeat identity samples to ensure they're well-learned (1% of data)
|
| 688 |
+
identity_target = max(500, len(all_samples) // 100)
|
| 689 |
+
identity_expanded = []
|
| 690 |
+
while len(identity_expanded) < identity_target:
|
| 691 |
+
identity_expanded.extend(identity)
|
| 692 |
+
identity_expanded = identity_expanded[:identity_target]
|
| 693 |
+
all_samples.extend(identity_expanded)
|
| 694 |
+
|
| 695 |
+
print(f"\n--- Raw collection complete ---")
|
| 696 |
+
print(f" Total raw samples: {len(all_samples):,}")
|
| 697 |
+
print(f" Est. tokens: {compute_tokens(all_samples):,}")
|
| 698 |
+
|
| 699 |
+
# Shuffle before dedup to mix sources
|
| 700 |
+
random.seed(42)
|
| 701 |
+
random.shuffle(all_samples)
|
| 702 |
+
|
| 703 |
+
# Deduplication
|
| 704 |
+
print(f"\nDeduplicating...")
|
| 705 |
+
deduped = deduplicate(all_samples)
|
| 706 |
+
print(f" Before: {len(all_samples):,} β After: {len(deduped):,} "
|
| 707 |
+
f"(removed {len(all_samples) - len(deduped):,} dupes)")
|
| 708 |
+
|
| 709 |
+
# Check token count
|
| 710 |
+
tok_count = compute_tokens(deduped)
|
| 711 |
+
print(f" Est. tokens after dedup: {tok_count:,}")
|
| 712 |
+
|
| 713 |
+
# If we exceeded target, trim
|
| 714 |
+
if tok_count > args.target_tokens * 1.1:
|
| 715 |
+
# Keep identity samples, trim the rest
|
| 716 |
+
identity_fps = set(_fingerprint(s['output']) for s in identity)
|
| 717 |
+
identity_kept = [s for s in deduped if _fingerprint(s['output']) in identity_fps]
|
| 718 |
+
rest = [s for s in deduped if _fingerprint(s['output']) not in identity_fps]
|
| 719 |
+
random.shuffle(rest)
|
| 720 |
+
|
| 721 |
+
# Binary search for right cutoff
|
| 722 |
+
lo, hi = 0, len(rest)
|
| 723 |
+
while lo < hi:
|
| 724 |
+
mid = (lo + hi) // 2
|
| 725 |
+
if compute_tokens(rest[:mid] + identity_kept) < args.target_tokens:
|
| 726 |
+
lo = mid + 1
|
| 727 |
+
else:
|
| 728 |
+
hi = mid
|
| 729 |
+
rest = rest[:lo]
|
| 730 |
+
deduped = rest + identity_kept
|
| 731 |
+
random.shuffle(deduped)
|
| 732 |
+
tok_count = compute_tokens(deduped)
|
| 733 |
+
print(f" Trimmed to {len(deduped):,} samples ({tok_count:,} tokens)")
|
| 734 |
+
|
| 735 |
+
# Final shuffle
|
| 736 |
+
random.shuffle(deduped)
|
| 737 |
+
|
| 738 |
+
# Split train/val
|
| 739 |
+
val_size = max(500, int(len(deduped) * args.val_fraction))
|
| 740 |
+
val_data = deduped[:val_size]
|
| 741 |
+
train_data = deduped[val_size:]
|
| 742 |
+
|
| 743 |
+
print(f"\n Train: {len(train_data):,} samples ({compute_tokens(train_data):,} tokens)")
|
| 744 |
+
print(f" Val: {len(val_data):,} samples ({compute_tokens(val_data):,} tokens)")
|
| 745 |
+
|
| 746 |
+
# Write output
|
| 747 |
+
os.makedirs(args.output_dir, exist_ok=True)
|
| 748 |
+
train_path = os.path.join(args.output_dir, "train.json")
|
| 749 |
+
val_path = os.path.join(args.output_dir, "val.json")
|
| 750 |
+
|
| 751 |
+
with open(train_path, 'w', encoding='utf-8') as f:
|
| 752 |
+
json.dump(train_data, f, indent=2, ensure_ascii=False)
|
| 753 |
+
with open(val_path, 'w', encoding='utf-8') as f:
|
| 754 |
+
json.dump(val_data, f, indent=2, ensure_ascii=False)
|
| 755 |
+
|
| 756 |
+
elapsed = time.time() - t0
|
| 757 |
+
|
| 758 |
+
print(f"\n{'#'*60}")
|
| 759 |
+
print(f" DATASET BUILD COMPLETE")
|
| 760 |
+
print(f" Train: {train_path} ({len(train_data):,} samples)")
|
| 761 |
+
print(f" Val: {val_path} ({len(val_data):,} samples)")
|
| 762 |
+
print(f" Total tokens: ~{compute_tokens(deduped):,}")
|
| 763 |
+
print(f" Time: {elapsed:.0f}s")
|
| 764 |
+
print(f"{'#'*60}")
|
| 765 |
+
print(f"\nNext: litgpt finetune_full --config Base/configs/finetune_100m_english_instruct.yaml")
|
| 766 |
+
|
| 767 |
+
|
| 768 |
+
if __name__ == "__main__":
|
| 769 |
+
main()
|
Base/scripts/chat.py
ADDED
|
@@ -0,0 +1,68 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Interactive chat with the finetuned model. Model loads once, then you can ask questions back-to-back."""
|
| 2 |
+
|
| 3 |
+
import argparse
|
| 4 |
+
import torch
|
| 5 |
+
from pathlib import Path
|
| 6 |
+
from litgpt import Tokenizer
|
| 7 |
+
from litgpt.config import Config
|
| 8 |
+
from litgpt.model import GPT
|
| 9 |
+
|
| 10 |
+
|
| 11 |
+
def main():
|
| 12 |
+
parser = argparse.ArgumentParser()
|
| 13 |
+
parser.add_argument("--checkpoint_dir", type=str, default="Base/out/finetune/custom-100m-english-instruct/final")
|
| 14 |
+
parser.add_argument("--tokenizer_dir", type=str, default="Base/checkpoints/EleutherAI/pythia-160m")
|
| 15 |
+
parser.add_argument("--max_new_tokens", type=int, default=50)
|
| 16 |
+
parser.add_argument("--temperature", type=float, default=0.2)
|
| 17 |
+
parser.add_argument("--top_k", type=int, default=1)
|
| 18 |
+
args = parser.parse_args()
|
| 19 |
+
|
| 20 |
+
ckpt = Path(args.checkpoint_dir)
|
| 21 |
+
print(f"Loading model from {ckpt}...")
|
| 22 |
+
cfg = Config.from_checkpoint(ckpt)
|
| 23 |
+
model = GPT(cfg)
|
| 24 |
+
sd = torch.load(str(ckpt / "lit_model.pth"), map_location="cpu", weights_only=False)
|
| 25 |
+
if "model" in sd:
|
| 26 |
+
sd = sd["model"]
|
| 27 |
+
model.load_state_dict(sd, strict=False)
|
| 28 |
+
model = model.to("cuda").eval()
|
| 29 |
+
tok = Tokenizer(Path(args.tokenizer_dir))
|
| 30 |
+
print(f"Model loaded! ({sum(p.numel() for p in model.parameters()):,} params)")
|
| 31 |
+
print(f"Settings: temp={args.temperature}, top_k={args.top_k}, max_tokens={args.max_new_tokens}")
|
| 32 |
+
print("Type your question and press Enter. Type 'quit' or 'exit' to stop.\n")
|
| 33 |
+
|
| 34 |
+
while True:
|
| 35 |
+
try:
|
| 36 |
+
question = input("You: ").strip()
|
| 37 |
+
except (EOFError, KeyboardInterrupt):
|
| 38 |
+
print("\nBye!")
|
| 39 |
+
break
|
| 40 |
+
if not question:
|
| 41 |
+
continue
|
| 42 |
+
if question.lower() in ("quit", "exit", "q"):
|
| 43 |
+
print("Bye!")
|
| 44 |
+
break
|
| 45 |
+
|
| 46 |
+
alpaca = (
|
| 47 |
+
"Below is an instruction that describes a task. "
|
| 48 |
+
"Write a response that appropriately completes the request.\n\n"
|
| 49 |
+
f"### Instruction:\n{question}\n\n### Response:\n"
|
| 50 |
+
)
|
| 51 |
+
ids = tok.encode(alpaca, device="cuda").unsqueeze(0)
|
| 52 |
+
with torch.no_grad():
|
| 53 |
+
for _ in range(args.max_new_tokens):
|
| 54 |
+
logits = model(ids[:, -cfg.block_size:])
|
| 55 |
+
logits = logits[:, -1, :] / args.temperature
|
| 56 |
+
v, _ = torch.topk(logits, min(args.top_k, logits.size(-1)))
|
| 57 |
+
logits[logits < v[:, [-1]]] = float("-inf")
|
| 58 |
+
probs = torch.softmax(logits, dim=-1)
|
| 59 |
+
nxt = torch.multinomial(probs, num_samples=1)
|
| 60 |
+
ids = torch.cat([ids, nxt], dim=1)
|
| 61 |
+
if nxt.item() == tok.eos_id:
|
| 62 |
+
break
|
| 63 |
+
resp = tok.decode(ids[0]).split("### Response:")[-1].strip()
|
| 64 |
+
print(f"\nLUNA: {resp}\n")
|
| 65 |
+
|
| 66 |
+
|
| 67 |
+
if __name__ == "__main__":
|
| 68 |
+
main()
|
Base/scripts/clean_and_merge_pretrain.py
ADDED
|
@@ -0,0 +1,497 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""
|
| 3 |
+
Clean filtered_english_new_100m parquets + merge ALL pretrain litdata into one.
|
| 4 |
+
|
| 5 |
+
Steps:
|
| 6 |
+
1. Read 99,397 docs from filtered_english_new_100m parquets
|
| 7 |
+
2. Deep clean + smart filter (multiprocessed, 30 cores)
|
| 8 |
+
3. Tokenize β litdata_new_100m_clean
|
| 9 |
+
4. Merge into final unified pretrain dataset:
|
| 10 |
+
litdata_3b_clean (~2.94B tokens)
|
| 11 |
+
+ litdata_english_500m (~514M tokens)
|
| 12 |
+
+ litdata_combined (~257M tokens β existing cleaned English)
|
| 13 |
+
+ litdata_new_100m_clean (freshly cleaned from parquets)
|
| 14 |
+
= litdata_pretrain_final
|
| 15 |
+
|
| 16 |
+
All data is pure English pretraining text (Wikipedia + FineWeb-Edu + web).
|
| 17 |
+
No instruction/finetune data included.
|
| 18 |
+
|
| 19 |
+
Output: Base/data/litdata_pretrain_final/
|
| 20 |
+
"""
|
| 21 |
+
|
| 22 |
+
import json
|
| 23 |
+
import os
|
| 24 |
+
import re
|
| 25 |
+
import shutil
|
| 26 |
+
import time
|
| 27 |
+
import unicodedata
|
| 28 |
+
from pathlib import Path
|
| 29 |
+
from multiprocessing import Pool, cpu_count
|
| 30 |
+
from collections import Counter
|
| 31 |
+
|
| 32 |
+
import numpy as np
|
| 33 |
+
import pyarrow.parquet as pq
|
| 34 |
+
from tokenizers import Tokenizer
|
| 35 |
+
|
| 36 |
+
ROOT = Path(__file__).resolve().parent.parent.parent
|
| 37 |
+
BLOCK_SIZE = 1025
|
| 38 |
+
DTYPE = np.int32
|
| 39 |
+
CHUNK_BYTES_TARGET = 64 * 1024 * 1024
|
| 40 |
+
EOS_TOKEN_ID = 0
|
| 41 |
+
NUM_WORKERS = max(1, cpu_count() - 2)
|
| 42 |
+
ENCODE_BATCH = 8000
|
| 43 |
+
|
| 44 |
+
DATA_DIR = ROOT / "Base" / "data"
|
| 45 |
+
PARQUET_DIR = DATA_DIR / "filtered_english_new_100m"
|
| 46 |
+
CLEAN_100M_DIR = DATA_DIR / "litdata_new_100m_clean"
|
| 47 |
+
FINAL_DIR = DATA_DIR / "litdata_pretrain_final"
|
| 48 |
+
TOKENIZER_PATH = str(ROOT / "Base" / "checkpoints" / "EleutherAI" / "pythia-160m" / "tokenizer.json")
|
| 49 |
+
|
| 50 |
+
# Datasets to merge (in order)
|
| 51 |
+
MERGE_SOURCES = [
|
| 52 |
+
DATA_DIR / "litdata_3b_clean",
|
| 53 |
+
DATA_DIR / "litdata_english_500m",
|
| 54 |
+
DATA_DIR / "litdata_combined",
|
| 55 |
+
]
|
| 56 |
+
|
| 57 |
+
|
| 58 |
+
# ==============================================================================
|
| 59 |
+
# CLEANING PIPELINE (same as reclean_3b.py)
|
| 60 |
+
# ==============================================================================
|
| 61 |
+
|
| 62 |
+
CONTROL_CHARS = [
|
| 63 |
+
"\x00", "\x01", "\x02", "\x03", "\x04", "\x05", "\x06", "\x07",
|
| 64 |
+
"\x08", "\x0b", "\x0c", "\x0e", "\x0f", "\x10", "\x11", "\x12",
|
| 65 |
+
"\x13", "\x14", "\x15", "\x16", "\x17", "\x18", "\x19", "\x1a",
|
| 66 |
+
"\x1b", "\x1c", "\x1d", "\x1e", "\x1f", "\x7f", "\ufeff", "\ufffd",
|
| 67 |
+
]
|
| 68 |
+
|
| 69 |
+
HTML_ENTITIES = [
|
| 70 |
+
("&", "&"), ("<", "<"), (">", ">"),
|
| 71 |
+
(""", '"'), ("'", "'"), ("'", "'"),
|
| 72 |
+
(" ", " "), ("—", " - "), ("–", "-"),
|
| 73 |
+
("…", "..."), ("«", '"'), ("»", '"'),
|
| 74 |
+
("•", "- "), ("·", " "), ("©", "(c)"),
|
| 75 |
+
("®", "(R)"), ("™", "(TM)"), ("°", " degrees"),
|
| 76 |
+
]
|
| 77 |
+
|
| 78 |
+
RE_URL = re.compile(r'https?://\S+|www\.\S+', re.I)
|
| 79 |
+
RE_EMAIL = re.compile(r'\b[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}\b')
|
| 80 |
+
RE_FILE_PATH = re.compile(r'(?:[A-Z]:\\|/(?:home|usr|var|etc|opt)/)\S+')
|
| 81 |
+
RE_HTML_TAG = re.compile(r'</?[a-zA-Z][a-zA-Z0-9]*(?:\s[^>]*)?\s*/?>')
|
| 82 |
+
RE_HTML_COMMENT = re.compile(r'<!--.*?-->', re.DOTALL)
|
| 83 |
+
RE_CODE_BLOCK = re.compile(r'```[\s\S]*?```')
|
| 84 |
+
RE_IMPORT = re.compile(r'^(?:import |from \S+ import |#include |using namespace |require\()', re.M)
|
| 85 |
+
RE_REPEATED_LINE = re.compile(r'^(.{20,})\n(?:\1\n?)+', re.M)
|
| 86 |
+
RE_REPEATED_PUNCT = re.compile(r'([!?.])\1{3,}')
|
| 87 |
+
RE_REPEATED_CHAR = re.compile(r'(.)\1{5,}')
|
| 88 |
+
RE_REPEATED_WORD = re.compile(r'\b(\w+)(?:\s+\1){2,}\b', re.I)
|
| 89 |
+
RE_MULTI_NEWLINE = re.compile(r'\n{4,}')
|
| 90 |
+
RE_MULTI_SPACE = re.compile(r'[ \t]{2,}')
|
| 91 |
+
RE_TRAILING_SPACE = re.compile(r'[ \t]+$', re.M)
|
| 92 |
+
RE_NO_SPACE_AFTER_PERIOD = re.compile(r'([.!?])([A-Z])')
|
| 93 |
+
RE_DOUBLE_PERIOD = re.compile(r'\.{2}(?!\.)')
|
| 94 |
+
RE_SPACE_BEFORE_PUNCT = re.compile(r'\s+([.,;:!?])')
|
| 95 |
+
|
| 96 |
+
RE_CJK = re.compile(r'[\u4e00-\u9fff\u3040-\u309f\u30a0-\u30ff\uac00-\ud7af]{3,}')
|
| 97 |
+
RE_ARABIC = re.compile(r'[\u0600-\u06ff]{5,}')
|
| 98 |
+
RE_CYRILLIC = re.compile(r'[\u0400-\u04ff]{5,}')
|
| 99 |
+
RE_DEVANAGARI = re.compile(r'[\u0900-\u097f]{5,}')
|
| 100 |
+
RE_RESIDUAL_CODE = re.compile(
|
| 101 |
+
r'(function\s*\(|var\s+\w+\s*=|console\.log|document\.get|if\s*\(\s*\w+\s*[!=]==)', re.I
|
| 102 |
+
)
|
| 103 |
+
RE_COOKIE_LINE = re.compile(r'^.*(?:cookie|cookies)\s+(?:policy|consent|notice|preferences|settings).*$', re.I | re.M)
|
| 104 |
+
RE_SUBSCRIBE_LINE = re.compile(r'^.*(?:subscribe|sign\s*up\s+(?:for|to)\s+(?:our|the)\s+newsletter|unsubscribe|opt[\s-]*out\s+of).*$', re.I | re.M)
|
| 105 |
+
RE_CLICKBAIT_LINE = re.compile(r'^.*(?:you\s+won\'?t\s+believe|click\s+here|read\s+more\s*\.{0,3}$|share\s+this\s+(?:article|post|story)|trending\s+now|sponsored\s+content|advertisement).*$', re.I | re.M)
|
| 106 |
+
RE_SOCIAL_LINE = re.compile(r'^.*(?:follow\s+us\s+on|share\s+on\s+(?:facebook|twitter|linkedin|instagram)|like\s+us\s+on|tweet\s+this).*$', re.I | re.M)
|
| 107 |
+
RE_NAV_LINE = re.compile(r'^.*(?:skip\s+to\s+(?:main\s+)?content|back\s+to\s+top|previous\s+article|next\s+article|related\s+(?:articles|posts)).*$', re.I | re.M)
|
| 108 |
+
RE_LOGIN_LINE = re.compile(r'^.*(?:log\s*in\s+to\s+(?:your|an)\s+account|create\s+(?:a\s+)?(?:free\s+)?account|forgot\s+(?:your\s+)?password|already\s+(?:a\s+)?member).*$', re.I | re.M)
|
| 109 |
+
RE_COMMENT_LINE = re.compile(r'^.*(?:leave\s+a\s+(?:comment|reply)|post\s+a\s+comment|\d+\s+comments?$|logged\s+in\s+as).*$', re.I | re.M)
|
| 110 |
+
RE_COPYRIGHT_LINE = re.compile(r'^.*(?:all\s+rights\s+reserved|\(c\)\s*\d{4}|copyright\s+\d{4}).*$', re.I | re.M)
|
| 111 |
+
|
| 112 |
+
|
| 113 |
+
def deep_clean(text):
|
| 114 |
+
if not text or len(text.strip()) < 30:
|
| 115 |
+
return None
|
| 116 |
+
text = unicodedata.normalize("NFKC", text)
|
| 117 |
+
for ch in CONTROL_CHARS:
|
| 118 |
+
text = text.replace(ch, "")
|
| 119 |
+
for old, new in HTML_ENTITIES:
|
| 120 |
+
text = text.replace(old, new)
|
| 121 |
+
text = RE_HTML_COMMENT.sub("", text)
|
| 122 |
+
text = RE_HTML_TAG.sub("", text)
|
| 123 |
+
text = RE_URL.sub("", text)
|
| 124 |
+
text = RE_EMAIL.sub("", text)
|
| 125 |
+
text = RE_FILE_PATH.sub("", text)
|
| 126 |
+
text = RE_CODE_BLOCK.sub("", text)
|
| 127 |
+
text = RE_REPEATED_LINE.sub(r'\1', text)
|
| 128 |
+
text = RE_REPEATED_PUNCT.sub(r'\1\1\1', text)
|
| 129 |
+
text = RE_REPEATED_CHAR.sub(r'\1\1\1', text)
|
| 130 |
+
text = RE_REPEATED_WORD.sub(r'\1', text)
|
| 131 |
+
text = text.replace('\t', ' ')
|
| 132 |
+
text = RE_TRAILING_SPACE.sub('', text)
|
| 133 |
+
text = RE_MULTI_SPACE.sub(' ', text)
|
| 134 |
+
text = RE_MULTI_NEWLINE.sub('\n\n\n', text)
|
| 135 |
+
text = RE_DOUBLE_PERIOD.sub('.', text)
|
| 136 |
+
text = RE_NO_SPACE_AFTER_PERIOD.sub(r'\1 \2', text)
|
| 137 |
+
text = RE_SPACE_BEFORE_PUNCT.sub(r'\1', text)
|
| 138 |
+
text = text.replace('\u2018', "'").replace('\u2019', "'")
|
| 139 |
+
text = text.replace('\u201c', '"').replace('\u201d', '"')
|
| 140 |
+
text = text.replace('\u2013', '-').replace('\u2014', ' - ')
|
| 141 |
+
text = text.replace('\u2026', '...')
|
| 142 |
+
text = text.replace('\u2022', '- ')
|
| 143 |
+
text = text.replace('\u00b7', ' ')
|
| 144 |
+
text = text.replace('\u00a0', ' ')
|
| 145 |
+
lines = text.split('\n')
|
| 146 |
+
clean_lines = []
|
| 147 |
+
for line in lines:
|
| 148 |
+
line = line.strip()
|
| 149 |
+
if not line:
|
| 150 |
+
clean_lines.append('')
|
| 151 |
+
continue
|
| 152 |
+
if len(line) > 10:
|
| 153 |
+
alpha_count = sum(1 for c in line if c.isalpha())
|
| 154 |
+
if alpha_count / len(line) < 0.40:
|
| 155 |
+
continue
|
| 156 |
+
if line.count('|') > 3 or line.count('{') > 2 or line.count('}') > 2:
|
| 157 |
+
continue
|
| 158 |
+
if RE_IMPORT.match(line):
|
| 159 |
+
continue
|
| 160 |
+
if line and line[0].isalpha() and line[0].islower():
|
| 161 |
+
if not clean_lines or clean_lines[-1] == '' or clean_lines[-1].rstrip().endswith(('.', '!', '?', ':')):
|
| 162 |
+
line = line[0].upper() + line[1:]
|
| 163 |
+
clean_lines.append(line)
|
| 164 |
+
text = '\n'.join(clean_lines)
|
| 165 |
+
text = text.strip()
|
| 166 |
+
paragraphs = text.split('\n\n')
|
| 167 |
+
seen = set()
|
| 168 |
+
unique_paragraphs = []
|
| 169 |
+
for p in paragraphs:
|
| 170 |
+
p_stripped = p.strip()
|
| 171 |
+
if not p_stripped:
|
| 172 |
+
continue
|
| 173 |
+
p_key = ' '.join(p_stripped.lower().split())
|
| 174 |
+
if p_key not in seen:
|
| 175 |
+
seen.add(p_key)
|
| 176 |
+
unique_paragraphs.append(p_stripped)
|
| 177 |
+
text = '\n\n'.join(unique_paragraphs)
|
| 178 |
+
text = text.strip()
|
| 179 |
+
if len(text) < 50:
|
| 180 |
+
return None
|
| 181 |
+
if len(text.split()) < 10:
|
| 182 |
+
return None
|
| 183 |
+
ascii_count = sum(1 for c in text if ord(c) < 128)
|
| 184 |
+
if ascii_count / max(len(text), 1) < 0.85:
|
| 185 |
+
return None
|
| 186 |
+
return text
|
| 187 |
+
|
| 188 |
+
|
| 189 |
+
def smart_filter(text):
|
| 190 |
+
words = text.split()
|
| 191 |
+
word_count = len(words)
|
| 192 |
+
if word_count < 50:
|
| 193 |
+
return False, text, "too short"
|
| 194 |
+
scripts = []
|
| 195 |
+
if RE_CJK.search(text): scripts.append("CJK")
|
| 196 |
+
if RE_ARABIC.search(text): scripts.append("Arabic")
|
| 197 |
+
if RE_CYRILLIC.search(text): scripts.append("Cyrillic")
|
| 198 |
+
if RE_DEVANAGARI.search(text): scripts.append("Devanagari")
|
| 199 |
+
if scripts:
|
| 200 |
+
return False, text, "non-English"
|
| 201 |
+
if word_count > 50:
|
| 202 |
+
unique_ratio = len(set(w.lower() for w in words)) / word_count
|
| 203 |
+
if unique_ratio < 0.20:
|
| 204 |
+
return False, text, "repetitive"
|
| 205 |
+
code_matches = RE_RESIDUAL_CODE.findall(text)
|
| 206 |
+
if len(code_matches) >= 5:
|
| 207 |
+
return False, text, "residual code"
|
| 208 |
+
for pattern in [RE_COOKIE_LINE, RE_SUBSCRIBE_LINE, RE_CLICKBAIT_LINE,
|
| 209 |
+
RE_SOCIAL_LINE, RE_NAV_LINE, RE_LOGIN_LINE,
|
| 210 |
+
RE_COMMENT_LINE, RE_COPYRIGHT_LINE]:
|
| 211 |
+
text = pattern.sub('', text)
|
| 212 |
+
lines = text.split('\n')
|
| 213 |
+
clean_lines = []
|
| 214 |
+
for line in lines:
|
| 215 |
+
stripped = line.strip()
|
| 216 |
+
if stripped and len(stripped) > 10:
|
| 217 |
+
digit_count = sum(1 for c in stripped if c.isdigit() or c in ' ,.\t-+/%$')
|
| 218 |
+
if digit_count / len(stripped) > 0.80:
|
| 219 |
+
continue
|
| 220 |
+
clean_lines.append(line)
|
| 221 |
+
text = '\n'.join(clean_lines)
|
| 222 |
+
text = re.sub(r'\n{3,}', '\n\n', text)
|
| 223 |
+
text = text.strip()
|
| 224 |
+
if len(text.split()) < 50:
|
| 225 |
+
return False, text, "too short after stripping"
|
| 226 |
+
return True, text, None
|
| 227 |
+
|
| 228 |
+
|
| 229 |
+
def clean_and_filter(text):
|
| 230 |
+
"""Combined pipeline for multiprocessing. Returns (text|None, reason|None, boilerplate)."""
|
| 231 |
+
result = deep_clean(text)
|
| 232 |
+
if result is None:
|
| 233 |
+
return None, "deep_clean_drop", 0
|
| 234 |
+
keep, stripped, reason = smart_filter(result)
|
| 235 |
+
if not keep:
|
| 236 |
+
return None, reason, 0
|
| 237 |
+
return stripped, None, len(result) - len(stripped)
|
| 238 |
+
|
| 239 |
+
|
| 240 |
+
# ==============================================================================
|
| 241 |
+
# LITDATA I/O
|
| 242 |
+
# ==============================================================================
|
| 243 |
+
|
| 244 |
+
def write_litdata_chunks(output_dir, token_stream, config):
|
| 245 |
+
os.makedirs(output_dir, exist_ok=True)
|
| 246 |
+
dtype_size = DTYPE().itemsize
|
| 247 |
+
tokens_per_chunk = (CHUNK_BYTES_TARGET // dtype_size // BLOCK_SIZE) * BLOCK_SIZE
|
| 248 |
+
chunks_metadata = []
|
| 249 |
+
pos = 0
|
| 250 |
+
chunk_idx = 0
|
| 251 |
+
while pos < len(token_stream):
|
| 252 |
+
remaining = len(token_stream) - pos
|
| 253 |
+
chunk_tokens = min(tokens_per_chunk, remaining)
|
| 254 |
+
num_blocks = chunk_tokens // BLOCK_SIZE
|
| 255 |
+
if num_blocks == 0:
|
| 256 |
+
break
|
| 257 |
+
actual_tokens = num_blocks * BLOCK_SIZE
|
| 258 |
+
chunk_data = token_stream[pos:pos + actual_tokens]
|
| 259 |
+
filename = f"chunk-0-{chunk_idx}.bin"
|
| 260 |
+
filepath = os.path.join(output_dir, filename)
|
| 261 |
+
header_num = np.array([num_blocks], dtype=np.uint32)
|
| 262 |
+
offsets = np.arange(num_blocks + 1, dtype=np.uint32) * (BLOCK_SIZE * dtype_size)
|
| 263 |
+
header = np.concatenate([header_num, offsets])
|
| 264 |
+
with open(filepath, "wb") as f:
|
| 265 |
+
header.tofile(f)
|
| 266 |
+
chunk_data.tofile(f)
|
| 267 |
+
meta = {
|
| 268 |
+
"chunk_bytes": int(header.nbytes + chunk_data.nbytes),
|
| 269 |
+
"chunk_size": num_blocks,
|
| 270 |
+
"dim": int(actual_tokens),
|
| 271 |
+
"filename": filename,
|
| 272 |
+
}
|
| 273 |
+
chunks_metadata.append(meta)
|
| 274 |
+
pos += actual_tokens
|
| 275 |
+
chunk_idx += 1
|
| 276 |
+
if chunk_idx % 5 == 0 or pos >= len(token_stream):
|
| 277 |
+
print(f" Written chunk {chunk_idx} ({pos:,}/{len(token_stream):,} tokens)")
|
| 278 |
+
index = {"chunks": chunks_metadata, "config": config, "updated_at": str(time.time())}
|
| 279 |
+
with open(os.path.join(output_dir, "index.json"), "w") as f:
|
| 280 |
+
json.dump(index, f, indent=2)
|
| 281 |
+
return chunks_metadata
|
| 282 |
+
|
| 283 |
+
|
| 284 |
+
def merge_litdata_dirs(sources, dest_dir):
|
| 285 |
+
"""Copy chunks from multiple litdata dirs into one, renumbering sequentially."""
|
| 286 |
+
os.makedirs(str(dest_dir), exist_ok=True)
|
| 287 |
+
all_chunks = []
|
| 288 |
+
config = None
|
| 289 |
+
chunk_offset = 0
|
| 290 |
+
|
| 291 |
+
for src in sources:
|
| 292 |
+
idx_path = src / "index.json"
|
| 293 |
+
if not idx_path.exists():
|
| 294 |
+
print(f" WARNING: {src} has no index.json, skipping")
|
| 295 |
+
continue
|
| 296 |
+
with open(idx_path) as f:
|
| 297 |
+
index = json.load(f)
|
| 298 |
+
if config is None:
|
| 299 |
+
config = index.get("config", {"block_size": BLOCK_SIZE, "vocab_size": 50277})
|
| 300 |
+
|
| 301 |
+
src_tokens = sum(c["dim"] for c in index["chunks"])
|
| 302 |
+
for chunk_meta in index["chunks"]:
|
| 303 |
+
old_fn = chunk_meta["filename"]
|
| 304 |
+
new_fn = f"chunk-0-{chunk_offset}.bin"
|
| 305 |
+
shutil.copy2(str(src / old_fn), str(dest_dir / new_fn))
|
| 306 |
+
new_meta = dict(chunk_meta)
|
| 307 |
+
new_meta["filename"] = new_fn
|
| 308 |
+
all_chunks.append(new_meta)
|
| 309 |
+
chunk_offset += 1
|
| 310 |
+
|
| 311 |
+
print(f" {src.name}: {len(index['chunks'])} chunks, {src_tokens:,} tokens")
|
| 312 |
+
|
| 313 |
+
combined_index = {
|
| 314 |
+
"chunks": all_chunks,
|
| 315 |
+
"config": config,
|
| 316 |
+
"updated_at": str(time.time()),
|
| 317 |
+
}
|
| 318 |
+
with open(str(dest_dir / "index.json"), "w") as f:
|
| 319 |
+
json.dump(combined_index, f, indent=2)
|
| 320 |
+
|
| 321 |
+
total = sum(c["dim"] for c in all_chunks)
|
| 322 |
+
return all_chunks, total
|
| 323 |
+
|
| 324 |
+
|
| 325 |
+
# ==============================================================================
|
| 326 |
+
# MAIN
|
| 327 |
+
# ==============================================================================
|
| 328 |
+
|
| 329 |
+
def main():
|
| 330 |
+
t_start = time.time()
|
| 331 |
+
|
| 332 |
+
print("Loading tokenizer...")
|
| 333 |
+
tokenizer = Tokenizer.from_file(TOKENIZER_PATH)
|
| 334 |
+
|
| 335 |
+
config = {
|
| 336 |
+
"block_size": BLOCK_SIZE,
|
| 337 |
+
"vocab_size": tokenizer.get_vocab_size(),
|
| 338 |
+
}
|
| 339 |
+
|
| 340 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 341 |
+
# STEP 1: Read parquets from filtered_english_new_100m
|
| 342 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 343 |
+
print(f"\n{'='*75}")
|
| 344 |
+
print(f" STEP 1: READING PARQUETS FROM {PARQUET_DIR.name}")
|
| 345 |
+
print(f"{'='*75}")
|
| 346 |
+
|
| 347 |
+
parquet_files = sorted(PARQUET_DIR.glob("*.parquet"))
|
| 348 |
+
all_texts = []
|
| 349 |
+
for pf in parquet_files:
|
| 350 |
+
table = pq.read_table(str(pf), columns=["text"])
|
| 351 |
+
for val in table.column("text"):
|
| 352 |
+
all_texts.append(val.as_py())
|
| 353 |
+
print(f" Read {len(all_texts):,} documents from {len(parquet_files)} parquet files")
|
| 354 |
+
|
| 355 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 356 |
+
# STEP 2: Deep clean + smart filter (multiprocessed)
|
| 357 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 358 |
+
print(f"\n{'='*75}")
|
| 359 |
+
print(f" STEP 2: CLEANING {len(all_texts):,} DOCUMENTS ({NUM_WORKERS} workers)")
|
| 360 |
+
print(f"{'='*75}")
|
| 361 |
+
|
| 362 |
+
t_clean = time.time()
|
| 363 |
+
pool = Pool(processes=NUM_WORKERS)
|
| 364 |
+
results = pool.map(clean_and_filter, all_texts, chunksize=256)
|
| 365 |
+
pool.close()
|
| 366 |
+
pool.join()
|
| 367 |
+
del all_texts
|
| 368 |
+
|
| 369 |
+
cleaned_texts = []
|
| 370 |
+
drop_reasons = Counter()
|
| 371 |
+
total_boilerplate = 0
|
| 372 |
+
for cleaned, reason, bp in results:
|
| 373 |
+
if cleaned is not None:
|
| 374 |
+
cleaned_texts.append(cleaned)
|
| 375 |
+
total_boilerplate += bp
|
| 376 |
+
else:
|
| 377 |
+
drop_reasons[reason] += 1
|
| 378 |
+
del results
|
| 379 |
+
|
| 380 |
+
total_dropped = sum(drop_reasons.values())
|
| 381 |
+
clean_time = time.time() - t_clean
|
| 382 |
+
print(f" Cleaned in {clean_time:.1f}s")
|
| 383 |
+
print(f" Kept: {len(cleaned_texts):,} | Dropped: {total_dropped:,}")
|
| 384 |
+
if drop_reasons:
|
| 385 |
+
print(f" Drop reasons:")
|
| 386 |
+
for reason, count in sorted(drop_reasons.items(), key=lambda x: -x[1]):
|
| 387 |
+
print(f" {reason:<35} {count:>6,}")
|
| 388 |
+
print(f" Boilerplate stripped: {total_boilerplate:,} chars")
|
| 389 |
+
|
| 390 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 391 |
+
# STEP 3: Tokenize + write litdata
|
| 392 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 393 |
+
print(f"\n{'='*75}")
|
| 394 |
+
print(f" STEP 3: TOKENIZING {len(cleaned_texts):,} DOCUMENTS")
|
| 395 |
+
print(f"{'='*75}")
|
| 396 |
+
|
| 397 |
+
t_tok = time.time()
|
| 398 |
+
all_token_ids = []
|
| 399 |
+
total_tokens = 0
|
| 400 |
+
for i in range(0, len(cleaned_texts), ENCODE_BATCH):
|
| 401 |
+
batch = cleaned_texts[i:i+ENCODE_BATCH]
|
| 402 |
+
encoded = tokenizer.encode_batch(batch, add_special_tokens=False)
|
| 403 |
+
for enc in encoded:
|
| 404 |
+
ids = enc.ids
|
| 405 |
+
all_token_ids.extend(ids)
|
| 406 |
+
all_token_ids.append(EOS_TOKEN_ID)
|
| 407 |
+
total_tokens += len(ids) + 1
|
| 408 |
+
done = min(i + ENCODE_BATCH, len(cleaned_texts))
|
| 409 |
+
if done % 20000 == 0 or done == len(cleaned_texts):
|
| 410 |
+
print(f" Tokenized {done:,}/{len(cleaned_texts):,} ({total_tokens:,} tokens)")
|
| 411 |
+
del cleaned_texts
|
| 412 |
+
|
| 413 |
+
token_array = np.array(all_token_ids, dtype=DTYPE)
|
| 414 |
+
del all_token_ids
|
| 415 |
+
tok_time = time.time() - t_tok
|
| 416 |
+
print(f" Tokenized in {tok_time:.1f}s β {total_tokens:,} tokens")
|
| 417 |
+
|
| 418 |
+
print(f"\n Writing litdata to {CLEAN_100M_DIR}...")
|
| 419 |
+
chunks_new = write_litdata_chunks(str(CLEAN_100M_DIR), token_array, config)
|
| 420 |
+
new_tokens = sum(c["dim"] for c in chunks_new)
|
| 421 |
+
del token_array
|
| 422 |
+
print(f" Written {len(chunks_new)} chunks, {new_tokens:,} tokens")
|
| 423 |
+
|
| 424 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 425 |
+
# STEP 4: Merge ALL pretrain data into one
|
| 426 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 427 |
+
print(f"\n{'='*75}")
|
| 428 |
+
print(f" STEP 4: MERGING ALL PRETRAIN DATA β {FINAL_DIR.name}")
|
| 429 |
+
print(f"{'='*75}\n")
|
| 430 |
+
|
| 431 |
+
all_sources = MERGE_SOURCES + [CLEAN_100M_DIR]
|
| 432 |
+
all_chunks, total_final = merge_litdata_dirs(all_sources, FINAL_DIR)
|
| 433 |
+
|
| 434 |
+
print(f"\n Final: {len(all_chunks)} chunks, {total_final:,} tokens")
|
| 435 |
+
|
| 436 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 437 |
+
# REPORT
|
| 438 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 439 |
+
total_time = time.time() - t_start
|
| 440 |
+
|
| 441 |
+
# Gather per-source token counts
|
| 442 |
+
source_info = []
|
| 443 |
+
for src in all_sources:
|
| 444 |
+
if (src / "index.json").exists():
|
| 445 |
+
with open(src / "index.json") as f:
|
| 446 |
+
idx = json.load(f)
|
| 447 |
+
tokens = sum(c["dim"] for c in idx["chunks"])
|
| 448 |
+
source_info.append((src.name, tokens, len(idx["chunks"])))
|
| 449 |
+
|
| 450 |
+
report = []
|
| 451 |
+
report.append(f"{'='*75}")
|
| 452 |
+
report.append(f" FINAL PRETRAIN DATASET β BUILD REPORT")
|
| 453 |
+
report.append(f"{'='*75}")
|
| 454 |
+
report.append(f"")
|
| 455 |
+
report.append(f" Total time: {total_time:.0f}s ({total_time/60:.1f} min)")
|
| 456 |
+
report.append(f"")
|
| 457 |
+
report.append(f" FRESHLY CLEANED: filtered_english_new_100m")
|
| 458 |
+
report.append(f" {'-'*60}")
|
| 459 |
+
report.append(f" Input docs: 99,397")
|
| 460 |
+
report.append(f" Kept: {len(cleaned_texts) if 'cleaned_texts' in dir() else new_tokens}")
|
| 461 |
+
report.append(f" Dropped: {total_dropped}")
|
| 462 |
+
if drop_reasons:
|
| 463 |
+
for reason, count in sorted(drop_reasons.items(), key=lambda x: -x[1]):
|
| 464 |
+
report.append(f" {reason}: {count}")
|
| 465 |
+
report.append(f" Boilerplate: {total_boilerplate:,} chars stripped")
|
| 466 |
+
report.append(f" Tokens: {new_tokens:,}")
|
| 467 |
+
report.append(f"")
|
| 468 |
+
report.append(f" MERGED SOURCES")
|
| 469 |
+
report.append(f" {'-'*60}")
|
| 470 |
+
for name, tokens, nchunks in source_info:
|
| 471 |
+
report.append(f" {name:<30} {tokens:>15,} tokens ({nchunks} chunks)")
|
| 472 |
+
report.append(f" {'-'*60}")
|
| 473 |
+
report.append(f" {'TOTAL':<30} {total_final:>15,} tokens ({len(all_chunks)} chunks)")
|
| 474 |
+
report.append(f"")
|
| 475 |
+
report.append(f" OUTPUT")
|
| 476 |
+
report.append(f" {'-'*60}")
|
| 477 |
+
report.append(f" Location: {FINAL_DIR}")
|
| 478 |
+
report.append(f" Chunks: {len(all_chunks)}")
|
| 479 |
+
report.append(f" Tokens: {total_final:,}")
|
| 480 |
+
report.append(f" Format: litdata binary (int32, BLOCK_SIZE=1025, EOS=0)")
|
| 481 |
+
report.append(f"")
|
| 482 |
+
report.append(f" DATA TYPE: Pure English pretraining text")
|
| 483 |
+
report.append(f" Sources: Wikipedia, FineWeb-Edu, web crawl")
|
| 484 |
+
report.append(f" NO instruction/finetune data included")
|
| 485 |
+
report.append(f"{'='*75}")
|
| 486 |
+
|
| 487 |
+
full_report = '\n'.join(report)
|
| 488 |
+
print(f"\n{full_report}")
|
| 489 |
+
|
| 490 |
+
with open(FINAL_DIR / "BUILD_REPORT.txt", "w", encoding="utf-8") as f:
|
| 491 |
+
f.write(full_report)
|
| 492 |
+
print(f"\n Report saved to: {FINAL_DIR / 'BUILD_REPORT.txt'}")
|
| 493 |
+
print(f" Done! Final unified pretrain dataset ready.")
|
| 494 |
+
|
| 495 |
+
|
| 496 |
+
if __name__ == "__main__":
|
| 497 |
+
main()
|
Base/scripts/consolidate_datasets.py
ADDED
|
@@ -0,0 +1,194 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Consolidate all pretraining and finetuning datasets into single files.
|
| 3 |
+
Creates:
|
| 4 |
+
Base/data/consolidated_pretrain/ -> all_pretrain_data.parquet (3B + English)
|
| 5 |
+
Base/Datasets/consolidated_finetune/ -> all_finetune_data.json (v1 + English instruct)
|
| 6 |
+
Also counts tokens using the Pythia tokenizer (batch mode for speed).
|
| 7 |
+
"""
|
| 8 |
+
|
| 9 |
+
import os
|
| 10 |
+
import json
|
| 11 |
+
import glob
|
| 12 |
+
import pandas as pd
|
| 13 |
+
from pathlib import Path
|
| 14 |
+
from tokenizers import Tokenizer
|
| 15 |
+
|
| 16 |
+
ROOT = Path(__file__).resolve().parent.parent.parent # LUNA root
|
| 17 |
+
|
| 18 |
+
# ββ Tokenizer ββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 19 |
+
print("Loading tokenizer...")
|
| 20 |
+
tokenizer = Tokenizer.from_file(
|
| 21 |
+
str(ROOT / "Base" / "checkpoints" / "EleutherAI" / "pythia-160m" / "tokenizer.json")
|
| 22 |
+
)
|
| 23 |
+
|
| 24 |
+
|
| 25 |
+
BATCH_SIZE = 10000 # encode_batch processes this many at once
|
| 26 |
+
|
| 27 |
+
|
| 28 |
+
def count_tokens_text(texts):
|
| 29 |
+
"""Count total tokens using fast batch encoding."""
|
| 30 |
+
total = 0
|
| 31 |
+
for i in range(0, len(texts), BATCH_SIZE):
|
| 32 |
+
batch = texts[i : i + BATCH_SIZE]
|
| 33 |
+
encoded = tokenizer.encode_batch(batch, add_special_tokens=False)
|
| 34 |
+
total += sum(len(e.ids) for e in encoded)
|
| 35 |
+
done = min(i + BATCH_SIZE, len(texts))
|
| 36 |
+
if done % 100000 == 0 or done == len(texts):
|
| 37 |
+
print(f" Tokenized {done:,}/{len(texts):,} documents ({total:,} tokens)")
|
| 38 |
+
return total
|
| 39 |
+
|
| 40 |
+
|
| 41 |
+
# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 42 |
+
# 1. PRETRAINING DATA (filtered_3b + filtered_english)
|
| 43 |
+
# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 44 |
+
print("\n" + "=" * 60)
|
| 45 |
+
print("CONSOLIDATING PRETRAINING DATA")
|
| 46 |
+
print("=" * 60)
|
| 47 |
+
|
| 48 |
+
pretrain_out = ROOT / "Base" / "data" / "consolidated_pretrain"
|
| 49 |
+
pretrain_out.mkdir(parents=True, exist_ok=True)
|
| 50 |
+
|
| 51 |
+
# Read filtered_3b (15 parquets)
|
| 52 |
+
dir_3b = ROOT / "Base" / "data" / "filtered_3b"
|
| 53 |
+
files_3b = sorted(glob.glob(str(dir_3b / "*.parquet")))
|
| 54 |
+
print(f"\nReading filtered_3b: {len(files_3b)} files...")
|
| 55 |
+
dfs_3b = [pd.read_parquet(f) for f in files_3b]
|
| 56 |
+
df_3b = pd.concat(dfs_3b, ignore_index=True)
|
| 57 |
+
print(f" -> {len(df_3b):,} documents")
|
| 58 |
+
|
| 59 |
+
# Read filtered_english (fineweb + wiki parquets)
|
| 60 |
+
dir_en = ROOT / "Base" / "data" / "filtered_english"
|
| 61 |
+
files_en = sorted(glob.glob(str(dir_en / "*.parquet")))
|
| 62 |
+
print(f"\nReading filtered_english: {len(files_en)} files...")
|
| 63 |
+
dfs_en = [pd.read_parquet(f) for f in files_en]
|
| 64 |
+
df_en = pd.concat(dfs_en, ignore_index=True)
|
| 65 |
+
print(f" -> {len(df_en):,} documents")
|
| 66 |
+
|
| 67 |
+
# Add source labels
|
| 68 |
+
df_3b["source"] = "filtered_3b"
|
| 69 |
+
df_en_list = []
|
| 70 |
+
for f in files_en:
|
| 71 |
+
fname = os.path.basename(f)
|
| 72 |
+
tmp = pd.read_parquet(f)
|
| 73 |
+
if fname.startswith("fineweb"):
|
| 74 |
+
tmp["source"] = "english_fineweb"
|
| 75 |
+
elif fname.startswith("wiki"):
|
| 76 |
+
tmp["source"] = "english_wiki"
|
| 77 |
+
else:
|
| 78 |
+
tmp["source"] = "english_other"
|
| 79 |
+
df_en_list.append(tmp)
|
| 80 |
+
df_en_labeled = pd.concat(df_en_list, ignore_index=True)
|
| 81 |
+
|
| 82 |
+
# Combine all pretraining data
|
| 83 |
+
df_pretrain = pd.concat([df_3b, df_en_labeled], ignore_index=True)
|
| 84 |
+
out_path = pretrain_out / "all_pretrain_data.parquet"
|
| 85 |
+
df_pretrain.to_parquet(str(out_path), index=False)
|
| 86 |
+
print(f"\nSaved combined pretrain: {out_path}")
|
| 87 |
+
print(f" Total documents: {len(df_pretrain):,}")
|
| 88 |
+
|
| 89 |
+
# Token counting
|
| 90 |
+
print("\nCounting pretraining tokens (this may take a while)...")
|
| 91 |
+
pretrain_tokens_3b = count_tokens_text(df_3b["text"].tolist())
|
| 92 |
+
print(f" filtered_3b tokens: {pretrain_tokens_3b:,}")
|
| 93 |
+
|
| 94 |
+
pretrain_tokens_en = count_tokens_text(df_en["text"].tolist())
|
| 95 |
+
print(f" filtered_english tokens: {pretrain_tokens_en:,}")
|
| 96 |
+
|
| 97 |
+
pretrain_total = pretrain_tokens_3b + pretrain_tokens_en
|
| 98 |
+
print(f" TOTAL pretrain tokens: {pretrain_total:,}")
|
| 99 |
+
|
| 100 |
+
|
| 101 |
+
# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 102 |
+
# 2. FINETUNING DATA (finetune + finetune_english, train+val each)
|
| 103 |
+
# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 104 |
+
print("\n" + "=" * 60)
|
| 105 |
+
print("CONSOLIDATING FINETUNING DATA")
|
| 106 |
+
print("=" * 60)
|
| 107 |
+
|
| 108 |
+
finetune_out = ROOT / "Base" / "Datasets" / "consolidated_finetune"
|
| 109 |
+
finetune_out.mkdir(parents=True, exist_ok=True)
|
| 110 |
+
|
| 111 |
+
# Read v1 finetune
|
| 112 |
+
ft_v1_train = json.loads((ROOT / "Base" / "Datasets" / "finetune" / "train.json").read_text(encoding="utf-8"))
|
| 113 |
+
ft_v1_val = json.loads((ROOT / "Base" / "Datasets" / "finetune" / "val.json").read_text(encoding="utf-8"))
|
| 114 |
+
print(f"\nFinetune v1: train={len(ft_v1_train):,} val={len(ft_v1_val):,}")
|
| 115 |
+
|
| 116 |
+
# Read English finetune
|
| 117 |
+
ft_en_train = json.loads((ROOT / "Base" / "Datasets" / "finetune_english" / "train.json").read_text(encoding="utf-8"))
|
| 118 |
+
ft_en_val = json.loads((ROOT / "Base" / "Datasets" / "finetune_english" / "val.json").read_text(encoding="utf-8"))
|
| 119 |
+
print(f"Finetune EN: train={len(ft_en_train):,} val={len(ft_en_val):,}")
|
| 120 |
+
|
| 121 |
+
# Add source tags
|
| 122 |
+
for item in ft_v1_train + ft_v1_val:
|
| 123 |
+
item["source"] = "finetune_v1"
|
| 124 |
+
for item in ft_en_train + ft_en_val:
|
| 125 |
+
item["source"] = "finetune_english"
|
| 126 |
+
|
| 127 |
+
# Combine
|
| 128 |
+
all_finetune = ft_v1_train + ft_v1_val + ft_en_train + ft_en_val
|
| 129 |
+
out_path_ft = finetune_out / "all_finetune_data.json"
|
| 130 |
+
with open(out_path_ft, "w", encoding="utf-8") as f:
|
| 131 |
+
json.dump(all_finetune, f, ensure_ascii=False, indent=2)
|
| 132 |
+
print(f"\nSaved combined finetune: {out_path_ft}")
|
| 133 |
+
print(f" Total samples: {len(all_finetune):,}")
|
| 134 |
+
|
| 135 |
+
# Token counting for finetuning
|
| 136 |
+
print("\nCounting finetuning tokens...")
|
| 137 |
+
|
| 138 |
+
|
| 139 |
+
def finetune_text(item):
|
| 140 |
+
"""Reconstruct the full text that gets tokenized during finetuning."""
|
| 141 |
+
parts = []
|
| 142 |
+
if item.get("instruction"):
|
| 143 |
+
parts.append(item["instruction"])
|
| 144 |
+
if item.get("input"):
|
| 145 |
+
parts.append(item["input"])
|
| 146 |
+
if item.get("output"):
|
| 147 |
+
parts.append(item["output"])
|
| 148 |
+
return " ".join(parts)
|
| 149 |
+
|
| 150 |
+
|
| 151 |
+
ft_v1_texts = [finetune_text(x) for x in ft_v1_train + ft_v1_val]
|
| 152 |
+
ft_en_texts = [finetune_text(x) for x in ft_en_train + ft_en_val]
|
| 153 |
+
|
| 154 |
+
ft_v1_tokens = count_tokens_text(ft_v1_texts)
|
| 155 |
+
print(f" finetune_v1 tokens: {ft_v1_tokens:,}")
|
| 156 |
+
|
| 157 |
+
ft_en_tokens = count_tokens_text(ft_en_texts)
|
| 158 |
+
print(f" finetune_english tokens: {ft_en_tokens:,}")
|
| 159 |
+
|
| 160 |
+
ft_total = ft_v1_tokens + ft_en_tokens
|
| 161 |
+
print(f" TOTAL finetune tokens: {ft_total:,}")
|
| 162 |
+
|
| 163 |
+
|
| 164 |
+
# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 165 |
+
# SUMMARY
|
| 166 |
+
# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 167 |
+
print("\n" + "=" * 60)
|
| 168 |
+
print("FINAL SUMMARY")
|
| 169 |
+
print("=" * 60)
|
| 170 |
+
|
| 171 |
+
summary = f"""
|
| 172 |
+
PRETRAINING DATA (Base/data/consolidated_pretrain/all_pretrain_data.parquet)
|
| 173 |
+
Source: filtered_3b -> {len(df_3b):>10,} docs | {pretrain_tokens_3b:>15,} tokens
|
| 174 |
+
Source: filtered_english -> {len(df_en):>10,} docs | {pretrain_tokens_en:>15,} tokens
|
| 175 |
+
βββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 176 |
+
TOTAL -> {len(df_pretrain):>10,} docs | {pretrain_total:>15,} tokens
|
| 177 |
+
|
| 178 |
+
FINETUNING DATA (Base/Datasets/consolidated_finetune/all_finetune_data.json)
|
| 179 |
+
Source: finetune_v1 -> {len(ft_v1_train)+len(ft_v1_val):>10,} samples | {ft_v1_tokens:>15,} tokens
|
| 180 |
+
Source: finetune_english -> {len(ft_en_train)+len(ft_en_val):>10,} samples | {ft_en_tokens:>15,} tokens
|
| 181 |
+
βββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 182 |
+
TOTAL -> {len(all_finetune):>10,} samples | {ft_total:>15,} tokens
|
| 183 |
+
|
| 184 |
+
GRAND TOTAL TOKENS: {pretrain_total + ft_total:,}
|
| 185 |
+
"""
|
| 186 |
+
print(summary)
|
| 187 |
+
|
| 188 |
+
# Save summary as text file too
|
| 189 |
+
with open(pretrain_out.parent.parent / "data" / "consolidated_pretrain" / "SUMMARY.txt", "w") as f:
|
| 190 |
+
f.write(summary)
|
| 191 |
+
with open(finetune_out / "SUMMARY.txt", "w") as f:
|
| 192 |
+
f.write(summary)
|
| 193 |
+
|
| 194 |
+
print("Summary saved to both consolidated folders.")
|
Base/scripts/consolidate_litdata.py
ADDED
|
@@ -0,0 +1,187 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Consolidate the actual training data (litdata binary chunks) into single folders.
|
| 3 |
+
|
| 4 |
+
Pretraining: litdata_3b + litdata_english -> consolidated_pretrain_litdata/
|
| 5 |
+
Finetuning: finetune/ + finetune_english/ -> consolidated_finetune/
|
| 6 |
+
|
| 7 |
+
Token counts are computed directly from the litdata index.json metadata.
|
| 8 |
+
"""
|
| 9 |
+
|
| 10 |
+
import json
|
| 11 |
+
import shutil
|
| 12 |
+
import os
|
| 13 |
+
from pathlib import Path
|
| 14 |
+
|
| 15 |
+
ROOT = Path(__file__).resolve().parent.parent.parent # LUNA root
|
| 16 |
+
DATA = ROOT / "Base" / "data"
|
| 17 |
+
DATASETS = ROOT / "Base" / "Datasets"
|
| 18 |
+
|
| 19 |
+
|
| 20 |
+
def read_litdata_index(litdata_dir):
|
| 21 |
+
"""Read index.json and return chunks list + config."""
|
| 22 |
+
with open(litdata_dir / "index.json", "r") as f:
|
| 23 |
+
index = json.load(f)
|
| 24 |
+
return index["chunks"], index["config"]
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
def count_tokens_from_index(chunks):
|
| 28 |
+
"""Sum up all tokens from chunk dim fields."""
|
| 29 |
+
return sum(c["dim"] for c in chunks)
|
| 30 |
+
|
| 31 |
+
|
| 32 |
+
# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 33 |
+
# 1. PRETRAINING: litdata_3b + litdata_english
|
| 34 |
+
# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 35 |
+
print("=" * 65)
|
| 36 |
+
print(" PRETRAINING DATA CONSOLIDATION")
|
| 37 |
+
print("=" * 65)
|
| 38 |
+
|
| 39 |
+
# Read both indexes
|
| 40 |
+
chunks_3b, config = read_litdata_index(DATA / "litdata_3b")
|
| 41 |
+
chunks_en, _ = read_litdata_index(DATA / "litdata_english")
|
| 42 |
+
|
| 43 |
+
tokens_3b = count_tokens_from_index(chunks_3b)
|
| 44 |
+
tokens_en = count_tokens_from_index(chunks_en)
|
| 45 |
+
tokens_pretrain = tokens_3b + tokens_en
|
| 46 |
+
|
| 47 |
+
print(f"\n litdata_3b: {len(chunks_3b):>4} chunks | {tokens_3b:>15,} tokens")
|
| 48 |
+
print(f" litdata_english: {len(chunks_en):>4} chunks | {tokens_en:>15,} tokens")
|
| 49 |
+
print(f" {'β' * 55}")
|
| 50 |
+
print(f" TOTAL PRETRAIN: {len(chunks_3b)+len(chunks_en):>4} chunks | {tokens_pretrain:>15,} tokens")
|
| 51 |
+
print(f" ({tokens_pretrain / 1e9:.3f} B tokens)")
|
| 52 |
+
|
| 53 |
+
# Create consolidated litdata folder
|
| 54 |
+
out_pretrain = DATA / "consolidated_pretrain_litdata"
|
| 55 |
+
out_pretrain.mkdir(parents=True, exist_ok=True)
|
| 56 |
+
|
| 57 |
+
# Copy chunks from litdata_3b (keep original names as chunk-0-{0..N})
|
| 58 |
+
print(f"\n Copying litdata_3b chunks ({len(chunks_3b)} files)...")
|
| 59 |
+
new_chunks = []
|
| 60 |
+
for i, chunk in enumerate(chunks_3b):
|
| 61 |
+
src = DATA / "litdata_3b" / chunk["filename"]
|
| 62 |
+
new_name = f"chunk-0-{i}.bin"
|
| 63 |
+
dst = out_pretrain / new_name
|
| 64 |
+
if not dst.exists():
|
| 65 |
+
shutil.copy2(str(src), str(dst))
|
| 66 |
+
new_chunks.append({**chunk, "filename": new_name})
|
| 67 |
+
if (i + 1) % 50 == 0 or i == len(chunks_3b) - 1:
|
| 68 |
+
print(f" {i+1}/{len(chunks_3b)} copied")
|
| 69 |
+
|
| 70 |
+
# Copy chunks from litdata_english (renumber continuing from 3b)
|
| 71 |
+
offset = len(chunks_3b)
|
| 72 |
+
print(f"\n Copying litdata_english chunks ({len(chunks_en)} files)...")
|
| 73 |
+
for i, chunk in enumerate(chunks_en):
|
| 74 |
+
src = DATA / "litdata_english" / chunk["filename"]
|
| 75 |
+
new_name = f"chunk-0-{offset + i}.bin"
|
| 76 |
+
dst = out_pretrain / new_name
|
| 77 |
+
if not dst.exists():
|
| 78 |
+
shutil.copy2(str(src), str(dst))
|
| 79 |
+
new_chunks.append({**chunk, "filename": new_name})
|
| 80 |
+
print(f" {i+1}/{len(chunks_en)} copied")
|
| 81 |
+
|
| 82 |
+
# Write combined index.json
|
| 83 |
+
combined_index = {"chunks": new_chunks, "config": config}
|
| 84 |
+
with open(out_pretrain / "index.json", "w") as f:
|
| 85 |
+
json.dump(combined_index, f, indent=2)
|
| 86 |
+
print(f"\n Saved: {out_pretrain}")
|
| 87 |
+
|
| 88 |
+
|
| 89 |
+
# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 90 |
+
# 2. FINETUNING: finetune/ + finetune_english/
|
| 91 |
+
# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 92 |
+
print(f"\n{'=' * 65}")
|
| 93 |
+
print(" FINETUNING DATA CONSOLIDATION")
|
| 94 |
+
print("=" * 65)
|
| 95 |
+
|
| 96 |
+
# Load all finetune JSONs
|
| 97 |
+
ft_v1_train = json.loads((DATASETS / "finetune" / "train.json").read_text(encoding="utf-8"))
|
| 98 |
+
ft_v1_val = json.loads((DATASETS / "finetune" / "val.json").read_text(encoding="utf-8"))
|
| 99 |
+
ft_en_train = json.loads((DATASETS / "finetune_english" / "train.json").read_text(encoding="utf-8"))
|
| 100 |
+
ft_en_val = json.loads((DATASETS / "finetune_english" / "val.json").read_text(encoding="utf-8"))
|
| 101 |
+
|
| 102 |
+
ft_v1_all = ft_v1_train + ft_v1_val
|
| 103 |
+
ft_en_all = ft_en_train + ft_en_val
|
| 104 |
+
|
| 105 |
+
# Tag sources
|
| 106 |
+
for item in ft_v1_all:
|
| 107 |
+
item["source"] = "finetune_v1"
|
| 108 |
+
for item in ft_en_all:
|
| 109 |
+
item["source"] = "finetune_english"
|
| 110 |
+
|
| 111 |
+
# Count tokens using the same tokenizer
|
| 112 |
+
from tokenizers import Tokenizer
|
| 113 |
+
tokenizer = Tokenizer.from_file(
|
| 114 |
+
str(ROOT / "Base" / "checkpoints" / "EleutherAI" / "pythia-160m" / "tokenizer.json")
|
| 115 |
+
)
|
| 116 |
+
|
| 117 |
+
def finetune_text(item):
|
| 118 |
+
parts = []
|
| 119 |
+
if item.get("instruction"): parts.append(item["instruction"])
|
| 120 |
+
if item.get("input"): parts.append(item["input"])
|
| 121 |
+
if item.get("output"): parts.append(item["output"])
|
| 122 |
+
return " ".join(parts)
|
| 123 |
+
|
| 124 |
+
def count_tokens_batch(texts, label=""):
|
| 125 |
+
total = 0
|
| 126 |
+
BATCH = 10000
|
| 127 |
+
for i in range(0, len(texts), BATCH):
|
| 128 |
+
batch = texts[i:i+BATCH]
|
| 129 |
+
encoded = tokenizer.encode_batch(batch, add_special_tokens=False)
|
| 130 |
+
total += sum(len(e.ids) for e in encoded)
|
| 131 |
+
return total
|
| 132 |
+
|
| 133 |
+
print(f"\n finetune_v1: train={len(ft_v1_train):>7,} val={len(ft_v1_val):>6,} total={len(ft_v1_all):>7,}")
|
| 134 |
+
print(f" finetune_english: train={len(ft_en_train):>7,} val={len(ft_en_val):>6,} total={len(ft_en_all):>7,}")
|
| 135 |
+
|
| 136 |
+
print("\n Counting finetune tokens...")
|
| 137 |
+
ft_v1_tokens = count_tokens_batch([finetune_text(x) for x in ft_v1_all], "v1")
|
| 138 |
+
ft_en_tokens = count_tokens_batch([finetune_text(x) for x in ft_en_all], "english")
|
| 139 |
+
ft_total = ft_v1_tokens + ft_en_tokens
|
| 140 |
+
|
| 141 |
+
print(f"\n finetune_v1 tokens: {ft_v1_tokens:>12,}")
|
| 142 |
+
print(f" finetune_english tokens: {ft_en_tokens:>12,}")
|
| 143 |
+
print(f" {'β' * 55}")
|
| 144 |
+
print(f" TOTAL FINETUNE: {ft_total:>12,} tokens")
|
| 145 |
+
|
| 146 |
+
# Save consolidated finetune
|
| 147 |
+
out_finetune = DATASETS / "consolidated_finetune"
|
| 148 |
+
out_finetune.mkdir(parents=True, exist_ok=True)
|
| 149 |
+
|
| 150 |
+
all_finetune = ft_v1_all + ft_en_all
|
| 151 |
+
with open(out_finetune / "all_finetune_data.json", "w", encoding="utf-8") as f:
|
| 152 |
+
json.dump(all_finetune, f, ensure_ascii=False, indent=2)
|
| 153 |
+
print(f"\n Saved: {out_finetune / 'all_finetune_data.json'}")
|
| 154 |
+
print(f" Total samples: {len(all_finetune):,}")
|
| 155 |
+
|
| 156 |
+
|
| 157 |
+
# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 158 |
+
# FINAL SUMMARY
|
| 159 |
+
# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 160 |
+
print(f"\n{'=' * 65}")
|
| 161 |
+
print(" FINAL SUMMARY")
|
| 162 |
+
print("=" * 65)
|
| 163 |
+
|
| 164 |
+
summary = f"""
|
| 165 |
+
PRETRAINING (Base/data/consolidated_pretrain_litdata/)
|
| 166 |
+
litdata_3b : {len(chunks_3b):>4} chunks | {tokens_3b:>15,} tokens
|
| 167 |
+
litdata_english : {len(chunks_en):>4} chunks | {tokens_en:>15,} tokens
|
| 168 |
+
βββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 169 |
+
TOTAL : {len(chunks_3b)+len(chunks_en):>4} chunks | {tokens_pretrain:>15,} tokens ({tokens_pretrain/1e9:.3f}B)
|
| 170 |
+
|
| 171 |
+
FINETUNING (Base/Datasets/consolidated_finetune/)
|
| 172 |
+
finetune_v1 : {len(ft_v1_all):>7,} samples | {ft_v1_tokens:>12,} tokens
|
| 173 |
+
finetune_english : {len(ft_en_all):>7,} samples | {ft_en_tokens:>12,} tokens
|
| 174 |
+
βββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 175 |
+
TOTAL : {len(all_finetune):>7,} samples | {ft_total:>12,} tokens
|
| 176 |
+
|
| 177 |
+
GRAND TOTAL TOKENS: {tokens_pretrain + ft_total:,}
|
| 178 |
+
"""
|
| 179 |
+
print(summary)
|
| 180 |
+
|
| 181 |
+
# Save summary
|
| 182 |
+
with open(out_pretrain / "SUMMARY.txt", "w", encoding="utf-8") as f:
|
| 183 |
+
f.write(summary)
|
| 184 |
+
with open(out_finetune / "SUMMARY.txt", "w", encoding="utf-8") as f:
|
| 185 |
+
f.write(summary)
|
| 186 |
+
|
| 187 |
+
print("Done! Summary saved to both consolidated folders.")
|
Base/scripts/dedup_5b_pretrain.py
ADDED
|
@@ -0,0 +1,298 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""
|
| 3 |
+
DEDUPLICATE litdata_pretrain_final β remove near-duplicate documents.
|
| 4 |
+
|
| 5 |
+
Strategy:
|
| 6 |
+
1. Scan all 308 chunks, hash each document (first 200 tokens)
|
| 7 |
+
2. Keep first occurrence, mark subsequent duplicates for removal
|
| 8 |
+
3. Rebuild chunks with duplicates removed (same format, new files)
|
| 9 |
+
4. Update index.json
|
| 10 |
+
|
| 11 |
+
Memory-efficient: processes 10 chunks at a time, uses hash set.
|
| 12 |
+
"""
|
| 13 |
+
|
| 14 |
+
import json
|
| 15 |
+
import os
|
| 16 |
+
import time
|
| 17 |
+
import hashlib
|
| 18 |
+
from pathlib import Path
|
| 19 |
+
|
| 20 |
+
import numpy as np
|
| 21 |
+
|
| 22 |
+
ROOT = Path(__file__).resolve().parent.parent.parent
|
| 23 |
+
FINAL_DIR = ROOT / "Base" / "data" / "litdata_pretrain_final"
|
| 24 |
+
BLOCK_SIZE = 1025
|
| 25 |
+
DTYPE = np.int32
|
| 26 |
+
EOS_TOKEN_ID = 0
|
| 27 |
+
CHUNK_BYTES_TARGET = 64 * 1024 * 1024
|
| 28 |
+
HASH_WINDOW = 200 # first N tokens for hash
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
def read_chunk(filepath):
|
| 32 |
+
with open(filepath, "rb") as f:
|
| 33 |
+
raw = f.read()
|
| 34 |
+
num_blocks = np.frombuffer(raw[:4], dtype=np.uint32)[0]
|
| 35 |
+
header_size = 4 + (num_blocks + 1) * 4
|
| 36 |
+
data_bytes = raw[header_size:]
|
| 37 |
+
expected_tokens = num_blocks * BLOCK_SIZE
|
| 38 |
+
expected_bytes = expected_tokens * DTYPE().itemsize
|
| 39 |
+
tokens = np.frombuffer(data_bytes[:expected_bytes], dtype=DTYPE)
|
| 40 |
+
return tokens, int(num_blocks)
|
| 41 |
+
|
| 42 |
+
|
| 43 |
+
def extract_documents(tokens):
|
| 44 |
+
"""Extract individual documents separated by EOS."""
|
| 45 |
+
eos_positions = np.where(tokens == EOS_TOKEN_ID)[0]
|
| 46 |
+
docs = []
|
| 47 |
+
start = 0
|
| 48 |
+
for eos_pos in eos_positions:
|
| 49 |
+
if eos_pos > start:
|
| 50 |
+
docs.append(tokens[start:eos_pos])
|
| 51 |
+
start = eos_pos + 1
|
| 52 |
+
# Trailing partial (no EOS at end β crosses chunk boundary)
|
| 53 |
+
if start < len(tokens):
|
| 54 |
+
remaining = tokens[start:]
|
| 55 |
+
if len(remaining) > 0:
|
| 56 |
+
docs.append(remaining)
|
| 57 |
+
return docs
|
| 58 |
+
|
| 59 |
+
|
| 60 |
+
def doc_hash(token_array):
|
| 61 |
+
"""Hash first HASH_WINDOW tokens of a document."""
|
| 62 |
+
key = token_array[:HASH_WINDOW].tobytes()
|
| 63 |
+
return hashlib.md5(key).hexdigest()
|
| 64 |
+
|
| 65 |
+
|
| 66 |
+
def write_chunk(filepath, tokens_array, block_size=BLOCK_SIZE):
|
| 67 |
+
"""Write a litdata chunk from a flat token array."""
|
| 68 |
+
num_blocks = len(tokens_array) // block_size
|
| 69 |
+
if num_blocks == 0:
|
| 70 |
+
return None
|
| 71 |
+
actual = num_blocks * block_size
|
| 72 |
+
data = np.array(tokens_array[:actual], dtype=DTYPE)
|
| 73 |
+
|
| 74 |
+
header_num = np.array([num_blocks], dtype=np.uint32)
|
| 75 |
+
offsets = np.arange(num_blocks + 1, dtype=np.uint32) * (block_size * DTYPE().itemsize)
|
| 76 |
+
header = np.concatenate([header_num, offsets])
|
| 77 |
+
|
| 78 |
+
with open(filepath, "wb") as f:
|
| 79 |
+
header.tofile(f)
|
| 80 |
+
data.tofile(f)
|
| 81 |
+
|
| 82 |
+
return {
|
| 83 |
+
"chunk_bytes": int(header.nbytes + data.nbytes),
|
| 84 |
+
"chunk_size": num_blocks,
|
| 85 |
+
"dim": int(actual),
|
| 86 |
+
"filename": os.path.basename(filepath),
|
| 87 |
+
}
|
| 88 |
+
|
| 89 |
+
|
| 90 |
+
class StreamingDeduplicator:
|
| 91 |
+
"""Accumulates deduplicated tokens and writes chunks."""
|
| 92 |
+
|
| 93 |
+
def __init__(self, output_dir, backup_suffix="_dedup"):
|
| 94 |
+
self.output_dir = Path(output_dir)
|
| 95 |
+
self.dtype_size = DTYPE().itemsize
|
| 96 |
+
self.tokens_per_chunk = (CHUNK_BYTES_TARGET // self.dtype_size // BLOCK_SIZE) * BLOCK_SIZE
|
| 97 |
+
self.buffer = []
|
| 98 |
+
self.chunk_idx = 0
|
| 99 |
+
self.chunks_meta = []
|
| 100 |
+
self.total_tokens = 0
|
| 101 |
+
|
| 102 |
+
def add_doc(self, doc_tokens):
|
| 103 |
+
self.buffer.extend(doc_tokens.tolist())
|
| 104 |
+
self.buffer.append(EOS_TOKEN_ID)
|
| 105 |
+
while len(self.buffer) >= self.tokens_per_chunk:
|
| 106 |
+
self._flush()
|
| 107 |
+
|
| 108 |
+
def _flush(self):
|
| 109 |
+
if len(self.buffer) < BLOCK_SIZE:
|
| 110 |
+
return
|
| 111 |
+
take = min(len(self.buffer), self.tokens_per_chunk)
|
| 112 |
+
num_blocks = take // BLOCK_SIZE
|
| 113 |
+
if num_blocks == 0:
|
| 114 |
+
return
|
| 115 |
+
actual = num_blocks * BLOCK_SIZE
|
| 116 |
+
|
| 117 |
+
data = np.array(self.buffer[:actual], dtype=DTYPE)
|
| 118 |
+
self.buffer = self.buffer[actual:]
|
| 119 |
+
|
| 120 |
+
filename = f"chunk-0-{self.chunk_idx}.bin"
|
| 121 |
+
filepath = self.output_dir / filename
|
| 122 |
+
|
| 123 |
+
header_num = np.array([num_blocks], dtype=np.uint32)
|
| 124 |
+
offsets = np.arange(num_blocks + 1, dtype=np.uint32) * (BLOCK_SIZE * self.dtype_size)
|
| 125 |
+
header = np.concatenate([header_num, offsets])
|
| 126 |
+
|
| 127 |
+
with open(filepath, "wb") as f:
|
| 128 |
+
header.tofile(f)
|
| 129 |
+
data.tofile(f)
|
| 130 |
+
|
| 131 |
+
meta = {
|
| 132 |
+
"chunk_bytes": int(header.nbytes + data.nbytes),
|
| 133 |
+
"chunk_size": num_blocks,
|
| 134 |
+
"dim": int(actual),
|
| 135 |
+
"filename": filename,
|
| 136 |
+
}
|
| 137 |
+
self.chunks_meta.append(meta)
|
| 138 |
+
self.total_tokens += actual
|
| 139 |
+
self.chunk_idx += 1
|
| 140 |
+
if self.chunk_idx % 25 == 0:
|
| 141 |
+
print(f" Written {self.chunk_idx} deduped chunks ({self.total_tokens:,} tokens)")
|
| 142 |
+
|
| 143 |
+
def finalize(self):
|
| 144 |
+
while len(self.buffer) >= BLOCK_SIZE:
|
| 145 |
+
self._flush()
|
| 146 |
+
discarded = len(self.buffer)
|
| 147 |
+
self.buffer = []
|
| 148 |
+
return self.total_tokens, discarded
|
| 149 |
+
|
| 150 |
+
|
| 151 |
+
def main():
|
| 152 |
+
t_start = time.time()
|
| 153 |
+
|
| 154 |
+
with open(FINAL_DIR / "index.json") as f:
|
| 155 |
+
index = json.load(f)
|
| 156 |
+
chunks_meta = index["chunks"]
|
| 157 |
+
config = index.get("config", {})
|
| 158 |
+
num_chunks = len(chunks_meta)
|
| 159 |
+
|
| 160 |
+
original_tokens = sum(c["dim"] for c in chunks_meta)
|
| 161 |
+
|
| 162 |
+
print(f"{'='*75}")
|
| 163 |
+
print(f" DEDUPLICATING litdata_pretrain_final")
|
| 164 |
+
print(f"{'='*75}")
|
| 165 |
+
print(f" Chunks: {num_chunks}")
|
| 166 |
+
print(f" Original tokens: {original_tokens:,}")
|
| 167 |
+
print()
|
| 168 |
+
|
| 169 |
+
# ββ PASS 1: Scan all chunks, collect hashes, identify duplicates ββββββ
|
| 170 |
+
print(f" PASS 1: Scanning all chunks for duplicates...")
|
| 171 |
+
seen_hashes = set()
|
| 172 |
+
dup_count = 0
|
| 173 |
+
keep_count = 0
|
| 174 |
+
total_docs = 0
|
| 175 |
+
# We need to track which is first occurrence
|
| 176 |
+
# But since we process sequentially, just check "in seen_hashes"
|
| 177 |
+
|
| 178 |
+
# Use a temp dir to write deduplicated data, then swap
|
| 179 |
+
TEMP_DIR = FINAL_DIR.parent / "litdata_pretrain_dedup_temp"
|
| 180 |
+
if TEMP_DIR.exists():
|
| 181 |
+
import shutil
|
| 182 |
+
shutil.rmtree(str(TEMP_DIR))
|
| 183 |
+
os.makedirs(str(TEMP_DIR))
|
| 184 |
+
|
| 185 |
+
writer = StreamingDeduplicator(TEMP_DIR)
|
| 186 |
+
|
| 187 |
+
for ci, meta in enumerate(chunks_meta):
|
| 188 |
+
filepath = FINAL_DIR / meta["filename"]
|
| 189 |
+
tokens, num_blocks = read_chunk(filepath)
|
| 190 |
+
docs = extract_documents(tokens)
|
| 191 |
+
|
| 192 |
+
chunk_dups = 0
|
| 193 |
+
chunk_kept = 0
|
| 194 |
+
|
| 195 |
+
for doc in docs:
|
| 196 |
+
total_docs += 1
|
| 197 |
+
if len(doc) < 10:
|
| 198 |
+
# Very short fragments β keep (usually chunk-boundary partials)
|
| 199 |
+
writer.add_doc(doc)
|
| 200 |
+
keep_count += 1
|
| 201 |
+
chunk_kept += 1
|
| 202 |
+
continue
|
| 203 |
+
|
| 204 |
+
h = doc_hash(doc)
|
| 205 |
+
if h in seen_hashes:
|
| 206 |
+
dup_count += 1
|
| 207 |
+
chunk_dups += 1
|
| 208 |
+
else:
|
| 209 |
+
seen_hashes.add(h)
|
| 210 |
+
writer.add_doc(doc)
|
| 211 |
+
keep_count += 1
|
| 212 |
+
chunk_kept += 1
|
| 213 |
+
|
| 214 |
+
if (ci + 1) % 25 == 0 or ci == num_chunks - 1:
|
| 215 |
+
print(f" Chunk {ci+1}/{num_chunks}: total docs={total_docs:,}, kept={keep_count:,}, dupes removed={dup_count:,}")
|
| 216 |
+
|
| 217 |
+
# Finalize
|
| 218 |
+
final_tokens, discarded = writer.finalize()
|
| 219 |
+
new_chunks = writer.chunk_idx
|
| 220 |
+
|
| 221 |
+
print(f"\n PASS 1 COMPLETE:")
|
| 222 |
+
print(f" Total documents scanned: {total_docs:,}")
|
| 223 |
+
print(f" Documents kept: {keep_count:,}")
|
| 224 |
+
print(f" Duplicates removed: {dup_count:,} ({100*dup_count/max(total_docs,1):.2f}%)")
|
| 225 |
+
print(f" Unique hashes: {len(seen_hashes):,}")
|
| 226 |
+
print(f" Tokens after dedup: {final_tokens:,}")
|
| 227 |
+
print(f" Token reduction: {original_tokens - final_tokens:,} ({100*(original_tokens-final_tokens)/original_tokens:.2f}%)")
|
| 228 |
+
print(f" Chunks after dedup: {new_chunks}")
|
| 229 |
+
print(f" Discarded partial: {discarded} tokens")
|
| 230 |
+
|
| 231 |
+
# ββ PASS 2: Swap temp into final ββββββββββββββββββββββββββββββββββββββ
|
| 232 |
+
print(f"\n PASS 2: Replacing original with deduplicated data...")
|
| 233 |
+
|
| 234 |
+
# Remove old chunk files
|
| 235 |
+
for meta in chunks_meta:
|
| 236 |
+
old_file = FINAL_DIR / meta["filename"]
|
| 237 |
+
if old_file.exists():
|
| 238 |
+
os.remove(str(old_file))
|
| 239 |
+
|
| 240 |
+
# Move new chunk files from temp to final
|
| 241 |
+
import shutil
|
| 242 |
+
for meta in writer.chunks_meta:
|
| 243 |
+
src = TEMP_DIR / meta["filename"]
|
| 244 |
+
dst = FINAL_DIR / meta["filename"]
|
| 245 |
+
shutil.move(str(src), str(dst))
|
| 246 |
+
|
| 247 |
+
# Remove temp dir
|
| 248 |
+
shutil.rmtree(str(TEMP_DIR))
|
| 249 |
+
|
| 250 |
+
# Update index.json
|
| 251 |
+
new_index = {
|
| 252 |
+
"chunks": writer.chunks_meta,
|
| 253 |
+
"config": config,
|
| 254 |
+
"updated_at": str(time.time()),
|
| 255 |
+
}
|
| 256 |
+
with open(FINAL_DIR / "index.json", "w") as f:
|
| 257 |
+
json.dump(new_index, f, indent=2)
|
| 258 |
+
|
| 259 |
+
elapsed = time.time() - t_start
|
| 260 |
+
|
| 261 |
+
# ββ Report ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 262 |
+
report = []
|
| 263 |
+
report.append(f"{'='*75}")
|
| 264 |
+
report.append(f" DEDUPLICATION REPORT β litdata_pretrain_final")
|
| 265 |
+
report.append(f"{'='*75}")
|
| 266 |
+
report.append(f"")
|
| 267 |
+
report.append(f" Time: {elapsed:.0f}s ({elapsed/60:.1f} min)")
|
| 268 |
+
report.append(f"")
|
| 269 |
+
report.append(f" BEFORE:")
|
| 270 |
+
report.append(f" Chunks: {num_chunks}")
|
| 271 |
+
report.append(f" Tokens: {original_tokens:,}")
|
| 272 |
+
report.append(f" Documents: {total_docs:,}")
|
| 273 |
+
report.append(f"")
|
| 274 |
+
report.append(f" AFTER:")
|
| 275 |
+
report.append(f" Chunks: {new_chunks}")
|
| 276 |
+
report.append(f" Tokens: {final_tokens:,}")
|
| 277 |
+
report.append(f" Documents: {keep_count:,}")
|
| 278 |
+
report.append(f"")
|
| 279 |
+
report.append(f" REMOVED:")
|
| 280 |
+
report.append(f" Duplicate docs: {dup_count:,} ({100*dup_count/max(total_docs,1):.2f}%)")
|
| 281 |
+
report.append(f" Tokens removed: {original_tokens - final_tokens:,} ({100*(original_tokens-final_tokens)/original_tokens:.2f}%)")
|
| 282 |
+
report.append(f"")
|
| 283 |
+
report.append(f" Format: litdata binary (int32, BLOCK_SIZE={BLOCK_SIZE}, EOS={EOS_TOKEN_ID})")
|
| 284 |
+
report.append(f" Location: {FINAL_DIR}")
|
| 285 |
+
report.append(f"{'='*75}")
|
| 286 |
+
|
| 287 |
+
full_report = '\n'.join(report)
|
| 288 |
+
print(f"\n{full_report}")
|
| 289 |
+
|
| 290 |
+
with open(FINAL_DIR / "DEDUP_REPORT.txt", "w", encoding="utf-8") as f:
|
| 291 |
+
f.write(full_report)
|
| 292 |
+
|
| 293 |
+
print(f"\n Saved to: {FINAL_DIR / 'DEDUP_REPORT.txt'}")
|
| 294 |
+
print(f" Done! Dataset is now clean and deduplicated.")
|
| 295 |
+
|
| 296 |
+
|
| 297 |
+
if __name__ == "__main__":
|
| 298 |
+
main()
|
Base/scripts/filter_datasets.py
ADDED
|
@@ -0,0 +1,151 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Filter parquet datasets based on language_score and token_count.
|
| 3 |
+
Creates two filtered sets:
|
| 4 |
+
- Small (~10M tokens) for quick test runs
|
| 5 |
+
- Full (~3B tokens) for full pretraining
|
| 6 |
+
|
| 7 |
+
Usage:
|
| 8 |
+
python Base/scripts/filter_datasets.py --mode small # ~10M tokens
|
| 9 |
+
python Base/scripts/filter_datasets.py --mode full # ~3B tokens
|
| 10 |
+
python Base/scripts/filter_datasets.py --mode both # both sets
|
| 11 |
+
"""
|
| 12 |
+
|
| 13 |
+
import argparse
|
| 14 |
+
import os
|
| 15 |
+
import time
|
| 16 |
+
from pathlib import Path
|
| 17 |
+
|
| 18 |
+
import duckdb
|
| 19 |
+
import pyarrow as pa
|
| 20 |
+
import pyarrow.parquet as pq
|
| 21 |
+
|
| 22 |
+
|
| 23 |
+
def filter_datasets(input_dir: str, output_dir: str, max_tokens: int, label: str):
|
| 24 |
+
"""Filter parquet files with SQL query and save until max_tokens reached."""
|
| 25 |
+
os.makedirs(output_dir, exist_ok=True)
|
| 26 |
+
|
| 27 |
+
parquet_files = sorted(Path(input_dir).glob("*.parquet"))
|
| 28 |
+
if not parquet_files:
|
| 29 |
+
print(f"No parquet files found in {input_dir}")
|
| 30 |
+
return 0
|
| 31 |
+
|
| 32 |
+
print(f"\n{'='*60}")
|
| 33 |
+
print(f"Filtering for [{label}] set β target: {max_tokens:,} tokens")
|
| 34 |
+
print(f"Source: {input_dir} ({len(parquet_files)} files)")
|
| 35 |
+
print(f"Output: {output_dir}")
|
| 36 |
+
print(f"{'='*60}\n")
|
| 37 |
+
|
| 38 |
+
conn = duckdb.connect()
|
| 39 |
+
|
| 40 |
+
total_tokens = 0
|
| 41 |
+
total_rows = 0
|
| 42 |
+
file_idx = 0
|
| 43 |
+
start_time = time.time()
|
| 44 |
+
|
| 45 |
+
for pf in parquet_files:
|
| 46 |
+
if total_tokens >= max_tokens:
|
| 47 |
+
break
|
| 48 |
+
|
| 49 |
+
# Use parameterized path for safety β duckdb read_parquet needs direct path
|
| 50 |
+
pf_str = str(pf).replace("\\", "/")
|
| 51 |
+
query = f"""
|
| 52 |
+
SELECT text, token_count
|
| 53 |
+
FROM read_parquet('{pf_str}')
|
| 54 |
+
WHERE language_score >= 0.96
|
| 55 |
+
AND token_count < 4096
|
| 56 |
+
AND token_count > 64
|
| 57 |
+
"""
|
| 58 |
+
|
| 59 |
+
df = conn.execute(query).fetchdf()
|
| 60 |
+
|
| 61 |
+
if df.empty:
|
| 62 |
+
print(f" {pf.name}: 0 qualifying rows, skipping")
|
| 63 |
+
continue
|
| 64 |
+
|
| 65 |
+
# Calculate how many rows we can take
|
| 66 |
+
remaining = max_tokens - total_tokens
|
| 67 |
+
cumsum = df["token_count"].cumsum()
|
| 68 |
+
cutoff_mask = cumsum <= remaining
|
| 69 |
+
|
| 70 |
+
if cutoff_mask.sum() == 0:
|
| 71 |
+
# Even one row exceeds remaining β take just one row to finish
|
| 72 |
+
df_subset = df.iloc[:1]
|
| 73 |
+
else:
|
| 74 |
+
df_subset = df[cutoff_mask]
|
| 75 |
+
|
| 76 |
+
subset_tokens = int(df_subset["token_count"].sum())
|
| 77 |
+
total_tokens += subset_tokens
|
| 78 |
+
total_rows += len(df_subset)
|
| 79 |
+
|
| 80 |
+
# Save only the 'text' column as parquet
|
| 81 |
+
output_path = Path(output_dir) / f"filtered_{file_idx:04d}.parquet"
|
| 82 |
+
table = pa.table({"text": df_subset["text"].values})
|
| 83 |
+
pq.write_table(table, str(output_path))
|
| 84 |
+
file_idx += 1
|
| 85 |
+
|
| 86 |
+
print(
|
| 87 |
+
f" {pf.name}: {len(df)} qualifying β kept {len(df_subset)} rows "
|
| 88 |
+
f"({subset_tokens:,} tokens) | cumulative: {total_tokens:,}"
|
| 89 |
+
)
|
| 90 |
+
|
| 91 |
+
conn.close()
|
| 92 |
+
elapsed = time.time() - start_time
|
| 93 |
+
|
| 94 |
+
print(f"\n--- {label} filtering complete ---")
|
| 95 |
+
print(f"Total rows: {total_rows:,}")
|
| 96 |
+
print(f"Total tokens: {total_tokens:,}")
|
| 97 |
+
print(f"Output files: {file_idx}")
|
| 98 |
+
print(f"Time: {elapsed:.1f}s")
|
| 99 |
+
|
| 100 |
+
return total_tokens
|
| 101 |
+
|
| 102 |
+
|
| 103 |
+
def main():
|
| 104 |
+
parser = argparse.ArgumentParser(description="Filter parquet datasets for pretraining")
|
| 105 |
+
parser.add_argument(
|
| 106 |
+
"--input_dir",
|
| 107 |
+
type=str,
|
| 108 |
+
default="Base/Datasets",
|
| 109 |
+
help="Directory containing source parquet files",
|
| 110 |
+
)
|
| 111 |
+
parser.add_argument(
|
| 112 |
+
"--mode",
|
| 113 |
+
choices=["small", "full", "both"],
|
| 114 |
+
default="small",
|
| 115 |
+
help="Which dataset to create: small (~10M tokens), full (~3B tokens), or both",
|
| 116 |
+
)
|
| 117 |
+
parser.add_argument(
|
| 118 |
+
"--small_tokens",
|
| 119 |
+
type=int,
|
| 120 |
+
default=10_000_000,
|
| 121 |
+
help="Target token count for small test set (default: 10M)",
|
| 122 |
+
)
|
| 123 |
+
parser.add_argument(
|
| 124 |
+
"--full_tokens",
|
| 125 |
+
type=int,
|
| 126 |
+
default=3_000_000_000,
|
| 127 |
+
help="Target token count for full training set (default: 3B)",
|
| 128 |
+
)
|
| 129 |
+
args = parser.parse_args()
|
| 130 |
+
|
| 131 |
+
if args.mode in ("small", "both"):
|
| 132 |
+
filter_datasets(
|
| 133 |
+
input_dir=args.input_dir,
|
| 134 |
+
output_dir="Base/data/filtered_10m",
|
| 135 |
+
max_tokens=args.small_tokens,
|
| 136 |
+
label="SMALL (10M)",
|
| 137 |
+
)
|
| 138 |
+
|
| 139 |
+
if args.mode in ("full", "both"):
|
| 140 |
+
filter_datasets(
|
| 141 |
+
input_dir=args.input_dir,
|
| 142 |
+
output_dir="Base/data/filtered_3b",
|
| 143 |
+
max_tokens=args.full_tokens,
|
| 144 |
+
label="FULL (3B)",
|
| 145 |
+
)
|
| 146 |
+
|
| 147 |
+
print("\nDone! Next step: run prepare_litdata.py to tokenize and create LitData format.")
|
| 148 |
+
|
| 149 |
+
|
| 150 |
+
if __name__ == "__main__":
|
| 151 |
+
main()
|
Base/scripts/merge_and_build.py
ADDED
|
@@ -0,0 +1,701 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""
|
| 3 |
+
Merge existing litdata datasets + build a NEW 100M token batch + merge all.
|
| 4 |
+
|
| 5 |
+
Steps:
|
| 6 |
+
1. Merge litdata_english_100m + litdata_english_clean β litdata_combined
|
| 7 |
+
2. Download 100M tokens of FRESH data (skipping everything already used)
|
| 8 |
+
3. Deep clean + smart cleanup
|
| 9 |
+
4. Tokenize β temp litdata chunks
|
| 10 |
+
5. Merge temp chunks into litdata_combined β final unified dataset
|
| 11 |
+
|
| 12 |
+
Skip values:
|
| 13 |
+
- WIKI_SKIP = 80,000 (first dataset used ~19.5K, second used 25K+53.5K = 78.5K)
|
| 14 |
+
- FINEWEB_SKIP = 70,000 (first used ~19.4K, second used 25K+43.5K = 68.5K)
|
| 15 |
+
|
| 16 |
+
Final output: Base/data/litdata_combined/
|
| 17 |
+
"""
|
| 18 |
+
|
| 19 |
+
import json
|
| 20 |
+
import os
|
| 21 |
+
import re
|
| 22 |
+
import shutil
|
| 23 |
+
import time
|
| 24 |
+
import unicodedata
|
| 25 |
+
from pathlib import Path
|
| 26 |
+
from collections import Counter
|
| 27 |
+
|
| 28 |
+
import numpy as np
|
| 29 |
+
import pyarrow as pa
|
| 30 |
+
import pyarrow.parquet as pq
|
| 31 |
+
from tokenizers import Tokenizer
|
| 32 |
+
|
| 33 |
+
ROOT = Path(__file__).resolve().parent.parent.parent
|
| 34 |
+
BLOCK_SIZE = 1025
|
| 35 |
+
DTYPE = np.int32
|
| 36 |
+
CHUNK_BYTES_TARGET = 64 * 1024 * 1024
|
| 37 |
+
EOS_TOKEN_ID = 0
|
| 38 |
+
TARGET_TOKENS = 100_000_000
|
| 39 |
+
TOKENS_PER_WORD = 1.3
|
| 40 |
+
|
| 41 |
+
# Skip counts β must exceed ALL previously used qualifying articles
|
| 42 |
+
WIKI_SKIP = 80_000
|
| 43 |
+
FINEWEB_SKIP = 70_000
|
| 44 |
+
|
| 45 |
+
WIKI_SHARE = 0.55
|
| 46 |
+
FINEWEB_MIN_SCORE = 4.0
|
| 47 |
+
|
| 48 |
+
# Paths
|
| 49 |
+
DATA_DIR = ROOT / "Base" / "data"
|
| 50 |
+
LITDATA_100M = DATA_DIR / "litdata_english_100m"
|
| 51 |
+
LITDATA_CLEAN = DATA_DIR / "litdata_english_clean"
|
| 52 |
+
COMBINED_DIR = DATA_DIR / "litdata_combined"
|
| 53 |
+
TEMP_DIR = DATA_DIR / "litdata_new_100m_temp"
|
| 54 |
+
PARQUET_DIR = DATA_DIR / "filtered_english_new_100m"
|
| 55 |
+
TOKENIZER_PATH = ROOT / "Base" / "checkpoints" / "EleutherAI" / "pythia-160m" / "tokenizer.json"
|
| 56 |
+
|
| 57 |
+
print("Loading tokenizer...")
|
| 58 |
+
tokenizer = Tokenizer.from_file(str(TOKENIZER_PATH))
|
| 59 |
+
|
| 60 |
+
|
| 61 |
+
# ==============================================================================
|
| 62 |
+
# TEXT QUALITY FILTERS
|
| 63 |
+
# ==============================================================================
|
| 64 |
+
|
| 65 |
+
def clean_text(text):
|
| 66 |
+
text = unicodedata.normalize("NFKC", text)
|
| 67 |
+
text = re.sub(r'[\x00-\x08\x0b\x0c\x0e-\x1f\x7f-\x9f]', '', text)
|
| 68 |
+
text = re.sub(r'\n{3,}', '\n\n', text)
|
| 69 |
+
text = re.sub(r'[ \t]+', ' ', text)
|
| 70 |
+
text = '\n'.join(line.strip() for line in text.split('\n'))
|
| 71 |
+
return text.strip()
|
| 72 |
+
|
| 73 |
+
|
| 74 |
+
def is_high_quality(text, min_chars=500, min_words=80):
|
| 75 |
+
if len(text) < min_chars:
|
| 76 |
+
return False
|
| 77 |
+
words = text.split()
|
| 78 |
+
num_words = len(words)
|
| 79 |
+
if num_words < min_words:
|
| 80 |
+
return False
|
| 81 |
+
alpha = sum(c.isalpha() for c in text)
|
| 82 |
+
if alpha / max(len(text), 1) < 0.65:
|
| 83 |
+
return False
|
| 84 |
+
avg_word_len = sum(len(w) for w in words) / num_words
|
| 85 |
+
if avg_word_len < 2.5 or avg_word_len > 15:
|
| 86 |
+
return False
|
| 87 |
+
url_hits = text.count('http://') + text.count('https://')
|
| 88 |
+
if url_hits > num_words * 0.03:
|
| 89 |
+
return False
|
| 90 |
+
sentences = re.split(r'[.!?]+', text)
|
| 91 |
+
real_sentences = [s.strip() for s in sentences if len(s.strip()) > 10]
|
| 92 |
+
if len(real_sentences) < 3:
|
| 93 |
+
return False
|
| 94 |
+
lines = [ln.strip() for ln in text.split('\n') if ln.strip()]
|
| 95 |
+
if len(lines) > 5:
|
| 96 |
+
unique_ratio = len(set(lines)) / len(lines)
|
| 97 |
+
if unique_ratio < 0.5:
|
| 98 |
+
return False
|
| 99 |
+
return True
|
| 100 |
+
|
| 101 |
+
|
| 102 |
+
_WIKI_SKIP_PATTERNS = re.compile(
|
| 103 |
+
r'(disambiguation|list of|lists of|index of|outline of|'
|
| 104 |
+
r'wikipedia:|template:|category:|portal:|module:|mediawiki:)',
|
| 105 |
+
re.IGNORECASE
|
| 106 |
+
)
|
| 107 |
+
|
| 108 |
+
|
| 109 |
+
# ==============================================================================
|
| 110 |
+
# DEEP CLEANING PIPELINE
|
| 111 |
+
# ==============================================================================
|
| 112 |
+
|
| 113 |
+
CONTROL_CHARS = [
|
| 114 |
+
"\x00", "\x01", "\x02", "\x03", "\x04", "\x05", "\x06", "\x07",
|
| 115 |
+
"\x08", "\x0b", "\x0c", "\x0e", "\x0f", "\x10", "\x11", "\x12",
|
| 116 |
+
"\x13", "\x14", "\x15", "\x16", "\x17", "\x18", "\x19", "\x1a",
|
| 117 |
+
"\x1b", "\x1c", "\x1d", "\x1e", "\x1f", "\x7f", "\ufeff", "\ufffd",
|
| 118 |
+
]
|
| 119 |
+
|
| 120 |
+
HTML_ENTITIES = [
|
| 121 |
+
("&", "&"), ("<", "<"), (">", ">"),
|
| 122 |
+
(""", '"'), ("'", "'"), ("'", "'"),
|
| 123 |
+
(" ", " "), ("—", " - "), ("–", "-"),
|
| 124 |
+
("…", "..."), ("«", '"'), ("»", '"'),
|
| 125 |
+
("•", "- "), ("·", " "), ("©", "(c)"),
|
| 126 |
+
("®", "(R)"), ("™", "(TM)"), ("°", " degrees"),
|
| 127 |
+
]
|
| 128 |
+
|
| 129 |
+
RE_URL = re.compile(r'https?://\S+|www\.\S+', re.I)
|
| 130 |
+
RE_EMAIL = re.compile(r'\b[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}\b')
|
| 131 |
+
RE_FILE_PATH = re.compile(r'(?:[A-Z]:\\|/(?:home|usr|var|etc|opt)/)\S+')
|
| 132 |
+
RE_HTML_TAG = re.compile(r'</?[a-zA-Z][a-zA-Z0-9]*(?:\s[^>]*)?\s*/?>')
|
| 133 |
+
RE_HTML_COMMENT = re.compile(r'<!--.*?-->', re.DOTALL)
|
| 134 |
+
RE_CODE_BLOCK = re.compile(r'```[\s\S]*?```')
|
| 135 |
+
RE_IMPORT = re.compile(r'^(?:import |from \S+ import |#include |using namespace |require\()', re.M)
|
| 136 |
+
RE_REPEATED_LINE = re.compile(r'^(.{20,})\n(?:\1\n?)+', re.M)
|
| 137 |
+
RE_REPEATED_PUNCT = re.compile(r'([!?.])\1{3,}')
|
| 138 |
+
RE_REPEATED_CHAR = re.compile(r'(.)\1{5,}')
|
| 139 |
+
RE_REPEATED_WORD = re.compile(r'\b(\w+)(?:\s+\1){2,}\b', re.I)
|
| 140 |
+
RE_MULTI_NEWLINE = re.compile(r'\n{4,}')
|
| 141 |
+
RE_MULTI_SPACE = re.compile(r'[ \t]{2,}')
|
| 142 |
+
RE_TRAILING_SPACE = re.compile(r'[ \t]+$', re.M)
|
| 143 |
+
RE_NO_SPACE_AFTER_PERIOD = re.compile(r'([.!?])([A-Z])')
|
| 144 |
+
RE_DOUBLE_PERIOD = re.compile(r'\.{2}(?!\.)')
|
| 145 |
+
RE_SPACE_BEFORE_PUNCT = re.compile(r'\s+([.,;:!?])')
|
| 146 |
+
|
| 147 |
+
RE_CJK = re.compile(r'[\u4e00-\u9fff\u3040-\u309f\u30a0-\u30ff\uac00-\ud7af]{3,}')
|
| 148 |
+
RE_ARABIC = re.compile(r'[\u0600-\u06ff]{5,}')
|
| 149 |
+
RE_CYRILLIC = re.compile(r'[\u0400-\u04ff]{5,}')
|
| 150 |
+
RE_DEVANAGARI = re.compile(r'[\u0900-\u097f]{5,}')
|
| 151 |
+
RE_RESIDUAL_CODE = re.compile(r'(function\s*\(|var\s+\w+\s*=|console\.log|document\.get|if\s*\(\s*\w+\s*[!=]==)', re.I)
|
| 152 |
+
|
| 153 |
+
RE_COOKIE_LINE = re.compile(r'^.*(?:cookie|cookies)\s+(?:policy|consent|notice|preferences|settings).*$', re.I | re.M)
|
| 154 |
+
RE_SUBSCRIBE_LINE = re.compile(r'^.*(?:subscribe|sign\s*up\s+(?:for|to)\s+(?:our|the)\s+newsletter|unsubscribe|opt[\s-]*out\s+of).*$', re.I | re.M)
|
| 155 |
+
RE_CLICKBAIT_LINE = re.compile(r'^.*(?:you\s+won\'?t\s+believe|click\s+here|read\s+more\s*\.{0,3}$|share\s+this\s+(?:article|post|story)|trending\s+now|sponsored\s+content|advertisement).*$', re.I | re.M)
|
| 156 |
+
RE_SOCIAL_LINE = re.compile(r'^.*(?:follow\s+us\s+on|share\s+on\s+(?:facebook|twitter|linkedin|instagram)|like\s+us\s+on|tweet\s+this).*$', re.I | re.M)
|
| 157 |
+
RE_NAV_LINE = re.compile(r'^.*(?:skip\s+to\s+(?:main\s+)?content|back\s+to\s+top|previous\s+article|next\s+article|related\s+(?:articles|posts)).*$', re.I | re.M)
|
| 158 |
+
RE_LOGIN_LINE = re.compile(r'^.*(?:log\s*in\s+to\s+(?:your|an)\s+account|create\s+(?:a\s+)?(?:free\s+)?account|forgot\s+(?:your\s+)?password|already\s+(?:a\s+)?member).*$', re.I | re.M)
|
| 159 |
+
RE_COMMENT_LINE = re.compile(r'^.*(?:leave\s+a\s+(?:comment|reply)|post\s+a\s+comment|\d+\s+comments?$|logged\s+in\s+as).*$', re.I | re.M)
|
| 160 |
+
RE_COPYRIGHT_LINE = re.compile(r'^.*(?:all\s+rights\s+reserved|\(c\)\s*\d{4}|copyright\s+\d{4}).*$', re.I | re.M)
|
| 161 |
+
|
| 162 |
+
|
| 163 |
+
def deep_clean(text):
|
| 164 |
+
if not text or len(text.strip()) < 30:
|
| 165 |
+
return None
|
| 166 |
+
text = unicodedata.normalize("NFKC", text)
|
| 167 |
+
for ch in CONTROL_CHARS:
|
| 168 |
+
text = text.replace(ch, "")
|
| 169 |
+
for old, new in HTML_ENTITIES:
|
| 170 |
+
text = text.replace(old, new)
|
| 171 |
+
text = RE_HTML_COMMENT.sub("", text)
|
| 172 |
+
text = RE_HTML_TAG.sub("", text)
|
| 173 |
+
text = RE_URL.sub("", text)
|
| 174 |
+
text = RE_EMAIL.sub("", text)
|
| 175 |
+
text = RE_FILE_PATH.sub("", text)
|
| 176 |
+
text = RE_CODE_BLOCK.sub("", text)
|
| 177 |
+
text = RE_REPEATED_LINE.sub(r'\1', text)
|
| 178 |
+
text = RE_REPEATED_PUNCT.sub(r'\1\1\1', text)
|
| 179 |
+
text = RE_REPEATED_CHAR.sub(r'\1\1\1', text)
|
| 180 |
+
text = RE_REPEATED_WORD.sub(r'\1', text)
|
| 181 |
+
text = text.replace('\t', ' ')
|
| 182 |
+
text = RE_TRAILING_SPACE.sub('', text)
|
| 183 |
+
text = RE_MULTI_SPACE.sub(' ', text)
|
| 184 |
+
text = RE_MULTI_NEWLINE.sub('\n\n\n', text)
|
| 185 |
+
text = RE_DOUBLE_PERIOD.sub('.', text)
|
| 186 |
+
text = RE_NO_SPACE_AFTER_PERIOD.sub(r'\1 \2', text)
|
| 187 |
+
text = RE_SPACE_BEFORE_PUNCT.sub(r'\1', text)
|
| 188 |
+
text = text.replace('\u2018', "'").replace('\u2019', "'")
|
| 189 |
+
text = text.replace('\u201c', '"').replace('\u201d', '"')
|
| 190 |
+
text = text.replace('\u2013', '-').replace('\u2014', ' - ')
|
| 191 |
+
text = text.replace('\u2026', '...')
|
| 192 |
+
text = text.replace('\u2022', '- ')
|
| 193 |
+
text = text.replace('\u00b7', ' ')
|
| 194 |
+
text = text.replace('\u00a0', ' ')
|
| 195 |
+
lines = text.split('\n')
|
| 196 |
+
clean_lines = []
|
| 197 |
+
for line in lines:
|
| 198 |
+
line = line.strip()
|
| 199 |
+
if not line:
|
| 200 |
+
clean_lines.append('')
|
| 201 |
+
continue
|
| 202 |
+
if len(line) > 10:
|
| 203 |
+
alpha_count = sum(1 for c in line if c.isalpha())
|
| 204 |
+
if alpha_count / len(line) < 0.40:
|
| 205 |
+
continue
|
| 206 |
+
if line.count('|') > 3 or line.count('{') > 2 or line.count('}') > 2:
|
| 207 |
+
continue
|
| 208 |
+
if RE_IMPORT.match(line):
|
| 209 |
+
continue
|
| 210 |
+
if line and line[0].isalpha() and line[0].islower():
|
| 211 |
+
if not clean_lines or clean_lines[-1] == '' or clean_lines[-1].rstrip().endswith(('.', '!', '?', ':')):
|
| 212 |
+
line = line[0].upper() + line[1:]
|
| 213 |
+
clean_lines.append(line)
|
| 214 |
+
text = '\n'.join(clean_lines)
|
| 215 |
+
text = text.strip()
|
| 216 |
+
paragraphs = text.split('\n\n')
|
| 217 |
+
seen = set()
|
| 218 |
+
unique_paragraphs = []
|
| 219 |
+
for p in paragraphs:
|
| 220 |
+
p_stripped = p.strip()
|
| 221 |
+
if not p_stripped:
|
| 222 |
+
continue
|
| 223 |
+
p_key = ' '.join(p_stripped.lower().split())
|
| 224 |
+
if p_key not in seen:
|
| 225 |
+
seen.add(p_key)
|
| 226 |
+
unique_paragraphs.append(p_stripped)
|
| 227 |
+
text = '\n\n'.join(unique_paragraphs)
|
| 228 |
+
text = text.strip()
|
| 229 |
+
if len(text) < 50:
|
| 230 |
+
return None
|
| 231 |
+
if len(text.split()) < 10:
|
| 232 |
+
return None
|
| 233 |
+
ascii_count = sum(1 for c in text if ord(c) < 128)
|
| 234 |
+
if ascii_count / max(len(text), 1) < 0.85:
|
| 235 |
+
return None
|
| 236 |
+
return text
|
| 237 |
+
|
| 238 |
+
|
| 239 |
+
def smart_filter(text):
|
| 240 |
+
words = text.split()
|
| 241 |
+
word_count = len(words)
|
| 242 |
+
if word_count < 50:
|
| 243 |
+
return False, text, f"too short ({word_count} words)"
|
| 244 |
+
scripts = []
|
| 245 |
+
if RE_CJK.search(text): scripts.append("CJK")
|
| 246 |
+
if RE_ARABIC.search(text): scripts.append("Arabic")
|
| 247 |
+
if RE_CYRILLIC.search(text): scripts.append("Cyrillic")
|
| 248 |
+
if RE_DEVANAGARI.search(text): scripts.append("Devanagari")
|
| 249 |
+
if scripts:
|
| 250 |
+
return False, text, f"non-English: {', '.join(scripts)}"
|
| 251 |
+
if word_count > 50:
|
| 252 |
+
unique_ratio = len(set(w.lower() for w in words)) / word_count
|
| 253 |
+
if unique_ratio < 0.20:
|
| 254 |
+
return False, text, f"repetitive ({unique_ratio:.3f})"
|
| 255 |
+
code_matches = RE_RESIDUAL_CODE.findall(text)
|
| 256 |
+
if len(code_matches) >= 5:
|
| 257 |
+
return False, text, f"residual code ({len(code_matches)})"
|
| 258 |
+
original_len = len(text)
|
| 259 |
+
for pattern in [RE_COOKIE_LINE, RE_SUBSCRIBE_LINE, RE_CLICKBAIT_LINE,
|
| 260 |
+
RE_SOCIAL_LINE, RE_NAV_LINE, RE_LOGIN_LINE,
|
| 261 |
+
RE_COMMENT_LINE, RE_COPYRIGHT_LINE]:
|
| 262 |
+
text = pattern.sub('', text)
|
| 263 |
+
lines = text.split('\n')
|
| 264 |
+
clean_lines = []
|
| 265 |
+
for line in lines:
|
| 266 |
+
stripped = line.strip()
|
| 267 |
+
if stripped and len(stripped) > 10:
|
| 268 |
+
digit_count = sum(1 for c in stripped if c.isdigit() or c in ' ,.\t-+/%$')
|
| 269 |
+
if digit_count / len(stripped) > 0.80:
|
| 270 |
+
continue
|
| 271 |
+
clean_lines.append(line)
|
| 272 |
+
text = '\n'.join(clean_lines)
|
| 273 |
+
text = re.sub(r'\n{3,}', '\n\n', text)
|
| 274 |
+
text = text.strip()
|
| 275 |
+
if len(text.split()) < 50:
|
| 276 |
+
return False, text, "too short after stripping"
|
| 277 |
+
return True, text, None
|
| 278 |
+
|
| 279 |
+
|
| 280 |
+
# ==============================================================================
|
| 281 |
+
# LITDATA I/O
|
| 282 |
+
# ==============================================================================
|
| 283 |
+
|
| 284 |
+
def write_litdata_chunks(output_dir, token_stream, config, start_chunk_idx=0):
|
| 285 |
+
"""Write token stream as litdata binary chunks, starting at given chunk index."""
|
| 286 |
+
os.makedirs(output_dir, exist_ok=True)
|
| 287 |
+
dtype_size = DTYPE().itemsize
|
| 288 |
+
tokens_per_chunk = CHUNK_BYTES_TARGET // dtype_size
|
| 289 |
+
tokens_per_chunk = (tokens_per_chunk // BLOCK_SIZE) * BLOCK_SIZE
|
| 290 |
+
chunks_metadata = []
|
| 291 |
+
pos = 0
|
| 292 |
+
chunk_idx = start_chunk_idx
|
| 293 |
+
while pos < len(token_stream):
|
| 294 |
+
remaining = len(token_stream) - pos
|
| 295 |
+
chunk_tokens = min(tokens_per_chunk, remaining)
|
| 296 |
+
num_blocks = chunk_tokens // BLOCK_SIZE
|
| 297 |
+
if num_blocks == 0:
|
| 298 |
+
break
|
| 299 |
+
actual_tokens = num_blocks * BLOCK_SIZE
|
| 300 |
+
chunk_data = token_stream[pos:pos + actual_tokens]
|
| 301 |
+
filename = f"chunk-0-{chunk_idx}.bin"
|
| 302 |
+
filepath = os.path.join(output_dir, filename)
|
| 303 |
+
header_num_items = np.array([num_blocks], dtype=np.uint32)
|
| 304 |
+
offsets = np.arange(num_blocks + 1, dtype=np.uint32) * (BLOCK_SIZE * dtype_size)
|
| 305 |
+
header = np.concatenate([header_num_items, offsets])
|
| 306 |
+
with open(filepath, "wb") as f:
|
| 307 |
+
header.tofile(f)
|
| 308 |
+
chunk_data.tofile(f)
|
| 309 |
+
meta = {
|
| 310 |
+
"chunk_bytes": int(header.nbytes + chunk_data.nbytes),
|
| 311 |
+
"chunk_size": num_blocks,
|
| 312 |
+
"dim": int(actual_tokens),
|
| 313 |
+
"filename": filename,
|
| 314 |
+
}
|
| 315 |
+
chunks_metadata.append(meta)
|
| 316 |
+
pos += actual_tokens
|
| 317 |
+
chunk_idx += 1
|
| 318 |
+
print(f" Written chunk {chunk_idx} ({pos:,}/{len(token_stream):,} tokens)")
|
| 319 |
+
return chunks_metadata
|
| 320 |
+
|
| 321 |
+
|
| 322 |
+
def merge_litdata(sources, dest_dir):
|
| 323 |
+
"""
|
| 324 |
+
Merge multiple litdata directories into one.
|
| 325 |
+
sources: list of Path objects to litdata directories
|
| 326 |
+
dest_dir: Path for the merged output
|
| 327 |
+
"""
|
| 328 |
+
os.makedirs(str(dest_dir), exist_ok=True)
|
| 329 |
+
all_chunks = []
|
| 330 |
+
config = None
|
| 331 |
+
chunk_offset = 0
|
| 332 |
+
|
| 333 |
+
for src in sources:
|
| 334 |
+
idx_path = src / "index.json"
|
| 335 |
+
with open(idx_path, "r") as f:
|
| 336 |
+
index = json.load(f)
|
| 337 |
+
|
| 338 |
+
if config is None:
|
| 339 |
+
config = index["config"]
|
| 340 |
+
|
| 341 |
+
for chunk_meta in index["chunks"]:
|
| 342 |
+
old_filename = chunk_meta["filename"]
|
| 343 |
+
new_filename = f"chunk-0-{chunk_offset}.bin"
|
| 344 |
+
|
| 345 |
+
# Copy the binary chunk file
|
| 346 |
+
src_file = src / old_filename
|
| 347 |
+
dst_file = dest_dir / new_filename
|
| 348 |
+
shutil.copy2(str(src_file), str(dst_file))
|
| 349 |
+
|
| 350 |
+
new_meta = dict(chunk_meta)
|
| 351 |
+
new_meta["filename"] = new_filename
|
| 352 |
+
all_chunks.append(new_meta)
|
| 353 |
+
chunk_offset += 1
|
| 354 |
+
|
| 355 |
+
src_tokens = sum(c["dim"] for c in index["chunks"])
|
| 356 |
+
print(f" Merged {src.name}: {len(index['chunks'])} chunks, {src_tokens:,} tokens")
|
| 357 |
+
|
| 358 |
+
# Write combined index
|
| 359 |
+
combined_index = {
|
| 360 |
+
"chunks": all_chunks,
|
| 361 |
+
"config": config,
|
| 362 |
+
"updated_at": str(time.time()),
|
| 363 |
+
}
|
| 364 |
+
with open(str(dest_dir / "index.json"), "w") as f:
|
| 365 |
+
json.dump(combined_index, f, indent=2)
|
| 366 |
+
|
| 367 |
+
total = sum(c["dim"] for c in all_chunks)
|
| 368 |
+
print(f" Total: {len(all_chunks)} chunks, {total:,} tokens")
|
| 369 |
+
return all_chunks, config
|
| 370 |
+
|
| 371 |
+
|
| 372 |
+
# ==============================================================================
|
| 373 |
+
# MAIN PIPELINE
|
| 374 |
+
# ==============================================================================
|
| 375 |
+
|
| 376 |
+
def main():
|
| 377 |
+
from datasets import load_dataset
|
| 378 |
+
|
| 379 |
+
t_start = time.time()
|
| 380 |
+
|
| 381 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 382 |
+
# STEP 1: Merge existing litdata_english_100m + litdata_english_clean
|
| 383 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 384 |
+
print(f"\n{'='*75}")
|
| 385 |
+
print(f" STEP 1: MERGING EXISTING DATASETS")
|
| 386 |
+
print(f" - {LITDATA_100M}")
|
| 387 |
+
print(f" - {LITDATA_CLEAN}")
|
| 388 |
+
print(f" β {COMBINED_DIR}")
|
| 389 |
+
print(f"{'='*75}\n")
|
| 390 |
+
|
| 391 |
+
merged_chunks, config = merge_litdata([LITDATA_100M, LITDATA_CLEAN], COMBINED_DIR)
|
| 392 |
+
merged_tokens = sum(c["dim"] for c in merged_chunks)
|
| 393 |
+
merged_chunk_count = len(merged_chunks)
|
| 394 |
+
print(f"\n Step 1 done: {merged_tokens:,} tokens in {merged_chunk_count} chunks")
|
| 395 |
+
|
| 396 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 397 |
+
# STEP 2: Download fresh 100M tokens
|
| 398 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 399 |
+
wiki_target = int(TARGET_TOKENS * WIKI_SHARE)
|
| 400 |
+
fineweb_target = TARGET_TOKENS - wiki_target
|
| 401 |
+
|
| 402 |
+
print(f"\n{'='*75}")
|
| 403 |
+
print(f" STEP 2: DOWNLOADING FRESH DATA")
|
| 404 |
+
print(f" Target: {TARGET_TOKENS:,} tokens ({wiki_target:,} Wiki + {fineweb_target:,} FineWeb)")
|
| 405 |
+
print(f" Skipping first {WIKI_SKIP:,} Wiki + {FINEWEB_SKIP:,} FineWeb (already used)")
|
| 406 |
+
print(f"{'='*75}")
|
| 407 |
+
|
| 408 |
+
os.makedirs(str(PARQUET_DIR), exist_ok=True)
|
| 409 |
+
|
| 410 |
+
# --- Wikipedia ---
|
| 411 |
+
print(f"\n [Wikipedia] Streaming English articles (skipping first {WIKI_SKIP:,})...")
|
| 412 |
+
ds_wiki = load_dataset(
|
| 413 |
+
"wikimedia/wikipedia", "20231101.en",
|
| 414 |
+
split="train", streaming=True,
|
| 415 |
+
trust_remote_code=False
|
| 416 |
+
)
|
| 417 |
+
|
| 418 |
+
wiki_texts = []
|
| 419 |
+
wiki_tokens = 0
|
| 420 |
+
wiki_seen = 0
|
| 421 |
+
wiki_skipped_quality = 0
|
| 422 |
+
wiki_skipped_meta = 0
|
| 423 |
+
t0 = time.time()
|
| 424 |
+
|
| 425 |
+
for article in ds_wiki:
|
| 426 |
+
title = (article.get("title") or "").strip()
|
| 427 |
+
raw = article.get("text") or ""
|
| 428 |
+
|
| 429 |
+
if _WIKI_SKIP_PATTERNS.search(title):
|
| 430 |
+
wiki_skipped_meta += 1
|
| 431 |
+
continue
|
| 432 |
+
|
| 433 |
+
cleaned = clean_text(raw)
|
| 434 |
+
if not is_high_quality(cleaned, min_chars=800, min_words=120):
|
| 435 |
+
wiki_skipped_quality += 1
|
| 436 |
+
continue
|
| 437 |
+
|
| 438 |
+
wiki_seen += 1
|
| 439 |
+
|
| 440 |
+
if wiki_seen <= WIKI_SKIP:
|
| 441 |
+
if wiki_seen % 10000 == 0:
|
| 442 |
+
print(f" Skipping... {wiki_seen:,}/{WIKI_SKIP:,}")
|
| 443 |
+
continue
|
| 444 |
+
|
| 445 |
+
full_text = f"{title}\n\n{cleaned}"
|
| 446 |
+
est_tok = int(len(full_text.split()) * TOKENS_PER_WORD)
|
| 447 |
+
wiki_texts.append(full_text)
|
| 448 |
+
wiki_tokens += est_tok
|
| 449 |
+
|
| 450 |
+
if len(wiki_texts) % 5000 == 0:
|
| 451 |
+
elapsed = time.time() - t0
|
| 452 |
+
print(f" Collected {len(wiki_texts):,} articles | ~{wiki_tokens:,} tokens | {elapsed:.0f}s")
|
| 453 |
+
|
| 454 |
+
if wiki_tokens >= wiki_target:
|
| 455 |
+
break
|
| 456 |
+
|
| 457 |
+
elapsed = time.time() - t0
|
| 458 |
+
print(f" [Wikipedia] Done: {len(wiki_texts):,} articles, ~{wiki_tokens:,} tokens in {elapsed:.0f}s")
|
| 459 |
+
print(f" (skipped {WIKI_SKIP:,} already-used + {wiki_skipped_quality:,} low-quality + {wiki_skipped_meta:,} meta)")
|
| 460 |
+
|
| 461 |
+
BATCH = 5000
|
| 462 |
+
for i in range(0, len(wiki_texts), BATCH):
|
| 463 |
+
batch = wiki_texts[i:i+BATCH]
|
| 464 |
+
fp = PARQUET_DIR / f"wiki_{i//BATCH:04d}.parquet"
|
| 465 |
+
pq.write_table(pa.table({"text": batch}), str(fp))
|
| 466 |
+
print(f" Saved {(len(wiki_texts)-1)//BATCH + 1} wiki parquet files")
|
| 467 |
+
|
| 468 |
+
# --- FineWeb-Edu ---
|
| 469 |
+
print(f"\n [FineWeb-Edu] Streaming (score >= {FINEWEB_MIN_SCORE}, skipping first {FINEWEB_SKIP:,})...")
|
| 470 |
+
ds_fineweb = load_dataset(
|
| 471 |
+
"HuggingFaceFW/fineweb-edu", "sample-10BT",
|
| 472 |
+
split="train", streaming=True,
|
| 473 |
+
trust_remote_code=False
|
| 474 |
+
)
|
| 475 |
+
|
| 476 |
+
fineweb_texts = []
|
| 477 |
+
fineweb_tokens = 0
|
| 478 |
+
fineweb_seen = 0
|
| 479 |
+
fineweb_skipped_quality = 0
|
| 480 |
+
fineweb_skipped_score = 0
|
| 481 |
+
t0 = time.time()
|
| 482 |
+
|
| 483 |
+
for doc in ds_fineweb:
|
| 484 |
+
score = doc.get("score", 0)
|
| 485 |
+
if not isinstance(score, (int, float)):
|
| 486 |
+
try:
|
| 487 |
+
score = float(score)
|
| 488 |
+
except (ValueError, TypeError):
|
| 489 |
+
continue
|
| 490 |
+
if score < FINEWEB_MIN_SCORE:
|
| 491 |
+
fineweb_skipped_score += 1
|
| 492 |
+
continue
|
| 493 |
+
|
| 494 |
+
raw = doc.get("text") or ""
|
| 495 |
+
cleaned = clean_text(raw)
|
| 496 |
+
if not is_high_quality(cleaned, min_chars=500, min_words=80):
|
| 497 |
+
fineweb_skipped_quality += 1
|
| 498 |
+
continue
|
| 499 |
+
|
| 500 |
+
fineweb_seen += 1
|
| 501 |
+
|
| 502 |
+
if fineweb_seen <= FINEWEB_SKIP:
|
| 503 |
+
if fineweb_seen % 10000 == 0:
|
| 504 |
+
print(f" Skipping... {fineweb_seen:,}/{FINEWEB_SKIP:,}")
|
| 505 |
+
continue
|
| 506 |
+
|
| 507 |
+
est_tok = int(len(cleaned.split()) * TOKENS_PER_WORD)
|
| 508 |
+
fineweb_texts.append(cleaned)
|
| 509 |
+
fineweb_tokens += est_tok
|
| 510 |
+
|
| 511 |
+
if len(fineweb_texts) % 5000 == 0:
|
| 512 |
+
elapsed = time.time() - t0
|
| 513 |
+
print(f" Collected {len(fineweb_texts):,} docs | ~{fineweb_tokens:,} tokens | {elapsed:.0f}s")
|
| 514 |
+
|
| 515 |
+
if fineweb_tokens >= fineweb_target:
|
| 516 |
+
break
|
| 517 |
+
|
| 518 |
+
elapsed = time.time() - t0
|
| 519 |
+
print(f" [FineWeb-Edu] Done: {len(fineweb_texts):,} docs, ~{fineweb_tokens:,} tokens in {elapsed:.0f}s")
|
| 520 |
+
print(f" (skipped {FINEWEB_SKIP:,} already-used + {fineweb_skipped_quality:,} low-quality + {fineweb_skipped_score:,} low-score)")
|
| 521 |
+
|
| 522 |
+
for i in range(0, len(fineweb_texts), BATCH):
|
| 523 |
+
batch = fineweb_texts[i:i+BATCH]
|
| 524 |
+
fp = PARQUET_DIR / f"fineweb_{i//BATCH:04d}.parquet"
|
| 525 |
+
pq.write_table(pa.table({"text": batch}), str(fp))
|
| 526 |
+
print(f" Saved {(len(fineweb_texts)-1)//BATCH + 1} fineweb parquet files")
|
| 527 |
+
|
| 528 |
+
all_texts = wiki_texts + fineweb_texts
|
| 529 |
+
total_est_tokens = wiki_tokens + fineweb_tokens
|
| 530 |
+
del wiki_texts, fineweb_texts
|
| 531 |
+
|
| 532 |
+
print(f"\n Step 2 done: {len(all_texts):,} documents, ~{total_est_tokens:,} estimated tokens")
|
| 533 |
+
|
| 534 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 535 |
+
# STEP 3: Deep clean every document
|
| 536 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 537 |
+
print(f"\n{'='*75}")
|
| 538 |
+
print(f" STEP 3: DEEP CLEANING {len(all_texts):,} DOCUMENTS")
|
| 539 |
+
print(f"{'='*75}")
|
| 540 |
+
|
| 541 |
+
t2 = time.time()
|
| 542 |
+
cleaned_texts = []
|
| 543 |
+
dropped_clean = 0
|
| 544 |
+
for i, text in enumerate(all_texts):
|
| 545 |
+
result = deep_clean(text)
|
| 546 |
+
if result is not None:
|
| 547 |
+
cleaned_texts.append(result)
|
| 548 |
+
else:
|
| 549 |
+
dropped_clean += 1
|
| 550 |
+
if (i + 1) % 10000 == 0 or i == len(all_texts) - 1:
|
| 551 |
+
print(f" Cleaned {i+1:,}/{len(all_texts):,} | kept={len(cleaned_texts):,} | dropped={dropped_clean:,}")
|
| 552 |
+
|
| 553 |
+
del all_texts
|
| 554 |
+
print(f" Deep clean done in {time.time()-t2:.1f}s")
|
| 555 |
+
print(f" Kept: {len(cleaned_texts):,} | Dropped: {dropped_clean:,}")
|
| 556 |
+
|
| 557 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 558 |
+
# STEP 4: Smart cleanup
|
| 559 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 560 |
+
print(f"\n{'='*75}")
|
| 561 |
+
print(f" STEP 4: SMART CLEANUP ON {len(cleaned_texts):,} DOCUMENTS")
|
| 562 |
+
print(f"{'='*75}")
|
| 563 |
+
|
| 564 |
+
t3 = time.time()
|
| 565 |
+
final_texts = []
|
| 566 |
+
removed_reasons = Counter()
|
| 567 |
+
total_boilerplate = 0
|
| 568 |
+
|
| 569 |
+
for i, text in enumerate(cleaned_texts):
|
| 570 |
+
keep, stripped_text, reason = smart_filter(text)
|
| 571 |
+
if keep:
|
| 572 |
+
final_texts.append(stripped_text)
|
| 573 |
+
total_boilerplate += len(text) - len(stripped_text)
|
| 574 |
+
else:
|
| 575 |
+
removed_reasons[reason.split('(')[0].strip().split(':')[0].strip()] += 1
|
| 576 |
+
if (i + 1) % 10000 == 0 or i == len(cleaned_texts) - 1:
|
| 577 |
+
print(f" Processed {i+1:,}/{len(cleaned_texts):,} | kept={len(final_texts):,}")
|
| 578 |
+
|
| 579 |
+
del cleaned_texts
|
| 580 |
+
print(f" Smart cleanup done in {time.time()-t3:.1f}s")
|
| 581 |
+
print(f" Final new documents: {len(final_texts):,}")
|
| 582 |
+
print(f" Boilerplate stripped: {total_boilerplate:,} chars")
|
| 583 |
+
if removed_reasons:
|
| 584 |
+
print(f" Removal reasons:")
|
| 585 |
+
for reason, count in sorted(removed_reasons.items(), key=lambda x: -x[1]):
|
| 586 |
+
print(f" {reason:<35} {count:>6,}")
|
| 587 |
+
|
| 588 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 589 |
+
# STEP 5: Tokenize new data
|
| 590 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 591 |
+
print(f"\n{'='*75}")
|
| 592 |
+
print(f" STEP 5: TOKENIZING {len(final_texts):,} NEW DOCUMENTS")
|
| 593 |
+
print(f"{'='*75}")
|
| 594 |
+
|
| 595 |
+
t4 = time.time()
|
| 596 |
+
all_token_ids = []
|
| 597 |
+
total_tokens = 0
|
| 598 |
+
ENCODE_BATCH = 10000
|
| 599 |
+
|
| 600 |
+
for i in range(0, len(final_texts), ENCODE_BATCH):
|
| 601 |
+
batch = final_texts[i:i+ENCODE_BATCH]
|
| 602 |
+
encoded = tokenizer.encode_batch(batch, add_special_tokens=False)
|
| 603 |
+
for enc in encoded:
|
| 604 |
+
ids = enc.ids
|
| 605 |
+
all_token_ids.extend(ids)
|
| 606 |
+
all_token_ids.append(EOS_TOKEN_ID)
|
| 607 |
+
total_tokens += len(ids) + 1
|
| 608 |
+
done = min(i + ENCODE_BATCH, len(final_texts))
|
| 609 |
+
if done % 20000 == 0 or done == len(final_texts):
|
| 610 |
+
print(f" Tokenized {done:,}/{len(final_texts):,} ({total_tokens:,} tokens)")
|
| 611 |
+
|
| 612 |
+
del final_texts
|
| 613 |
+
print(f" Tokenized in {time.time()-t4:.1f}s - {total_tokens:,} total new tokens")
|
| 614 |
+
|
| 615 |
+
# βββββββββββββββββββββββββββββββββοΏ½οΏ½βββββββββββββββββββββββββββ
|
| 616 |
+
# STEP 6: Write new chunks directly into combined dir
|
| 617 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 618 |
+
print(f"\n{'='*75}")
|
| 619 |
+
print(f" STEP 6: WRITING NEW CHUNKS INTO COMBINED DATASET")
|
| 620 |
+
print(f"{'='*75}")
|
| 621 |
+
|
| 622 |
+
token_array = np.array(all_token_ids, dtype=DTYPE)
|
| 623 |
+
del all_token_ids
|
| 624 |
+
|
| 625 |
+
litdata_config = {
|
| 626 |
+
"block_size": BLOCK_SIZE,
|
| 627 |
+
"vocab_size": tokenizer.get_vocab_size(),
|
| 628 |
+
}
|
| 629 |
+
|
| 630 |
+
# Write new chunks starting after the existing merged chunks
|
| 631 |
+
new_chunks = write_litdata_chunks(
|
| 632 |
+
str(COMBINED_DIR), token_array, litdata_config,
|
| 633 |
+
start_chunk_idx=merged_chunk_count,
|
| 634 |
+
)
|
| 635 |
+
new_token_count = sum(c["dim"] for c in new_chunks)
|
| 636 |
+
del token_array
|
| 637 |
+
|
| 638 |
+
# Update the combined index.json with the new chunks appended
|
| 639 |
+
all_final_chunks = merged_chunks + new_chunks
|
| 640 |
+
final_total_tokens = sum(c["dim"] for c in all_final_chunks)
|
| 641 |
+
final_index = {
|
| 642 |
+
"chunks": all_final_chunks,
|
| 643 |
+
"config": litdata_config,
|
| 644 |
+
"updated_at": str(time.time()),
|
| 645 |
+
}
|
| 646 |
+
with open(str(COMBINED_DIR / "index.json"), "w") as f:
|
| 647 |
+
json.dump(final_index, f, indent=2)
|
| 648 |
+
|
| 649 |
+
print(f"\n New data: {new_token_count:,} tokens in {len(new_chunks)} chunks")
|
| 650 |
+
print(f" Final combined: {final_total_tokens:,} tokens in {len(all_final_chunks)} chunks")
|
| 651 |
+
|
| 652 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 653 |
+
# FINAL REPORT
|
| 654 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 655 |
+
total_time = time.time() - t_start
|
| 656 |
+
|
| 657 |
+
report_lines = []
|
| 658 |
+
report_lines.append(f"\n{'='*75}")
|
| 659 |
+
report_lines.append(f" LITDATA_COMBINED - MERGE & BUILD REPORT")
|
| 660 |
+
report_lines.append(f"{'='*75}")
|
| 661 |
+
report_lines.append(f"\n Total time: {total_time:.0f}s ({total_time/60:.1f} min)")
|
| 662 |
+
report_lines.append(f"\n SOURCE DATASETS MERGED")
|
| 663 |
+
report_lines.append(f" {'-'*60}")
|
| 664 |
+
report_lines.append(f" litdata_english_100m + litdata_english_clean")
|
| 665 |
+
report_lines.append(f" Merged subtotal: {merged_tokens:,} tokens ({merged_chunk_count} chunks)")
|
| 666 |
+
report_lines.append(f"\n NEW 100M BATCH (articles {WIKI_SKIP+1:,}+)")
|
| 667 |
+
report_lines.append(f" {'-'*60}")
|
| 668 |
+
report_lines.append(f" Wikipedia: ~{wiki_tokens:,} est. tokens")
|
| 669 |
+
report_lines.append(f" FineWeb-Edu: ~{fineweb_tokens:,} est. tokens")
|
| 670 |
+
report_lines.append(f" Downloaded: ~{total_est_tokens:,} est. tokens")
|
| 671 |
+
report_lines.append(f" Deep clean: dropped {dropped_clean:,} docs")
|
| 672 |
+
report_lines.append(f" Smart filter: removed {sum(removed_reasons.values()):,} docs")
|
| 673 |
+
if removed_reasons:
|
| 674 |
+
for reason, count in sorted(removed_reasons.items(), key=lambda x: -x[1]):
|
| 675 |
+
report_lines.append(f" - {reason}: {count:,}")
|
| 676 |
+
report_lines.append(f" Boilerplate: stripped {total_boilerplate:,} chars")
|
| 677 |
+
report_lines.append(f" New tokens: {new_token_count:,} ({len(new_chunks)} chunks)")
|
| 678 |
+
report_lines.append(f"\n FINAL COMBINED OUTPUT")
|
| 679 |
+
report_lines.append(f" {'-'*60}")
|
| 680 |
+
report_lines.append(f" Location: {COMBINED_DIR}")
|
| 681 |
+
report_lines.append(f" Chunks: {len(all_final_chunks)}")
|
| 682 |
+
report_lines.append(f" Tokens: {final_total_tokens:,}")
|
| 683 |
+
report_lines.append(f" Format: litdata binary (int32, BLOCK_SIZE=1025, EOS=0)")
|
| 684 |
+
report_lines.append(f"\n ZERO OVERLAP GUARANTEE:")
|
| 685 |
+
report_lines.append(f" - litdata_english_clean: Wiki articles 1-19,523 + FineWeb docs 1-19,445")
|
| 686 |
+
report_lines.append(f" - litdata_english_100m: Wiki articles 25,001-78,482 + FineWeb docs 25,001-68,550")
|
| 687 |
+
report_lines.append(f" - New batch: Wiki articles {WIKI_SKIP+1:,}+ + FineWeb docs {FINEWEB_SKIP+1:,}+")
|
| 688 |
+
report_lines.append(f"\n{'='*75}")
|
| 689 |
+
|
| 690 |
+
full_report = '\n'.join(report_lines)
|
| 691 |
+
print(full_report)
|
| 692 |
+
|
| 693 |
+
report_path = COMBINED_DIR / "BUILD_REPORT.txt"
|
| 694 |
+
with open(report_path, "w", encoding="utf-8") as f:
|
| 695 |
+
f.write(full_report)
|
| 696 |
+
print(f"\n Report saved to: {report_path}")
|
| 697 |
+
print(f" Done! Final combined dataset ready for training.")
|
| 698 |
+
|
| 699 |
+
|
| 700 |
+
if __name__ == "__main__":
|
| 701 |
+
main()
|
Base/scripts/prepare_finetune_data.py
ADDED
|
@@ -0,0 +1,88 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Merge all JSONL instruction-tuning datasets into train.json and val.json
|
| 3 |
+
for litgpt finetune (JSON data module).
|
| 4 |
+
|
| 5 |
+
Each row must have: instruction, output, and optionally input.
|
| 6 |
+
"""
|
| 7 |
+
|
| 8 |
+
import json
|
| 9 |
+
import os
|
| 10 |
+
import random
|
| 11 |
+
import argparse
|
| 12 |
+
|
| 13 |
+
def main():
|
| 14 |
+
parser = argparse.ArgumentParser()
|
| 15 |
+
parser.add_argument("--input_dir", type=str, default=r"D:\ASTERIZER 2026\LitGPT\OLD_DATASETS")
|
| 16 |
+
parser.add_argument("--output_dir", type=str, default=r"D:\ASTERIZER 2026\LUNA\Base\Datasets\finetune")
|
| 17 |
+
parser.add_argument("--val_fraction", type=float, default=0.05, help="Fraction for validation split")
|
| 18 |
+
parser.add_argument("--seed", type=int, default=42)
|
| 19 |
+
args = parser.parse_args()
|
| 20 |
+
|
| 21 |
+
random.seed(args.seed)
|
| 22 |
+
os.makedirs(args.output_dir, exist_ok=True)
|
| 23 |
+
|
| 24 |
+
all_samples = []
|
| 25 |
+
file_counts = {}
|
| 26 |
+
|
| 27 |
+
for fname in sorted(os.listdir(args.input_dir)):
|
| 28 |
+
if not fname.endswith(".jsonl"):
|
| 29 |
+
continue
|
| 30 |
+
fpath = os.path.join(args.input_dir, fname)
|
| 31 |
+
count = 0
|
| 32 |
+
with open(fpath, "r", encoding="utf-8") as f:
|
| 33 |
+
for line in f:
|
| 34 |
+
line = line.strip()
|
| 35 |
+
if not line:
|
| 36 |
+
continue
|
| 37 |
+
row = json.loads(line)
|
| 38 |
+
# Keep only the keys litgpt expects
|
| 39 |
+
sample = {
|
| 40 |
+
"instruction": row.get("instruction", ""),
|
| 41 |
+
"input": row.get("input", ""),
|
| 42 |
+
"output": row.get("output", ""),
|
| 43 |
+
}
|
| 44 |
+
# Skip rows with empty instruction AND empty input
|
| 45 |
+
if not sample["instruction"].strip() and not sample["input"].strip():
|
| 46 |
+
continue
|
| 47 |
+
# Skip rows with empty output
|
| 48 |
+
if not sample["output"].strip():
|
| 49 |
+
continue
|
| 50 |
+
all_samples.append(sample)
|
| 51 |
+
count += 1
|
| 52 |
+
file_counts[fname] = count
|
| 53 |
+
print(f" {fname}: {count} samples")
|
| 54 |
+
|
| 55 |
+
print(f"\nTotal valid samples: {len(all_samples)}")
|
| 56 |
+
|
| 57 |
+
# Shuffle
|
| 58 |
+
random.shuffle(all_samples)
|
| 59 |
+
|
| 60 |
+
# Split
|
| 61 |
+
val_size = max(1, int(len(all_samples) * args.val_fraction))
|
| 62 |
+
val_data = all_samples[:val_size]
|
| 63 |
+
train_data = all_samples[val_size:]
|
| 64 |
+
|
| 65 |
+
print(f"Train: {len(train_data)}, Val: {val_size}")
|
| 66 |
+
|
| 67 |
+
# Write
|
| 68 |
+
train_path = os.path.join(args.output_dir, "train.json")
|
| 69 |
+
val_path = os.path.join(args.output_dir, "val.json")
|
| 70 |
+
|
| 71 |
+
with open(train_path, "w", encoding="utf-8") as f:
|
| 72 |
+
json.dump(train_data, f, ensure_ascii=False, indent=None)
|
| 73 |
+
with open(val_path, "w", encoding="utf-8") as f:
|
| 74 |
+
json.dump(val_data, f, ensure_ascii=False, indent=None)
|
| 75 |
+
|
| 76 |
+
print(f"\nSaved: {train_path}")
|
| 77 |
+
print(f"Saved: {val_path}")
|
| 78 |
+
|
| 79 |
+
# Show a few samples
|
| 80 |
+
print("\n--- Sample train entries ---")
|
| 81 |
+
for s in train_data[:3]:
|
| 82 |
+
print(f" instruction: {s['instruction'][:80]}")
|
| 83 |
+
print(f" input: {s['input'][:80]}")
|
| 84 |
+
print(f" output: {s['output'][:80]}")
|
| 85 |
+
print()
|
| 86 |
+
|
| 87 |
+
if __name__ == "__main__":
|
| 88 |
+
main()
|
Base/scripts/prepare_litdata.py
ADDED
|
@@ -0,0 +1,275 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Tokenize filtered parquet data and create LitData streaming format for litgpt.
|
| 3 |
+
|
| 4 |
+
Creates flat binary token chunks compatible with litdata's TokensLoader,
|
| 5 |
+
which litgpt's LitData data module uses for pretraining.
|
| 6 |
+
|
| 7 |
+
STREAMING MODE: Tokenizes texts in batches and flushes chunks to disk
|
| 8 |
+
incrementally, keeping RAM usage constant (~128MB buffer) regardless of
|
| 9 |
+
dataset size. Supports resuming from the last completed chunk.
|
| 10 |
+
|
| 11 |
+
Usage:
|
| 12 |
+
python Base/scripts/prepare_litdata.py --mode small
|
| 13 |
+
python Base/scripts/prepare_litdata.py --mode full
|
| 14 |
+
python Base/scripts/prepare_litdata.py --mode both
|
| 15 |
+
"""
|
| 16 |
+
|
| 17 |
+
import argparse
|
| 18 |
+
import json
|
| 19 |
+
import os
|
| 20 |
+
import time
|
| 21 |
+
from pathlib import Path
|
| 22 |
+
|
| 23 |
+
import numpy as np
|
| 24 |
+
import pyarrow.parquet as pq
|
| 25 |
+
|
| 26 |
+
from litgpt.tokenizer import Tokenizer
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
# Must match model config's block_size + 1 (extra token for targets)
|
| 30 |
+
BLOCK_SIZE = 1025
|
| 31 |
+
# 64 MB per chunk file
|
| 32 |
+
CHUNK_BYTES_TARGET = 64 * 1024 * 1024
|
| 33 |
+
# litdata's TokensLoader uses _TORCH_DTYPES_MAPPING[16] = torch.int32
|
| 34 |
+
DTYPE = np.int32
|
| 35 |
+
# litdata dtype index β 16 maps to torch.int32
|
| 36 |
+
DTYPE_INDEX = 16
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
def _write_chunk(output_dir, chunk_idx, block_buffer, num_blocks):
|
| 40 |
+
"""Write a single chunk file with litdata header + flat int32 data."""
|
| 41 |
+
dtype_size = DTYPE().itemsize
|
| 42 |
+
filename = f"chunk-0-{chunk_idx}.bin"
|
| 43 |
+
filepath = os.path.join(output_dir, filename)
|
| 44 |
+
|
| 45 |
+
chunk_data = block_buffer[:num_blocks * BLOCK_SIZE]
|
| 46 |
+
|
| 47 |
+
# Header: [num_items(uint32)] + [offsets 0..num_blocks(uint32)]
|
| 48 |
+
header_num_items = np.array([num_blocks], dtype=np.uint32)
|
| 49 |
+
offsets = np.arange(num_blocks + 1, dtype=np.uint32) * (BLOCK_SIZE * dtype_size)
|
| 50 |
+
header = np.concatenate([header_num_items, offsets])
|
| 51 |
+
|
| 52 |
+
with open(filepath, "wb") as f:
|
| 53 |
+
header.tofile(f)
|
| 54 |
+
chunk_data.tofile(f)
|
| 55 |
+
|
| 56 |
+
data_bytes = int(chunk_data.nbytes)
|
| 57 |
+
header_bytes = int(header.nbytes)
|
| 58 |
+
return {
|
| 59 |
+
"chunk_bytes": data_bytes + header_bytes,
|
| 60 |
+
"chunk_size": num_blocks,
|
| 61 |
+
"dim": int(len(chunk_data)),
|
| 62 |
+
"filename": filename,
|
| 63 |
+
}, header_bytes, data_bytes
|
| 64 |
+
|
| 65 |
+
|
| 66 |
+
def _load_resume_state(output_dir):
|
| 67 |
+
"""Check for existing chunks to support resuming."""
|
| 68 |
+
index_path = os.path.join(output_dir, "index.json.partial")
|
| 69 |
+
if not os.path.exists(index_path):
|
| 70 |
+
return [], 0, 0
|
| 71 |
+
with open(index_path) as f:
|
| 72 |
+
state = json.load(f)
|
| 73 |
+
chunks = state.get("chunks", [])
|
| 74 |
+
tokens_written = sum(c["dim"] for c in chunks)
|
| 75 |
+
return chunks, len(chunks), tokens_written
|
| 76 |
+
|
| 77 |
+
|
| 78 |
+
def _save_resume_state(output_dir, chunks_metadata):
|
| 79 |
+
"""Save partial progress for resume."""
|
| 80 |
+
path = os.path.join(output_dir, "index.json.partial")
|
| 81 |
+
with open(path, "w") as f:
|
| 82 |
+
json.dump({"chunks": chunks_metadata}, f)
|
| 83 |
+
|
| 84 |
+
|
| 85 |
+
def prepare_litdata(filtered_dir: str, output_dir: str, tokenizer_path: str, label: str):
|
| 86 |
+
"""Tokenize filtered text and write chunks in streaming fashion (constant RAM)."""
|
| 87 |
+
print(f"\n{'='*60}")
|
| 88 |
+
print(f"Preparing LitData [{label}]")
|
| 89 |
+
print(f"Input: {filtered_dir}")
|
| 90 |
+
print(f"Output: {output_dir}")
|
| 91 |
+
print(f"Tokenizer: {tokenizer_path}")
|
| 92 |
+
print(f"{'='*60}\n")
|
| 93 |
+
|
| 94 |
+
if not Path(filtered_dir).exists():
|
| 95 |
+
print(f"ERROR: Filtered directory not found: {filtered_dir}")
|
| 96 |
+
print("Run filter_datasets.py first!")
|
| 97 |
+
return
|
| 98 |
+
|
| 99 |
+
os.makedirs(output_dir, exist_ok=True)
|
| 100 |
+
tokenizer = Tokenizer(Path(tokenizer_path))
|
| 101 |
+
|
| 102 |
+
# How many tokens fit in one chunk (aligned to BLOCK_SIZE)
|
| 103 |
+
dtype_size = DTYPE().itemsize
|
| 104 |
+
tokens_per_chunk = CHUNK_BYTES_TARGET // dtype_size
|
| 105 |
+
tokens_per_chunk = (tokens_per_chunk // BLOCK_SIZE) * BLOCK_SIZE
|
| 106 |
+
|
| 107 |
+
# Check for resume state
|
| 108 |
+
chunks_metadata, chunk_idx, tokens_to_skip = _load_resume_state(output_dir)
|
| 109 |
+
if tokens_to_skip > 0:
|
| 110 |
+
print(f"RESUMING: Found {chunk_idx} existing chunks ({tokens_to_skip:,} tokens)")
|
| 111 |
+
print(f" Will skip {tokens_to_skip:,} tokens then continue writing chunks\n")
|
| 112 |
+
|
| 113 |
+
# Token buffer β holds at most one chunk worth of tokens (~64MB)
|
| 114 |
+
token_buf = np.empty(tokens_per_chunk + BLOCK_SIZE * 512, dtype=DTYPE)
|
| 115 |
+
buf_pos = 0 # current write position in buffer
|
| 116 |
+
total_tokens_written = sum(c["dim"] for c in chunks_metadata)
|
| 117 |
+
total_tokens_seen = 0
|
| 118 |
+
total_texts = 0
|
| 119 |
+
skipped_tokens = 0
|
| 120 |
+
|
| 121 |
+
parquet_files = sorted(Path(filtered_dir).glob("*.parquet"))
|
| 122 |
+
t0 = time.time()
|
| 123 |
+
|
| 124 |
+
print("Tokenizing & writing chunks (streaming)...")
|
| 125 |
+
for pf in parquet_files:
|
| 126 |
+
parquet_file = pq.ParquetFile(str(pf))
|
| 127 |
+
file_count = 0
|
| 128 |
+
for batch in parquet_file.iter_batches(batch_size=4096, columns=["text"]):
|
| 129 |
+
texts = batch.column("text").to_pylist()
|
| 130 |
+
for text in texts:
|
| 131 |
+
tokens = tokenizer.encode(text, bos=False, eos=True)
|
| 132 |
+
tok_array = tokens.numpy().astype(DTYPE) if hasattr(tokens, 'numpy') else np.array(tokens.tolist(), dtype=DTYPE)
|
| 133 |
+
n = len(tok_array)
|
| 134 |
+
total_tokens_seen += n
|
| 135 |
+
|
| 136 |
+
# If resuming, skip tokens already written
|
| 137 |
+
if skipped_tokens < tokens_to_skip:
|
| 138 |
+
remaining_skip = tokens_to_skip - skipped_tokens
|
| 139 |
+
if n <= remaining_skip:
|
| 140 |
+
skipped_tokens += n
|
| 141 |
+
file_count += 1
|
| 142 |
+
continue
|
| 143 |
+
else:
|
| 144 |
+
tok_array = tok_array[int(remaining_skip):]
|
| 145 |
+
skipped_tokens = tokens_to_skip
|
| 146 |
+
n = len(tok_array)
|
| 147 |
+
|
| 148 |
+
# Append to buffer
|
| 149 |
+
if buf_pos + n > len(token_buf):
|
| 150 |
+
# Grow buffer if needed (rare)
|
| 151 |
+
token_buf = np.concatenate([token_buf[:buf_pos], np.empty(max(n, BLOCK_SIZE * 512), dtype=DTYPE)])
|
| 152 |
+
|
| 153 |
+
token_buf[buf_pos:buf_pos + n] = tok_array
|
| 154 |
+
buf_pos += n
|
| 155 |
+
file_count += 1
|
| 156 |
+
|
| 157 |
+
# Flush full chunks from buffer
|
| 158 |
+
while buf_pos >= tokens_per_chunk:
|
| 159 |
+
num_blocks = tokens_per_chunk // BLOCK_SIZE
|
| 160 |
+
chunk_array = token_buf[:tokens_per_chunk].copy()
|
| 161 |
+
|
| 162 |
+
meta, hdr_b, data_b = _write_chunk(output_dir, chunk_idx, chunk_array, num_blocks)
|
| 163 |
+
chunks_metadata.append(meta)
|
| 164 |
+
total_tokens_written += meta["dim"]
|
| 165 |
+
_save_resume_state(output_dir, chunks_metadata)
|
| 166 |
+
|
| 167 |
+
print(f" chunk-0-{chunk_idx}.bin: {num_blocks} blocks, "
|
| 168 |
+
f"{meta['dim']:,} tokens (header: {hdr_b/1024:.1f} KB, "
|
| 169 |
+
f"data: {data_b/1024/1024:.1f} MB) | "
|
| 170 |
+
f"total: {total_tokens_written:,}")
|
| 171 |
+
chunk_idx += 1
|
| 172 |
+
|
| 173 |
+
# Shift remaining tokens to front of buffer
|
| 174 |
+
leftover = buf_pos - tokens_per_chunk
|
| 175 |
+
if leftover > 0:
|
| 176 |
+
token_buf[:leftover] = token_buf[tokens_per_chunk:tokens_per_chunk + leftover]
|
| 177 |
+
buf_pos = leftover
|
| 178 |
+
|
| 179 |
+
total_texts += file_count
|
| 180 |
+
parquet_file.close()
|
| 181 |
+
elapsed = time.time() - t0
|
| 182 |
+
print(f" [{pf.name}] {file_count:,} texts | "
|
| 183 |
+
f"total tokens written: {total_tokens_written:,} | "
|
| 184 |
+
f"elapsed: {elapsed:.0f}s")
|
| 185 |
+
|
| 186 |
+
# Flush remaining buffer as final chunk (if enough for at least 1 block)
|
| 187 |
+
remaining_blocks = buf_pos // BLOCK_SIZE
|
| 188 |
+
if remaining_blocks > 0:
|
| 189 |
+
final_tokens = remaining_blocks * BLOCK_SIZE
|
| 190 |
+
chunk_array = token_buf[:final_tokens].copy()
|
| 191 |
+
meta, hdr_b, data_b = _write_chunk(output_dir, chunk_idx, chunk_array, remaining_blocks)
|
| 192 |
+
chunks_metadata.append(meta)
|
| 193 |
+
total_tokens_written += meta["dim"]
|
| 194 |
+
print(f" chunk-0-{chunk_idx}.bin (final): {remaining_blocks} blocks, "
|
| 195 |
+
f"{meta['dim']:,} tokens (header: {hdr_b/1024:.1f} KB, "
|
| 196 |
+
f"data: {data_b/1024/1024:.1f} MB)")
|
| 197 |
+
chunk_idx += 1
|
| 198 |
+
|
| 199 |
+
# Write final index.json
|
| 200 |
+
index = {
|
| 201 |
+
"chunks": chunks_metadata,
|
| 202 |
+
"config": {
|
| 203 |
+
"chunk_bytes": CHUNK_BYTES_TARGET,
|
| 204 |
+
"chunk_size": None,
|
| 205 |
+
"compression": None,
|
| 206 |
+
"data_format": [f"no_header_tensor:{DTYPE_INDEX}"],
|
| 207 |
+
"data_spec": None,
|
| 208 |
+
"encryption": None,
|
| 209 |
+
"item_loader": "TokensLoader",
|
| 210 |
+
},
|
| 211 |
+
"updated_at": str(time.time()),
|
| 212 |
+
}
|
| 213 |
+
with open(os.path.join(output_dir, "index.json"), "w") as f:
|
| 214 |
+
json.dump(index, f, indent=2)
|
| 215 |
+
|
| 216 |
+
# Clean up partial state
|
| 217 |
+
partial_path = os.path.join(output_dir, "index.json.partial")
|
| 218 |
+
if os.path.exists(partial_path):
|
| 219 |
+
os.remove(partial_path)
|
| 220 |
+
|
| 221 |
+
elapsed = time.time() - t0
|
| 222 |
+
print(f"\n--- LitData preparation [{label}] complete ---")
|
| 223 |
+
print(f"Chunks: {chunk_idx}")
|
| 224 |
+
print(f"Total tokens: {total_tokens_written:,}")
|
| 225 |
+
print(f"Total blocks: {total_tokens_written // BLOCK_SIZE:,}")
|
| 226 |
+
print(f"Texts processed: {total_texts:,}")
|
| 227 |
+
print(f"Time: {elapsed:.0f}s ({total_tokens_written/elapsed:.0f} tok/s)")
|
| 228 |
+
print(f"Output: {output_dir}")
|
| 229 |
+
|
| 230 |
+
|
| 231 |
+
def main():
|
| 232 |
+
parser = argparse.ArgumentParser(description="Prepare LitData from filtered parquet")
|
| 233 |
+
parser.add_argument(
|
| 234 |
+
"--mode",
|
| 235 |
+
choices=["small", "full", "both", "english"],
|
| 236 |
+
default="small",
|
| 237 |
+
help="Which dataset to prepare (english = ultra-clean English corpus)",
|
| 238 |
+
)
|
| 239 |
+
parser.add_argument(
|
| 240 |
+
"--tokenizer_path",
|
| 241 |
+
type=str,
|
| 242 |
+
default="Base/checkpoints/EleutherAI/pythia-160m",
|
| 243 |
+
help="Path to tokenizer directory",
|
| 244 |
+
)
|
| 245 |
+
args = parser.parse_args()
|
| 246 |
+
|
| 247 |
+
if args.mode in ("small", "both"):
|
| 248 |
+
prepare_litdata(
|
| 249 |
+
filtered_dir="Base/data/filtered_10m",
|
| 250 |
+
output_dir="Base/data/litdata_10m",
|
| 251 |
+
tokenizer_path=args.tokenizer_path,
|
| 252 |
+
label="SMALL (10M)",
|
| 253 |
+
)
|
| 254 |
+
|
| 255 |
+
if args.mode in ("full", "both"):
|
| 256 |
+
prepare_litdata(
|
| 257 |
+
filtered_dir="Base/data/filtered_3b",
|
| 258 |
+
output_dir="Base/data/litdata_3b",
|
| 259 |
+
tokenizer_path=args.tokenizer_path,
|
| 260 |
+
label="FULL (3B)",
|
| 261 |
+
)
|
| 262 |
+
|
| 263 |
+
if args.mode == "english":
|
| 264 |
+
prepare_litdata(
|
| 265 |
+
filtered_dir="Base/data/filtered_english",
|
| 266 |
+
output_dir="Base/data/litdata_english",
|
| 267 |
+
tokenizer_path=args.tokenizer_path,
|
| 268 |
+
label="ENGLISH (ultra-clean)",
|
| 269 |
+
)
|
| 270 |
+
|
| 271 |
+
print("\nDone! Next step: run litgpt pretrain with the appropriate config.")
|
| 272 |
+
|
| 273 |
+
|
| 274 |
+
if __name__ == "__main__":
|
| 275 |
+
main()
|
Base/scripts/reclean_3b.py
ADDED
|
@@ -0,0 +1,593 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""
|
| 3 |
+
Memory-efficient, multiprocessed reclean of litdata_3b (~3B tokens).
|
| 4 |
+
|
| 5 |
+
Problem: The old script loaded ALL 3B tokens β decoded 4.3M docs β re-tokenized
|
| 6 |
+
into a Python list β OOM at ~12 GB list + ~15 GB decoded text.
|
| 7 |
+
|
| 8 |
+
Solution: Process in chunk batches (10 chunks β 168M tokens β 670 MB).
|
| 9 |
+
1. Read 10 chunks at a time
|
| 10 |
+
2. Split tokens into documents (carry partial docs across batches)
|
| 11 |
+
3. Decode with tokenizer.decode_batch (multithreaded Rust)
|
| 12 |
+
4. Deep clean + smart filter with multiprocessing Pool (all CPU cores)
|
| 13 |
+
5. Re-tokenize with encode_batch (multithreaded Rust)
|
| 14 |
+
6. Stream output to litdata chunks incrementally (never accumulate)
|
| 15 |
+
|
| 16 |
+
Peak memory: ~4-5 GB instead of 40+ GB.
|
| 17 |
+
Output: Base/data/litdata_3b_clean/ (separate, not merged)
|
| 18 |
+
"""
|
| 19 |
+
|
| 20 |
+
import json
|
| 21 |
+
import os
|
| 22 |
+
import re
|
| 23 |
+
import sys
|
| 24 |
+
import time
|
| 25 |
+
import unicodedata
|
| 26 |
+
from pathlib import Path
|
| 27 |
+
from multiprocessing import Pool, cpu_count
|
| 28 |
+
from collections import Counter
|
| 29 |
+
|
| 30 |
+
import numpy as np
|
| 31 |
+
from tokenizers import Tokenizer
|
| 32 |
+
|
| 33 |
+
ROOT = Path(__file__).resolve().parent.parent.parent
|
| 34 |
+
BLOCK_SIZE = 1025
|
| 35 |
+
DTYPE = np.int32
|
| 36 |
+
CHUNK_BYTES_TARGET = 64 * 1024 * 1024
|
| 37 |
+
EOS_TOKEN_ID = 0
|
| 38 |
+
|
| 39 |
+
CHUNKS_PER_BATCH = 10 # ~670 MB per batch
|
| 40 |
+
NUM_WORKERS = max(1, cpu_count() - 2) # leave 2 cores for main + I/O
|
| 41 |
+
DECODE_BATCH = 8000 # docs per decode_batch call
|
| 42 |
+
ENCODE_BATCH = 8000 # docs per encode_batch call
|
| 43 |
+
|
| 44 |
+
DATA_DIR = ROOT / "Base" / "data"
|
| 45 |
+
INPUT_DIR = DATA_DIR / "litdata_3b"
|
| 46 |
+
OUTPUT_DIR = DATA_DIR / "litdata_3b_clean"
|
| 47 |
+
TOKENIZER_PATH = str(ROOT / "Base" / "checkpoints" / "EleutherAI" / "pythia-160m" / "tokenizer.json")
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
# ==============================================================================
|
| 51 |
+
# TEXT CLEANING PIPELINE (deep_clean + smart_filter)
|
| 52 |
+
# ==============================================================================
|
| 53 |
+
|
| 54 |
+
CONTROL_CHARS = [
|
| 55 |
+
"\x00", "\x01", "\x02", "\x03", "\x04", "\x05", "\x06", "\x07",
|
| 56 |
+
"\x08", "\x0b", "\x0c", "\x0e", "\x0f", "\x10", "\x11", "\x12",
|
| 57 |
+
"\x13", "\x14", "\x15", "\x16", "\x17", "\x18", "\x19", "\x1a",
|
| 58 |
+
"\x1b", "\x1c", "\x1d", "\x1e", "\x1f", "\x7f", "\ufeff", "\ufffd",
|
| 59 |
+
]
|
| 60 |
+
|
| 61 |
+
HTML_ENTITIES = [
|
| 62 |
+
("&", "&"), ("<", "<"), (">", ">"),
|
| 63 |
+
(""", '"'), ("'", "'"), ("'", "'"),
|
| 64 |
+
(" ", " "), ("—", " - "), ("–", "-"),
|
| 65 |
+
("…", "..."), ("«", '"'), ("»", '"'),
|
| 66 |
+
("•", "- "), ("·", " "), ("©", "(c)"),
|
| 67 |
+
("®", "(R)"), ("™", "(TM)"), ("°", " degrees"),
|
| 68 |
+
]
|
| 69 |
+
|
| 70 |
+
RE_URL = re.compile(r'https?://\S+|www\.\S+', re.I)
|
| 71 |
+
RE_EMAIL = re.compile(r'\b[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}\b')
|
| 72 |
+
RE_FILE_PATH = re.compile(r'(?:[A-Z]:\\|/(?:home|usr|var|etc|opt)/)\S+')
|
| 73 |
+
RE_HTML_TAG = re.compile(r'</?[a-zA-Z][a-zA-Z0-9]*(?:\s[^>]*)?\s*/?>')
|
| 74 |
+
RE_HTML_COMMENT = re.compile(r'<!--.*?-->', re.DOTALL)
|
| 75 |
+
RE_CODE_BLOCK = re.compile(r'```[\s\S]*?```')
|
| 76 |
+
RE_IMPORT = re.compile(r'^(?:import |from \S+ import |#include |using namespace |require\()', re.M)
|
| 77 |
+
RE_REPEATED_LINE = re.compile(r'^(.{20,})\n(?:\1\n?)+', re.M)
|
| 78 |
+
RE_REPEATED_PUNCT = re.compile(r'([!?.])\1{3,}')
|
| 79 |
+
RE_REPEATED_CHAR = re.compile(r'(.)\1{5,}')
|
| 80 |
+
RE_REPEATED_WORD = re.compile(r'\b(\w+)(?:\s+\1){2,}\b', re.I)
|
| 81 |
+
RE_MULTI_NEWLINE = re.compile(r'\n{4,}')
|
| 82 |
+
RE_MULTI_SPACE = re.compile(r'[ \t]{2,}')
|
| 83 |
+
RE_TRAILING_SPACE = re.compile(r'[ \t]+$', re.M)
|
| 84 |
+
RE_NO_SPACE_AFTER_PERIOD = re.compile(r'([.!?])([A-Z])')
|
| 85 |
+
RE_DOUBLE_PERIOD = re.compile(r'\.{2}(?!\.)')
|
| 86 |
+
RE_SPACE_BEFORE_PUNCT = re.compile(r'\s+([.,;:!?])')
|
| 87 |
+
|
| 88 |
+
# Smart filter patterns
|
| 89 |
+
RE_CJK = re.compile(r'[\u4e00-\u9fff\u3040-\u309f\u30a0-\u30ff\uac00-\ud7af]{3,}')
|
| 90 |
+
RE_ARABIC = re.compile(r'[\u0600-\u06ff]{5,}')
|
| 91 |
+
RE_CYRILLIC = re.compile(r'[\u0400-\u04ff]{5,}')
|
| 92 |
+
RE_DEVANAGARI = re.compile(r'[\u0900-\u097f]{5,}')
|
| 93 |
+
RE_RESIDUAL_CODE = re.compile(
|
| 94 |
+
r'(function\s*\(|var\s+\w+\s*=|console\.log|document\.get|if\s*\(\s*\w+\s*[!=]==)', re.I
|
| 95 |
+
)
|
| 96 |
+
|
| 97 |
+
RE_COOKIE_LINE = re.compile(r'^.*(?:cookie|cookies)\s+(?:policy|consent|notice|preferences|settings).*$', re.I | re.M)
|
| 98 |
+
RE_SUBSCRIBE_LINE = re.compile(r'^.*(?:subscribe|sign\s*up\s+(?:for|to)\s+(?:our|the)\s+newsletter|unsubscribe|opt[\s-]*out\s+of).*$', re.I | re.M)
|
| 99 |
+
RE_CLICKBAIT_LINE = re.compile(r'^.*(?:you\s+won\'?t\s+believe|click\s+here|read\s+more\s*\.{0,3}$|share\s+this\s+(?:article|post|story)|trending\s+now|sponsored\s+content|advertisement).*$', re.I | re.M)
|
| 100 |
+
RE_SOCIAL_LINE = re.compile(r'^.*(?:follow\s+us\s+on|share\s+on\s+(?:facebook|twitter|linkedin|instagram)|like\s+us\s+on|tweet\s+this).*$', re.I | re.M)
|
| 101 |
+
RE_NAV_LINE = re.compile(r'^.*(?:skip\s+to\s+(?:main\s+)?content|back\s+to\s+top|previous\s+article|next\s+article|related\s+(?:articles|posts)).*$', re.I | re.M)
|
| 102 |
+
RE_LOGIN_LINE = re.compile(r'^.*(?:log\s*in\s+to\s+(?:your|an)\s+account|create\s+(?:a\s+)?(?:free\s+)?account|forgot\s+(?:your\s+)?password|already\s+(?:a\s+)?member).*$', re.I | re.M)
|
| 103 |
+
RE_COMMENT_LINE = re.compile(r'^.*(?:leave\s+a\s+(?:comment|reply)|post\s+a\s+comment|\d+\s+comments?$|logged\s+in\s+as).*$', re.I | re.M)
|
| 104 |
+
RE_COPYRIGHT_LINE = re.compile(r'^.*(?:all\s+rights\s+reserved|\(c\)\s*\d{4}|copyright\s+\d{4}).*$', re.I | re.M)
|
| 105 |
+
|
| 106 |
+
|
| 107 |
+
def deep_clean(text):
|
| 108 |
+
"""Full deep cleaning pipeline."""
|
| 109 |
+
if not text or len(text.strip()) < 30:
|
| 110 |
+
return None
|
| 111 |
+
|
| 112 |
+
text = unicodedata.normalize("NFKC", text)
|
| 113 |
+
for ch in CONTROL_CHARS:
|
| 114 |
+
text = text.replace(ch, "")
|
| 115 |
+
for old, new in HTML_ENTITIES:
|
| 116 |
+
text = text.replace(old, new)
|
| 117 |
+
|
| 118 |
+
text = RE_HTML_COMMENT.sub("", text)
|
| 119 |
+
text = RE_HTML_TAG.sub("", text)
|
| 120 |
+
text = RE_URL.sub("", text)
|
| 121 |
+
text = RE_EMAIL.sub("", text)
|
| 122 |
+
text = RE_FILE_PATH.sub("", text)
|
| 123 |
+
text = RE_CODE_BLOCK.sub("", text)
|
| 124 |
+
|
| 125 |
+
text = RE_REPEATED_LINE.sub(r'\1', text)
|
| 126 |
+
text = RE_REPEATED_PUNCT.sub(r'\1\1\1', text)
|
| 127 |
+
text = RE_REPEATED_CHAR.sub(r'\1\1\1', text)
|
| 128 |
+
text = RE_REPEATED_WORD.sub(r'\1', text)
|
| 129 |
+
|
| 130 |
+
text = text.replace('\t', ' ')
|
| 131 |
+
text = RE_TRAILING_SPACE.sub('', text)
|
| 132 |
+
text = RE_MULTI_SPACE.sub(' ', text)
|
| 133 |
+
text = RE_MULTI_NEWLINE.sub('\n\n\n', text)
|
| 134 |
+
|
| 135 |
+
text = RE_DOUBLE_PERIOD.sub('.', text)
|
| 136 |
+
text = RE_NO_SPACE_AFTER_PERIOD.sub(r'\1 \2', text)
|
| 137 |
+
text = RE_SPACE_BEFORE_PUNCT.sub(r'\1', text)
|
| 138 |
+
|
| 139 |
+
text = text.replace('\u2018', "'").replace('\u2019', "'")
|
| 140 |
+
text = text.replace('\u201c', '"').replace('\u201d', '"')
|
| 141 |
+
text = text.replace('\u2013', '-').replace('\u2014', ' - ')
|
| 142 |
+
text = text.replace('\u2026', '...')
|
| 143 |
+
text = text.replace('\u2022', '- ')
|
| 144 |
+
text = text.replace('\u00b7', ' ')
|
| 145 |
+
text = text.replace('\u00a0', ' ')
|
| 146 |
+
|
| 147 |
+
lines = text.split('\n')
|
| 148 |
+
clean_lines = []
|
| 149 |
+
for line in lines:
|
| 150 |
+
line = line.strip()
|
| 151 |
+
if not line:
|
| 152 |
+
clean_lines.append('')
|
| 153 |
+
continue
|
| 154 |
+
if len(line) > 10:
|
| 155 |
+
alpha_count = sum(1 for c in line if c.isalpha())
|
| 156 |
+
if alpha_count / len(line) < 0.40:
|
| 157 |
+
continue
|
| 158 |
+
if line.count('|') > 3 or line.count('{') > 2 or line.count('}') > 2:
|
| 159 |
+
continue
|
| 160 |
+
if RE_IMPORT.match(line):
|
| 161 |
+
continue
|
| 162 |
+
if line and line[0].isalpha() and line[0].islower():
|
| 163 |
+
if not clean_lines or clean_lines[-1] == '' or clean_lines[-1].rstrip().endswith(('.', '!', '?', ':')):
|
| 164 |
+
line = line[0].upper() + line[1:]
|
| 165 |
+
clean_lines.append(line)
|
| 166 |
+
|
| 167 |
+
text = '\n'.join(clean_lines)
|
| 168 |
+
text = text.strip()
|
| 169 |
+
|
| 170 |
+
paragraphs = text.split('\n\n')
|
| 171 |
+
seen = set()
|
| 172 |
+
unique_paragraphs = []
|
| 173 |
+
for p in paragraphs:
|
| 174 |
+
p_stripped = p.strip()
|
| 175 |
+
if not p_stripped:
|
| 176 |
+
continue
|
| 177 |
+
p_key = ' '.join(p_stripped.lower().split())
|
| 178 |
+
if p_key not in seen:
|
| 179 |
+
seen.add(p_key)
|
| 180 |
+
unique_paragraphs.append(p_stripped)
|
| 181 |
+
text = '\n\n'.join(unique_paragraphs)
|
| 182 |
+
|
| 183 |
+
text = text.strip()
|
| 184 |
+
if len(text) < 50:
|
| 185 |
+
return None
|
| 186 |
+
if len(text.split()) < 10:
|
| 187 |
+
return None
|
| 188 |
+
ascii_count = sum(1 for c in text if ord(c) < 128)
|
| 189 |
+
if ascii_count / max(len(text), 1) < 0.85:
|
| 190 |
+
return None
|
| 191 |
+
|
| 192 |
+
return text
|
| 193 |
+
|
| 194 |
+
|
| 195 |
+
def smart_filter(text):
|
| 196 |
+
"""Smart cleanup: returns (keep, cleaned_text, reason)."""
|
| 197 |
+
words = text.split()
|
| 198 |
+
word_count = len(words)
|
| 199 |
+
if word_count < 50:
|
| 200 |
+
return False, text, "too short"
|
| 201 |
+
|
| 202 |
+
scripts = []
|
| 203 |
+
if RE_CJK.search(text): scripts.append("CJK")
|
| 204 |
+
if RE_ARABIC.search(text): scripts.append("Arabic")
|
| 205 |
+
if RE_CYRILLIC.search(text): scripts.append("Cyrillic")
|
| 206 |
+
if RE_DEVANAGARI.search(text): scripts.append("Devanagari")
|
| 207 |
+
if scripts:
|
| 208 |
+
return False, text, "non-English"
|
| 209 |
+
|
| 210 |
+
if word_count > 50:
|
| 211 |
+
unique_ratio = len(set(w.lower() for w in words)) / word_count
|
| 212 |
+
if unique_ratio < 0.20:
|
| 213 |
+
return False, text, "repetitive"
|
| 214 |
+
|
| 215 |
+
code_matches = RE_RESIDUAL_CODE.findall(text)
|
| 216 |
+
if len(code_matches) >= 5:
|
| 217 |
+
return False, text, "residual code"
|
| 218 |
+
|
| 219 |
+
for pattern in [RE_COOKIE_LINE, RE_SUBSCRIBE_LINE, RE_CLICKBAIT_LINE,
|
| 220 |
+
RE_SOCIAL_LINE, RE_NAV_LINE, RE_LOGIN_LINE,
|
| 221 |
+
RE_COMMENT_LINE, RE_COPYRIGHT_LINE]:
|
| 222 |
+
text = pattern.sub('', text)
|
| 223 |
+
|
| 224 |
+
lines = text.split('\n')
|
| 225 |
+
clean_lines = []
|
| 226 |
+
for line in lines:
|
| 227 |
+
stripped = line.strip()
|
| 228 |
+
if stripped and len(stripped) > 10:
|
| 229 |
+
digit_count = sum(1 for c in stripped if c.isdigit() or c in ' ,.\t-+/%$')
|
| 230 |
+
if digit_count / len(stripped) > 0.80:
|
| 231 |
+
continue
|
| 232 |
+
clean_lines.append(line)
|
| 233 |
+
text = '\n'.join(clean_lines)
|
| 234 |
+
text = re.sub(r'\n{3,}', '\n\n', text)
|
| 235 |
+
text = text.strip()
|
| 236 |
+
|
| 237 |
+
if len(text.split()) < 50:
|
| 238 |
+
return False, text, "too short after stripping"
|
| 239 |
+
|
| 240 |
+
return True, text, None
|
| 241 |
+
|
| 242 |
+
|
| 243 |
+
def clean_and_filter(text):
|
| 244 |
+
"""Combined deep_clean + smart_filter for multiprocessing.
|
| 245 |
+
Returns (cleaned_text_or_None, drop_reason_or_None, boilerplate_chars_stripped).
|
| 246 |
+
"""
|
| 247 |
+
result = deep_clean(text)
|
| 248 |
+
if result is None:
|
| 249 |
+
return None, "deep_clean_drop", 0
|
| 250 |
+
|
| 251 |
+
keep, stripped, reason = smart_filter(result)
|
| 252 |
+
if not keep:
|
| 253 |
+
return None, reason, 0
|
| 254 |
+
|
| 255 |
+
boilerplate = len(result) - len(stripped)
|
| 256 |
+
return stripped, None, boilerplate
|
| 257 |
+
|
| 258 |
+
|
| 259 |
+
# ==============================================================================
|
| 260 |
+
# STREAMING CHUNK WRITER (never accumulates full token stream)
|
| 261 |
+
# ==============================================================================
|
| 262 |
+
|
| 263 |
+
class StreamingChunkWriter:
|
| 264 |
+
"""Writes litdata chunks incrementally. Flushes when buffer hits target size."""
|
| 265 |
+
|
| 266 |
+
def __init__(self, output_dir, config):
|
| 267 |
+
self.output_dir = Path(output_dir)
|
| 268 |
+
os.makedirs(str(self.output_dir), exist_ok=True)
|
| 269 |
+
self.config = config
|
| 270 |
+
self.dtype_size = DTYPE().itemsize
|
| 271 |
+
self.tokens_per_chunk = (CHUNK_BYTES_TARGET // self.dtype_size // BLOCK_SIZE) * BLOCK_SIZE
|
| 272 |
+
self.buffer = []
|
| 273 |
+
self.chunks_metadata = []
|
| 274 |
+
self.chunk_idx = 0
|
| 275 |
+
self.total_tokens = 0
|
| 276 |
+
|
| 277 |
+
def add_document(self, token_ids):
|
| 278 |
+
"""Add one document's token IDs + EOS. Flushes chunks as needed."""
|
| 279 |
+
self.buffer.extend(token_ids)
|
| 280 |
+
self.buffer.append(EOS_TOKEN_ID)
|
| 281 |
+
# Flush full chunks
|
| 282 |
+
while len(self.buffer) >= self.tokens_per_chunk:
|
| 283 |
+
self._flush_chunk()
|
| 284 |
+
|
| 285 |
+
def add_tokens_batch(self, encoded_batch):
|
| 286 |
+
"""Add a batch of encoded documents efficiently."""
|
| 287 |
+
for enc in encoded_batch:
|
| 288 |
+
ids = enc.ids
|
| 289 |
+
self.buffer.extend(ids)
|
| 290 |
+
self.buffer.append(EOS_TOKEN_ID)
|
| 291 |
+
while len(self.buffer) >= self.tokens_per_chunk:
|
| 292 |
+
self._flush_chunk()
|
| 293 |
+
|
| 294 |
+
def _flush_chunk(self):
|
| 295 |
+
"""Write one chunk from the buffer."""
|
| 296 |
+
if len(self.buffer) < BLOCK_SIZE:
|
| 297 |
+
return
|
| 298 |
+
|
| 299 |
+
take = min(len(self.buffer), self.tokens_per_chunk)
|
| 300 |
+
num_blocks = take // BLOCK_SIZE
|
| 301 |
+
if num_blocks == 0:
|
| 302 |
+
return
|
| 303 |
+
actual_tokens = num_blocks * BLOCK_SIZE
|
| 304 |
+
|
| 305 |
+
chunk_data = np.array(self.buffer[:actual_tokens], dtype=DTYPE)
|
| 306 |
+
self.buffer = self.buffer[actual_tokens:]
|
| 307 |
+
|
| 308 |
+
filename = f"chunk-0-{self.chunk_idx}.bin"
|
| 309 |
+
filepath = self.output_dir / filename
|
| 310 |
+
|
| 311 |
+
header_num = np.array([num_blocks], dtype=np.uint32)
|
| 312 |
+
offsets = np.arange(num_blocks + 1, dtype=np.uint32) * (BLOCK_SIZE * self.dtype_size)
|
| 313 |
+
header = np.concatenate([header_num, offsets])
|
| 314 |
+
|
| 315 |
+
with open(filepath, "wb") as f:
|
| 316 |
+
header.tofile(f)
|
| 317 |
+
chunk_data.tofile(f)
|
| 318 |
+
|
| 319 |
+
meta = {
|
| 320 |
+
"chunk_bytes": int(header.nbytes + chunk_data.nbytes),
|
| 321 |
+
"chunk_size": num_blocks,
|
| 322 |
+
"dim": int(actual_tokens),
|
| 323 |
+
"filename": filename,
|
| 324 |
+
}
|
| 325 |
+
self.chunks_metadata.append(meta)
|
| 326 |
+
self.total_tokens += actual_tokens
|
| 327 |
+
self.chunk_idx += 1
|
| 328 |
+
|
| 329 |
+
if self.chunk_idx % 25 == 0:
|
| 330 |
+
print(f" Flushed chunk {self.chunk_idx} ({self.total_tokens:,} tokens written)")
|
| 331 |
+
|
| 332 |
+
def finalize(self):
|
| 333 |
+
"""Flush remaining buffer and write index.json."""
|
| 334 |
+
while len(self.buffer) >= BLOCK_SIZE:
|
| 335 |
+
self._flush_chunk()
|
| 336 |
+
|
| 337 |
+
# Discard any remaining tokens that don't fill a block
|
| 338 |
+
discarded = len(self.buffer)
|
| 339 |
+
self.buffer = []
|
| 340 |
+
|
| 341 |
+
index = {
|
| 342 |
+
"chunks": self.chunks_metadata,
|
| 343 |
+
"config": self.config,
|
| 344 |
+
"updated_at": str(time.time()),
|
| 345 |
+
}
|
| 346 |
+
with open(self.output_dir / "index.json", "w") as f:
|
| 347 |
+
json.dump(index, f, indent=2)
|
| 348 |
+
|
| 349 |
+
return self.total_tokens, discarded
|
| 350 |
+
|
| 351 |
+
|
| 352 |
+
# ==============================================================================
|
| 353 |
+
# CHUNK READING
|
| 354 |
+
# ==============================================================================
|
| 355 |
+
|
| 356 |
+
def read_chunk_tokens(litdata_dir, chunk_meta):
|
| 357 |
+
"""Read a single chunk's token data (no header)."""
|
| 358 |
+
chunk_path = litdata_dir / chunk_meta["filename"]
|
| 359 |
+
n_blocks = chunk_meta["chunk_size"]
|
| 360 |
+
header_ints = 1 + n_blocks + 1
|
| 361 |
+
header_bytes = header_ints * 4
|
| 362 |
+
|
| 363 |
+
with open(chunk_path, "rb") as f:
|
| 364 |
+
f.seek(header_bytes)
|
| 365 |
+
data = np.fromfile(f, dtype=DTYPE, count=chunk_meta["dim"])
|
| 366 |
+
return data
|
| 367 |
+
|
| 368 |
+
|
| 369 |
+
def split_documents_from_tokens(token_array):
|
| 370 |
+
"""Split token array by EOS into list of per-document token lists (as Python lists)."""
|
| 371 |
+
eos_positions = np.where(token_array == EOS_TOKEN_ID)[0]
|
| 372 |
+
docs = []
|
| 373 |
+
start = 0
|
| 374 |
+
for eos_pos in eos_positions:
|
| 375 |
+
if eos_pos > start:
|
| 376 |
+
docs.append(token_array[start:eos_pos].tolist())
|
| 377 |
+
start = eos_pos + 1
|
| 378 |
+
# Return remaining tokens after last EOS (partial doc carry-over)
|
| 379 |
+
remainder = token_array[start:] if start < len(token_array) else np.array([], dtype=DTYPE)
|
| 380 |
+
return docs, remainder
|
| 381 |
+
|
| 382 |
+
|
| 383 |
+
# ==============================================================================
|
| 384 |
+
# MAIN PIPELINE
|
| 385 |
+
# ==============================================================================
|
| 386 |
+
|
| 387 |
+
def main():
|
| 388 |
+
t_start = time.time()
|
| 389 |
+
|
| 390 |
+
print("Loading tokenizer...")
|
| 391 |
+
tokenizer = Tokenizer.from_file(TOKENIZER_PATH)
|
| 392 |
+
|
| 393 |
+
print(f"\n{'='*75}")
|
| 394 |
+
print(f" RECLEAN litdata_3b (Memory-Efficient + Multiprocessing)")
|
| 395 |
+
print(f" Input: {INPUT_DIR}")
|
| 396 |
+
print(f" Output: {OUTPUT_DIR}")
|
| 397 |
+
print(f" Workers: {NUM_WORKERS} CPU cores")
|
| 398 |
+
print(f" Batch: {CHUNKS_PER_BATCH} chunks at a time")
|
| 399 |
+
print(f"{'='*75}")
|
| 400 |
+
|
| 401 |
+
# Read index
|
| 402 |
+
with open(INPUT_DIR / "index.json") as f:
|
| 403 |
+
index = json.load(f)
|
| 404 |
+
all_chunk_metas = index["chunks"]
|
| 405 |
+
total_input_tokens = sum(c["dim"] for c in all_chunk_metas)
|
| 406 |
+
num_chunks = len(all_chunk_metas)
|
| 407 |
+
|
| 408 |
+
print(f"\n Input: {num_chunks} chunks, {total_input_tokens:,} tokens")
|
| 409 |
+
|
| 410 |
+
# Output config (same format)
|
| 411 |
+
config = {
|
| 412 |
+
"block_size": BLOCK_SIZE,
|
| 413 |
+
"vocab_size": tokenizer.get_vocab_size(),
|
| 414 |
+
}
|
| 415 |
+
|
| 416 |
+
# Stats
|
| 417 |
+
total_docs_in = 0
|
| 418 |
+
total_docs_kept = 0
|
| 419 |
+
total_docs_dropped = 0
|
| 420 |
+
total_boilerplate = 0
|
| 421 |
+
drop_reasons = Counter()
|
| 422 |
+
|
| 423 |
+
# Initialize streaming writer
|
| 424 |
+
writer = StreamingChunkWriter(str(OUTPUT_DIR), config)
|
| 425 |
+
|
| 426 |
+
# Carry-over: partial document tokens spanning chunk boundaries
|
| 427 |
+
carry_over = np.array([], dtype=DTYPE)
|
| 428 |
+
|
| 429 |
+
# ββ Process in batches of chunks ββββββββββββββββββββββββββββββββββββββ
|
| 430 |
+
num_batches = (num_chunks + CHUNKS_PER_BATCH - 1) // CHUNKS_PER_BATCH
|
| 431 |
+
print(f" Processing in {num_batches} batches of {CHUNKS_PER_BATCH} chunks...\n")
|
| 432 |
+
|
| 433 |
+
# Create multiprocessing pool for cleaning
|
| 434 |
+
pool = Pool(processes=NUM_WORKERS)
|
| 435 |
+
|
| 436 |
+
for batch_idx in range(num_batches):
|
| 437 |
+
batch_start = batch_idx * CHUNKS_PER_BATCH
|
| 438 |
+
batch_end = min(batch_start + CHUNKS_PER_BATCH, num_chunks)
|
| 439 |
+
batch_metas = all_chunk_metas[batch_start:batch_end]
|
| 440 |
+
|
| 441 |
+
t_batch = time.time()
|
| 442 |
+
print(f" ββ Batch {batch_idx+1}/{num_batches} (chunks {batch_start}-{batch_end-1}) ββ")
|
| 443 |
+
|
| 444 |
+
# 1. Read chunk tokens
|
| 445 |
+
batch_tokens_list = []
|
| 446 |
+
for cm in batch_metas:
|
| 447 |
+
batch_tokens_list.append(read_chunk_tokens(INPUT_DIR, cm))
|
| 448 |
+
batch_tokens = np.concatenate(batch_tokens_list)
|
| 449 |
+
del batch_tokens_list
|
| 450 |
+
|
| 451 |
+
# Prepend carry-over from previous batch
|
| 452 |
+
if len(carry_over) > 0:
|
| 453 |
+
batch_tokens = np.concatenate([carry_over, batch_tokens])
|
| 454 |
+
carry_over = np.array([], dtype=DTYPE)
|
| 455 |
+
|
| 456 |
+
batch_token_count = len(batch_tokens)
|
| 457 |
+
|
| 458 |
+
# 2. Split into documents
|
| 459 |
+
doc_token_lists, carry_over = split_documents_from_tokens(batch_tokens)
|
| 460 |
+
del batch_tokens
|
| 461 |
+
num_docs = len(doc_token_lists)
|
| 462 |
+
total_docs_in += num_docs
|
| 463 |
+
|
| 464 |
+
# 3. Decode documents to text (tokenizer.decode_batch is multithreaded Rust)
|
| 465 |
+
t_dec = time.time()
|
| 466 |
+
raw_texts = tokenizer.decode_batch(doc_token_lists, skip_special_tokens=False)
|
| 467 |
+
del doc_token_lists
|
| 468 |
+
dec_time = time.time() - t_dec
|
| 469 |
+
|
| 470 |
+
# 4. Clean + smart filter in parallel across all cores
|
| 471 |
+
t_clean = time.time()
|
| 472 |
+
results = pool.map(clean_and_filter, raw_texts, chunksize=512)
|
| 473 |
+
del raw_texts
|
| 474 |
+
clean_time = time.time() - t_clean
|
| 475 |
+
|
| 476 |
+
# Collect cleaned texts
|
| 477 |
+
cleaned_texts = []
|
| 478 |
+
batch_dropped = 0
|
| 479 |
+
batch_boilerplate = 0
|
| 480 |
+
for cleaned, reason, bp in results:
|
| 481 |
+
if cleaned is not None:
|
| 482 |
+
cleaned_texts.append(cleaned)
|
| 483 |
+
batch_boilerplate += bp
|
| 484 |
+
else:
|
| 485 |
+
batch_dropped += 1
|
| 486 |
+
drop_reasons[reason] += 1
|
| 487 |
+
del results
|
| 488 |
+
|
| 489 |
+
batch_kept = len(cleaned_texts)
|
| 490 |
+
total_docs_kept += batch_kept
|
| 491 |
+
total_docs_dropped += batch_dropped
|
| 492 |
+
total_boilerplate += batch_boilerplate
|
| 493 |
+
|
| 494 |
+
# 5. Re-tokenize + stream to writer (encode_batch is multithreaded Rust)
|
| 495 |
+
t_tok = time.time()
|
| 496 |
+
for i in range(0, len(cleaned_texts), ENCODE_BATCH):
|
| 497 |
+
sub = cleaned_texts[i:i+ENCODE_BATCH]
|
| 498 |
+
encoded = tokenizer.encode_batch(sub, add_special_tokens=False)
|
| 499 |
+
writer.add_tokens_batch(encoded)
|
| 500 |
+
del cleaned_texts
|
| 501 |
+
tok_time = time.time() - t_tok
|
| 502 |
+
|
| 503 |
+
elapsed = time.time() - t_batch
|
| 504 |
+
print(f" {batch_token_count:,} tokens β {num_docs:,} docs β kept {batch_kept:,} / dropped {batch_dropped:,}")
|
| 505 |
+
print(f" decode {dec_time:.1f}s | clean {clean_time:.1f}s | tokenize {tok_time:.1f}s | total {elapsed:.1f}s")
|
| 506 |
+
print(f" Running: {total_docs_kept:,} kept, {total_docs_dropped:,} dropped, {writer.total_tokens:,} tokens written")
|
| 507 |
+
|
| 508 |
+
pool.close()
|
| 509 |
+
pool.join()
|
| 510 |
+
|
| 511 |
+
# Handle carry-over (last partial doc if any)
|
| 512 |
+
if len(carry_over) > 0:
|
| 513 |
+
text = tokenizer.decode(carry_over.tolist(), skip_special_tokens=False)
|
| 514 |
+
result = deep_clean(text)
|
| 515 |
+
if result is not None:
|
| 516 |
+
keep, stripped, reason = smart_filter(result)
|
| 517 |
+
if keep:
|
| 518 |
+
encoded = tokenizer.encode(stripped, add_special_tokens=False)
|
| 519 |
+
writer.buffer.extend(encoded.ids)
|
| 520 |
+
writer.buffer.append(EOS_TOKEN_ID)
|
| 521 |
+
total_docs_kept += 1
|
| 522 |
+
total_boilerplate += len(result) - len(stripped)
|
| 523 |
+
else:
|
| 524 |
+
total_docs_dropped += 1
|
| 525 |
+
drop_reasons[reason] += 1
|
| 526 |
+
else:
|
| 527 |
+
total_docs_dropped += 1
|
| 528 |
+
drop_reasons["deep_clean_drop"] += 1
|
| 529 |
+
total_docs_in += 1
|
| 530 |
+
|
| 531 |
+
# Finalize output
|
| 532 |
+
final_tokens, discarded = writer.finalize()
|
| 533 |
+
|
| 534 |
+
total_time = time.time() - t_start
|
| 535 |
+
|
| 536 |
+
# ββ Report ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 537 |
+
print(f"\n{'='*75}")
|
| 538 |
+
print(f" RECLEAN REPORT - litdata_3b")
|
| 539 |
+
print(f"{'='*75}")
|
| 540 |
+
print(f"\n Total time: {total_time:.0f}s ({total_time/60:.1f} min)")
|
| 541 |
+
print(f" Workers: {NUM_WORKERS} CPU cores")
|
| 542 |
+
print(f"\n INPUT")
|
| 543 |
+
print(f" {'-'*60}")
|
| 544 |
+
print(f" Chunks: {num_chunks}")
|
| 545 |
+
print(f" Tokens: {total_input_tokens:,}")
|
| 546 |
+
print(f" Documents: {total_docs_in:,}")
|
| 547 |
+
print(f"\n CLEANING")
|
| 548 |
+
print(f" {'-'*60}")
|
| 549 |
+
print(f" Kept: {total_docs_kept:,}")
|
| 550 |
+
print(f" Dropped: {total_docs_dropped:,} ({total_docs_dropped/(max(total_docs_in,1))*100:.2f}%)")
|
| 551 |
+
if drop_reasons:
|
| 552 |
+
print(f" Drop reasons:")
|
| 553 |
+
for reason, count in sorted(drop_reasons.items(), key=lambda x: -x[1]):
|
| 554 |
+
print(f" {reason:<35} {count:>8,}")
|
| 555 |
+
print(f" Boilerplate stripped: {total_boilerplate:,} chars")
|
| 556 |
+
print(f"\n OUTPUT")
|
| 557 |
+
print(f" {'-'*60}")
|
| 558 |
+
print(f" Location: {OUTPUT_DIR}")
|
| 559 |
+
print(f" Chunks: {writer.chunk_idx}")
|
| 560 |
+
print(f" Tokens: {final_tokens:,}")
|
| 561 |
+
print(f" Discarded: {discarded} trailing tokens (< 1 block)")
|
| 562 |
+
print(f" Format: litdata binary (int32, BLOCK_SIZE=1025, EOS=0)")
|
| 563 |
+
diff = total_input_tokens - final_tokens
|
| 564 |
+
print(f"\n Token change: {total_input_tokens:,} β {final_tokens:,}")
|
| 565 |
+
print(f" Difference: {diff:,} ({diff/max(total_input_tokens,1)*100:.2f}%)")
|
| 566 |
+
print(f"\n{'='*75}")
|
| 567 |
+
|
| 568 |
+
# Save report
|
| 569 |
+
report_lines = [
|
| 570 |
+
f"RECLEAN REPORT - litdata_3b",
|
| 571 |
+
f"Time: {total_time:.0f}s ({total_time/60:.1f} min)",
|
| 572 |
+
f"Workers: {NUM_WORKERS}",
|
| 573 |
+
f"",
|
| 574 |
+
f"INPUT: {num_chunks} chunks, {total_input_tokens:,} tokens, {total_docs_in:,} docs",
|
| 575 |
+
f"OUTPUT: {writer.chunk_idx} chunks, {final_tokens:,} tokens",
|
| 576 |
+
f"",
|
| 577 |
+
f"Docs kept: {total_docs_kept:,}",
|
| 578 |
+
f"Docs dropped: {total_docs_dropped:,}",
|
| 579 |
+
]
|
| 580 |
+
if drop_reasons:
|
| 581 |
+
for reason, count in sorted(drop_reasons.items(), key=lambda x: -x[1]):
|
| 582 |
+
report_lines.append(f" {reason}: {count:,}")
|
| 583 |
+
report_lines.append(f"Boilerplate stripped: {total_boilerplate:,} chars")
|
| 584 |
+
report_lines.append(f"Token change: {total_input_tokens:,} -> {final_tokens:,} ({diff/max(total_input_tokens,1)*100:.2f}%)")
|
| 585 |
+
|
| 586 |
+
with open(OUTPUT_DIR / "CLEAN_REPORT.txt", "w", encoding="utf-8") as f:
|
| 587 |
+
f.write("\n".join(report_lines))
|
| 588 |
+
print(f" Report saved to: {OUTPUT_DIR / 'CLEAN_REPORT.txt'}")
|
| 589 |
+
print(f" Done!")
|
| 590 |
+
|
| 591 |
+
|
| 592 |
+
if __name__ == "__main__":
|
| 593 |
+
main()
|
Base/scripts/reclean_english.py
ADDED
|
@@ -0,0 +1,610 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""
|
| 3 |
+
Reclean litdata_english ONLY (separate from the main reclean running on 3b).
|
| 4 |
+
|
| 5 |
+
Decodes all docs, applies English cleaning, re-tokenizes, rebuilds chunks.
|
| 6 |
+
Includes detailed BEFORE vs AFTER comparison report with sample diffs.
|
| 7 |
+
"""
|
| 8 |
+
|
| 9 |
+
import json
|
| 10 |
+
import os
|
| 11 |
+
import re
|
| 12 |
+
import time
|
| 13 |
+
import unicodedata
|
| 14 |
+
from pathlib import Path
|
| 15 |
+
from collections import Counter
|
| 16 |
+
|
| 17 |
+
import numpy as np
|
| 18 |
+
from tokenizers import Tokenizer
|
| 19 |
+
|
| 20 |
+
ROOT = Path(__file__).resolve().parent.parent.parent
|
| 21 |
+
BLOCK_SIZE = 1025
|
| 22 |
+
DTYPE = np.int32
|
| 23 |
+
CHUNK_BYTES_TARGET = 64 * 1024 * 1024
|
| 24 |
+
EOS_TOKEN_ID = 0
|
| 25 |
+
|
| 26 |
+
# -- Load tokenizer -----------------------------------------------------------
|
| 27 |
+
print("Loading tokenizer...")
|
| 28 |
+
tokenizer = Tokenizer.from_file(
|
| 29 |
+
str(ROOT / "Base" / "checkpoints" / "EleutherAI" / "pythia-160m" / "tokenizer.json")
|
| 30 |
+
)
|
| 31 |
+
|
| 32 |
+
# ==============================================================================
|
| 33 |
+
# TEXT CLEANING (same pipeline as reclean_litdata.py)
|
| 34 |
+
# ==============================================================================
|
| 35 |
+
|
| 36 |
+
CONTROL_CHARS = [
|
| 37 |
+
"\x00", "\x01", "\x02", "\x03", "\x04", "\x05", "\x06", "\x07",
|
| 38 |
+
"\x08", "\x0b", "\x0c", "\x0e", "\x0f", "\x10", "\x11", "\x12",
|
| 39 |
+
"\x13", "\x14", "\x15", "\x16", "\x17", "\x18", "\x19", "\x1a",
|
| 40 |
+
"\x1b", "\x1c", "\x1d", "\x1e", "\x1f", "\x7f",
|
| 41 |
+
"\ufeff", "\ufffd",
|
| 42 |
+
]
|
| 43 |
+
|
| 44 |
+
HTML_ENTITIES = [
|
| 45 |
+
("&", "&"), ("<", "<"), (">", ">"),
|
| 46 |
+
(""", '"'), ("'", "'"), ("'", "'"),
|
| 47 |
+
(" ", " "), ("—", " - "), ("–", "-"),
|
| 48 |
+
("…", "..."), ("«", '"'), ("»", '"'),
|
| 49 |
+
("•", "- "), ("·", " "), ("©", "(c)"),
|
| 50 |
+
("®", "(R)"), ("™", "(TM)"), ("°", " degrees"),
|
| 51 |
+
]
|
| 52 |
+
|
| 53 |
+
RE_URL = re.compile(r'https?://\S+|www\.\S+', re.I)
|
| 54 |
+
RE_EMAIL = re.compile(r'\b[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}\b')
|
| 55 |
+
RE_FILE_PATH = re.compile(r'(?:[A-Z]:\\|/(?:home|usr|var|etc|opt)/)\S+')
|
| 56 |
+
RE_HTML_TAG = re.compile(r'</?[a-zA-Z][a-zA-Z0-9]*(?:\s[^>]*)?\s*/?>')
|
| 57 |
+
RE_HTML_COMMENT = re.compile(r'<!--.*?-->', re.DOTALL)
|
| 58 |
+
RE_CODE_BLOCK = re.compile(r'```[\s\S]*?```')
|
| 59 |
+
RE_IMPORT = re.compile(r'^(?:import |from \S+ import |#include |using namespace |require\()', re.M)
|
| 60 |
+
RE_REPEATED_LINE = re.compile(r'^(.{20,})\n(?:\1\n?)+', re.M)
|
| 61 |
+
RE_REPEATED_PUNCT = re.compile(r'([!?.])\1{3,}')
|
| 62 |
+
RE_REPEATED_CHAR = re.compile(r'(.)\1{5,}')
|
| 63 |
+
RE_REPEATED_WORD = re.compile(r'\b(\w+)(?:\s+\1){2,}\b', re.I)
|
| 64 |
+
RE_MULTI_NEWLINE = re.compile(r'\n{4,}')
|
| 65 |
+
RE_MULTI_SPACE = re.compile(r'[ \t]{2,}')
|
| 66 |
+
RE_TRAILING_SPACE = re.compile(r'[ \t]+$', re.M)
|
| 67 |
+
RE_NO_SPACE_AFTER_PERIOD = re.compile(r'([.!?])([A-Z])')
|
| 68 |
+
RE_DOUBLE_PERIOD = re.compile(r'\.{2}(?!\.)')
|
| 69 |
+
RE_SPACE_BEFORE_PUNCT = re.compile(r'\s+([.,;:!?])')
|
| 70 |
+
|
| 71 |
+
|
| 72 |
+
# -- Stats trackers -----------------------------------------------------------
|
| 73 |
+
class CleaningStats:
|
| 74 |
+
def __init__(self):
|
| 75 |
+
self.control_chars_removed = 0
|
| 76 |
+
self.html_entities_fixed = 0
|
| 77 |
+
self.html_tags_removed = 0
|
| 78 |
+
self.urls_removed = 0
|
| 79 |
+
self.emails_removed = 0
|
| 80 |
+
self.code_blocks_removed = 0
|
| 81 |
+
self.repeated_lines_fixed = 0
|
| 82 |
+
self.repeated_punct_fixed = 0
|
| 83 |
+
self.repeated_words_fixed = 0
|
| 84 |
+
self.punctuation_fixed = 0
|
| 85 |
+
self.smart_quotes_fixed = 0
|
| 86 |
+
self.junk_lines_removed = 0
|
| 87 |
+
self.lines_capitalized = 0
|
| 88 |
+
self.duplicate_paragraphs_removed = 0
|
| 89 |
+
self.docs_dropped_short = 0
|
| 90 |
+
self.docs_dropped_nonenglish = 0
|
| 91 |
+
self.total_chars_before = 0
|
| 92 |
+
self.total_chars_after = 0
|
| 93 |
+
|
| 94 |
+
stats = CleaningStats()
|
| 95 |
+
|
| 96 |
+
|
| 97 |
+
def clean_text_tracked(text):
|
| 98 |
+
"""Clean text with detailed change tracking for the report."""
|
| 99 |
+
global stats
|
| 100 |
+
if not text or len(text.strip()) < 30:
|
| 101 |
+
stats.docs_dropped_short += 1
|
| 102 |
+
return None
|
| 103 |
+
|
| 104 |
+
original = text
|
| 105 |
+
stats.total_chars_before += len(text)
|
| 106 |
+
|
| 107 |
+
# 1. Unicode normalization
|
| 108 |
+
text = unicodedata.normalize("NFKC", text)
|
| 109 |
+
|
| 110 |
+
# 2. Control chars
|
| 111 |
+
for ch in CONTROL_CHARS:
|
| 112 |
+
count = text.count(ch)
|
| 113 |
+
if count:
|
| 114 |
+
stats.control_chars_removed += count
|
| 115 |
+
text = text.replace(ch, "")
|
| 116 |
+
|
| 117 |
+
# 3. HTML entities
|
| 118 |
+
for old, new in HTML_ENTITIES:
|
| 119 |
+
count = text.count(old)
|
| 120 |
+
if count:
|
| 121 |
+
stats.html_entities_fixed += count
|
| 122 |
+
text = text.replace(old, new)
|
| 123 |
+
|
| 124 |
+
# 4. HTML tags/comments
|
| 125 |
+
m = RE_HTML_COMMENT.findall(text)
|
| 126 |
+
stats.html_tags_removed += len(m)
|
| 127 |
+
text = RE_HTML_COMMENT.sub("", text)
|
| 128 |
+
m = RE_HTML_TAG.findall(text)
|
| 129 |
+
stats.html_tags_removed += len(m)
|
| 130 |
+
text = RE_HTML_TAG.sub("", text)
|
| 131 |
+
|
| 132 |
+
# 5. URLs, emails, file paths
|
| 133 |
+
stats.urls_removed += len(RE_URL.findall(text))
|
| 134 |
+
text = RE_URL.sub("", text)
|
| 135 |
+
stats.emails_removed += len(RE_EMAIL.findall(text))
|
| 136 |
+
text = RE_EMAIL.sub("", text)
|
| 137 |
+
text = RE_FILE_PATH.sub("", text)
|
| 138 |
+
|
| 139 |
+
# 6. Code blocks
|
| 140 |
+
stats.code_blocks_removed += len(RE_CODE_BLOCK.findall(text))
|
| 141 |
+
text = RE_CODE_BLOCK.sub("", text)
|
| 142 |
+
|
| 143 |
+
# 7. Repeated content
|
| 144 |
+
stats.repeated_lines_fixed += len(RE_REPEATED_LINE.findall(text))
|
| 145 |
+
text = RE_REPEATED_LINE.sub(r'\1', text)
|
| 146 |
+
stats.repeated_punct_fixed += len(RE_REPEATED_PUNCT.findall(text))
|
| 147 |
+
text = RE_REPEATED_PUNCT.sub(r'\1\1\1', text)
|
| 148 |
+
text = RE_REPEATED_CHAR.sub(r'\1\1\1', text)
|
| 149 |
+
stats.repeated_words_fixed += len(RE_REPEATED_WORD.findall(text))
|
| 150 |
+
text = RE_REPEATED_WORD.sub(r'\1', text)
|
| 151 |
+
|
| 152 |
+
# 8. Whitespace
|
| 153 |
+
text = text.replace('\t', ' ')
|
| 154 |
+
text = RE_TRAILING_SPACE.sub('', text)
|
| 155 |
+
text = RE_MULTI_SPACE.sub(' ', text)
|
| 156 |
+
text = RE_MULTI_NEWLINE.sub('\n\n\n', text)
|
| 157 |
+
|
| 158 |
+
# 9. Punctuation
|
| 159 |
+
p1 = len(RE_DOUBLE_PERIOD.findall(text))
|
| 160 |
+
p2 = len(RE_NO_SPACE_AFTER_PERIOD.findall(text))
|
| 161 |
+
p3 = len(RE_SPACE_BEFORE_PUNCT.findall(text))
|
| 162 |
+
stats.punctuation_fixed += p1 + p2 + p3
|
| 163 |
+
text = RE_DOUBLE_PERIOD.sub('.', text)
|
| 164 |
+
text = RE_NO_SPACE_AFTER_PERIOD.sub(r'\1 \2', text)
|
| 165 |
+
text = RE_SPACE_BEFORE_PUNCT.sub(r'\1', text)
|
| 166 |
+
|
| 167 |
+
# 10. Smart quotes -> ASCII
|
| 168 |
+
sq = 0
|
| 169 |
+
for ch in ['\u2018', '\u2019', '\u201c', '\u201d', '\u2013', '\u2014', '\u2026', '\u2022', '\u00b7', '\u00a0']:
|
| 170 |
+
c = text.count(ch)
|
| 171 |
+
if c:
|
| 172 |
+
sq += c
|
| 173 |
+
stats.smart_quotes_fixed += sq
|
| 174 |
+
text = text.replace('\u2018', "'").replace('\u2019', "'")
|
| 175 |
+
text = text.replace('\u201c', '"').replace('\u201d', '"')
|
| 176 |
+
text = text.replace('\u2013', '-').replace('\u2014', ' - ')
|
| 177 |
+
text = text.replace('\u2026', '...')
|
| 178 |
+
text = text.replace('\u2022', '- ')
|
| 179 |
+
text = text.replace('\u00b7', ' ')
|
| 180 |
+
text = text.replace('\u00a0', ' ')
|
| 181 |
+
|
| 182 |
+
# 11. Line-by-line cleaning
|
| 183 |
+
lines = text.split('\n')
|
| 184 |
+
clean_lines = []
|
| 185 |
+
for line in lines:
|
| 186 |
+
line = line.strip()
|
| 187 |
+
if not line:
|
| 188 |
+
clean_lines.append('')
|
| 189 |
+
continue
|
| 190 |
+
if len(line) > 10:
|
| 191 |
+
alpha_count = sum(1 for c in line if c.isalpha())
|
| 192 |
+
if alpha_count / len(line) < 0.40:
|
| 193 |
+
stats.junk_lines_removed += 1
|
| 194 |
+
continue
|
| 195 |
+
if line.count('|') > 3 or line.count('{') > 2 or line.count('}') > 2:
|
| 196 |
+
stats.junk_lines_removed += 1
|
| 197 |
+
continue
|
| 198 |
+
if RE_IMPORT.match(line):
|
| 199 |
+
stats.junk_lines_removed += 1
|
| 200 |
+
continue
|
| 201 |
+
if line and line[0].isalpha() and line[0].islower():
|
| 202 |
+
if not clean_lines or clean_lines[-1] == '' or clean_lines[-1].rstrip().endswith(('.', '!', '?', ':')):
|
| 203 |
+
line = line[0].upper() + line[1:]
|
| 204 |
+
stats.lines_capitalized += 1
|
| 205 |
+
clean_lines.append(line)
|
| 206 |
+
|
| 207 |
+
text = '\n'.join(clean_lines)
|
| 208 |
+
text = text.strip()
|
| 209 |
+
|
| 210 |
+
# 13. Deduplicate paragraphs
|
| 211 |
+
paragraphs = text.split('\n\n')
|
| 212 |
+
seen = set()
|
| 213 |
+
unique_paragraphs = []
|
| 214 |
+
for p in paragraphs:
|
| 215 |
+
p_stripped = p.strip()
|
| 216 |
+
if not p_stripped:
|
| 217 |
+
continue
|
| 218 |
+
p_key = ' '.join(p_stripped.lower().split())
|
| 219 |
+
if p_key not in seen:
|
| 220 |
+
seen.add(p_key)
|
| 221 |
+
unique_paragraphs.append(p_stripped)
|
| 222 |
+
else:
|
| 223 |
+
stats.duplicate_paragraphs_removed += 1
|
| 224 |
+
text = '\n\n'.join(unique_paragraphs)
|
| 225 |
+
|
| 226 |
+
# 14. Final quality gate
|
| 227 |
+
text = text.strip()
|
| 228 |
+
if len(text) < 50:
|
| 229 |
+
stats.docs_dropped_short += 1
|
| 230 |
+
return None
|
| 231 |
+
if len(text.split()) < 10:
|
| 232 |
+
stats.docs_dropped_short += 1
|
| 233 |
+
return None
|
| 234 |
+
ascii_count = sum(1 for c in text if ord(c) < 128)
|
| 235 |
+
if ascii_count / max(len(text), 1) < 0.85:
|
| 236 |
+
stats.docs_dropped_nonenglish += 1
|
| 237 |
+
return None
|
| 238 |
+
|
| 239 |
+
stats.total_chars_after += len(text)
|
| 240 |
+
return text
|
| 241 |
+
|
| 242 |
+
|
| 243 |
+
# ==============================================================================
|
| 244 |
+
# LITDATA I/O (same as reclean_litdata.py)
|
| 245 |
+
# ==============================================================================
|
| 246 |
+
|
| 247 |
+
def read_all_tokens(litdata_dir):
|
| 248 |
+
with open(litdata_dir / "index.json") as f:
|
| 249 |
+
index = json.load(f)
|
| 250 |
+
chunks = index["chunks"]
|
| 251 |
+
total_tokens = sum(c["dim"] for c in chunks)
|
| 252 |
+
print(f" Reading {len(chunks)} chunks ({total_tokens:,} tokens)...")
|
| 253 |
+
all_tokens = np.empty(total_tokens, dtype=DTYPE)
|
| 254 |
+
pos = 0
|
| 255 |
+
for i, chunk in enumerate(chunks):
|
| 256 |
+
chunk_path = litdata_dir / chunk["filename"]
|
| 257 |
+
n_blocks = chunk["chunk_size"]
|
| 258 |
+
header_ints = 1 + n_blocks + 1
|
| 259 |
+
header_bytes = header_ints * 4
|
| 260 |
+
with open(chunk_path, "rb") as f:
|
| 261 |
+
f.seek(header_bytes)
|
| 262 |
+
data = np.fromfile(f, dtype=DTYPE, count=chunk["dim"])
|
| 263 |
+
all_tokens[pos:pos + len(data)] = data
|
| 264 |
+
pos += len(data)
|
| 265 |
+
if (i + 1) == len(chunks):
|
| 266 |
+
print(f" Read {i+1}/{len(chunks)} chunks ({pos:,} tokens)")
|
| 267 |
+
return all_tokens[:pos]
|
| 268 |
+
|
| 269 |
+
|
| 270 |
+
def split_documents(token_stream):
|
| 271 |
+
eos_positions = np.where(token_stream == EOS_TOKEN_ID)[0]
|
| 272 |
+
docs = []
|
| 273 |
+
start = 0
|
| 274 |
+
for eos_pos in eos_positions:
|
| 275 |
+
if eos_pos > start:
|
| 276 |
+
docs.append(token_stream[start:eos_pos])
|
| 277 |
+
start = eos_pos + 1
|
| 278 |
+
if start < len(token_stream):
|
| 279 |
+
docs.append(token_stream[start:])
|
| 280 |
+
return docs
|
| 281 |
+
|
| 282 |
+
|
| 283 |
+
def write_litdata_chunks(output_dir, token_stream, config):
|
| 284 |
+
os.makedirs(output_dir, exist_ok=True)
|
| 285 |
+
dtype_size = DTYPE().itemsize
|
| 286 |
+
tokens_per_chunk = CHUNK_BYTES_TARGET // dtype_size
|
| 287 |
+
tokens_per_chunk = (tokens_per_chunk // BLOCK_SIZE) * BLOCK_SIZE
|
| 288 |
+
chunks_metadata = []
|
| 289 |
+
pos = 0
|
| 290 |
+
chunk_idx = 0
|
| 291 |
+
while pos < len(token_stream):
|
| 292 |
+
remaining = len(token_stream) - pos
|
| 293 |
+
chunk_tokens = min(tokens_per_chunk, remaining)
|
| 294 |
+
num_blocks = chunk_tokens // BLOCK_SIZE
|
| 295 |
+
if num_blocks == 0:
|
| 296 |
+
break
|
| 297 |
+
actual_tokens = num_blocks * BLOCK_SIZE
|
| 298 |
+
chunk_data = token_stream[pos:pos + actual_tokens]
|
| 299 |
+
filename = f"chunk-0-{chunk_idx}.bin"
|
| 300 |
+
filepath = os.path.join(output_dir, filename)
|
| 301 |
+
header_num_items = np.array([num_blocks], dtype=np.uint32)
|
| 302 |
+
offsets = np.arange(num_blocks + 1, dtype=np.uint32) * (BLOCK_SIZE * dtype_size)
|
| 303 |
+
header = np.concatenate([header_num_items, offsets])
|
| 304 |
+
with open(filepath, "wb") as f:
|
| 305 |
+
header.tofile(f)
|
| 306 |
+
chunk_data.tofile(f)
|
| 307 |
+
meta = {
|
| 308 |
+
"chunk_bytes": int(header.nbytes + chunk_data.nbytes),
|
| 309 |
+
"chunk_size": num_blocks,
|
| 310 |
+
"dim": int(actual_tokens),
|
| 311 |
+
"filename": filename,
|
| 312 |
+
}
|
| 313 |
+
chunks_metadata.append(meta)
|
| 314 |
+
pos += actual_tokens
|
| 315 |
+
chunk_idx += 1
|
| 316 |
+
print(f" Written chunk {chunk_idx} ({pos:,}/{len(token_stream):,} tokens)")
|
| 317 |
+
index = {"chunks": chunks_metadata, "config": config, "updated_at": str(time.time())}
|
| 318 |
+
with open(os.path.join(output_dir, "index.json"), "w") as f:
|
| 319 |
+
json.dump(index, f, indent=2)
|
| 320 |
+
return chunks_metadata
|
| 321 |
+
|
| 322 |
+
|
| 323 |
+
# ==============================================================================
|
| 324 |
+
# QUALITY METRICS (before/after comparison)
|
| 325 |
+
# ==============================================================================
|
| 326 |
+
|
| 327 |
+
def compute_quality_metrics(texts, label):
|
| 328 |
+
"""Compute quality metrics on a list of texts."""
|
| 329 |
+
if not texts:
|
| 330 |
+
return {}
|
| 331 |
+
ascii_ratios = []
|
| 332 |
+
alpha_ratios = []
|
| 333 |
+
word_counts = []
|
| 334 |
+
sentence_counts = []
|
| 335 |
+
unique_word_ratios = []
|
| 336 |
+
url_counts = 0
|
| 337 |
+
html_counts = 0
|
| 338 |
+
|
| 339 |
+
for t in texts:
|
| 340 |
+
chars = len(t)
|
| 341 |
+
if chars == 0:
|
| 342 |
+
continue
|
| 343 |
+
ascii_ratios.append(sum(1 for c in t if ord(c) < 128) / chars)
|
| 344 |
+
alpha_ratios.append(sum(1 for c in t if c.isalpha()) / chars)
|
| 345 |
+
words = t.split()
|
| 346 |
+
word_counts.append(len(words))
|
| 347 |
+
sentence_counts.append(len(re.findall(r'[.!?]+\s', t)))
|
| 348 |
+
if words:
|
| 349 |
+
unique_word_ratios.append(len(set(w.lower() for w in words)) / len(words))
|
| 350 |
+
url_counts += len(re.findall(r'https?://\S+', t))
|
| 351 |
+
if re.search(r'<(html|body|div|span|script|style)\b', t, re.I):
|
| 352 |
+
html_counts += 1
|
| 353 |
+
|
| 354 |
+
n = len(texts)
|
| 355 |
+
return {
|
| 356 |
+
"label": label,
|
| 357 |
+
"documents": n,
|
| 358 |
+
"avg_ascii_ratio": sum(ascii_ratios) / max(len(ascii_ratios), 1),
|
| 359 |
+
"avg_alpha_ratio": sum(alpha_ratios) / max(len(alpha_ratios), 1),
|
| 360 |
+
"avg_word_count": sum(word_counts) / max(len(word_counts), 1),
|
| 361 |
+
"avg_sentences": sum(sentence_counts) / max(len(sentence_counts), 1),
|
| 362 |
+
"avg_unique_word_ratio": sum(unique_word_ratios) / max(len(unique_word_ratios), 1),
|
| 363 |
+
"total_urls": url_counts,
|
| 364 |
+
"docs_with_html": html_counts,
|
| 365 |
+
"min_words": min(word_counts) if word_counts else 0,
|
| 366 |
+
"max_words": max(word_counts) if word_counts else 0,
|
| 367 |
+
}
|
| 368 |
+
|
| 369 |
+
|
| 370 |
+
def print_comparison(before_metrics, after_metrics):
|
| 371 |
+
"""Print a side-by-side comparison table."""
|
| 372 |
+
print(f"\n {'METRIC':<30} {'BEFORE':>15} {'AFTER':>15} {'CHANGE':>15}")
|
| 373 |
+
print(f" {'='*75}")
|
| 374 |
+
|
| 375 |
+
rows = [
|
| 376 |
+
("Documents", "documents", "d"),
|
| 377 |
+
("Avg ASCII ratio", "avg_ascii_ratio", "f4"),
|
| 378 |
+
("Avg alpha ratio", "avg_alpha_ratio", "f4"),
|
| 379 |
+
("Avg word count", "avg_word_count", "f1"),
|
| 380 |
+
("Avg sentences/doc", "avg_sentences", "f1"),
|
| 381 |
+
("Avg unique word ratio", "avg_unique_word_ratio", "f4"),
|
| 382 |
+
("URLs found", "total_urls", "d"),
|
| 383 |
+
("Docs with HTML", "docs_with_html", "d"),
|
| 384 |
+
("Min word count", "min_words", "d"),
|
| 385 |
+
("Max word count", "max_words", "d"),
|
| 386 |
+
]
|
| 387 |
+
|
| 388 |
+
for label, key, fmt in rows:
|
| 389 |
+
bv = before_metrics.get(key, 0)
|
| 390 |
+
av = after_metrics.get(key, 0)
|
| 391 |
+
if fmt == "d":
|
| 392 |
+
diff = av - bv
|
| 393 |
+
sign = "+" if diff > 0 else ""
|
| 394 |
+
print(f" {label:<30} {bv:>15,} {av:>15,} {sign}{diff:>14,}")
|
| 395 |
+
else:
|
| 396 |
+
diff = av - bv
|
| 397 |
+
sign = "+" if diff > 0 else ""
|
| 398 |
+
print(f" {label:<30} {bv:>15.4f} {av:>15.4f} {sign}{diff:>14.4f}")
|
| 399 |
+
|
| 400 |
+
|
| 401 |
+
# ==============================================================================
|
| 402 |
+
# MAIN
|
| 403 |
+
# ==============================================================================
|
| 404 |
+
|
| 405 |
+
if __name__ == "__main__":
|
| 406 |
+
t_start = time.time()
|
| 407 |
+
data_dir = ROOT / "Base" / "data"
|
| 408 |
+
input_dir = data_dir / "litdata_english"
|
| 409 |
+
output_dir = data_dir / "litdata_english_clean"
|
| 410 |
+
|
| 411 |
+
print(f"\n{'='*75}")
|
| 412 |
+
print(f" RECLEANING: litdata_english (English Knowledge)")
|
| 413 |
+
print(f" Input: {input_dir}")
|
| 414 |
+
print(f" Output: {output_dir}")
|
| 415 |
+
print(f"{'='*75}")
|
| 416 |
+
|
| 417 |
+
# 1. Read all tokens
|
| 418 |
+
t0 = time.time()
|
| 419 |
+
token_stream = read_all_tokens(input_dir)
|
| 420 |
+
print(f" Read {len(token_stream):,} tokens in {time.time()-t0:.1f}s")
|
| 421 |
+
|
| 422 |
+
# 2. Split into documents
|
| 423 |
+
t1 = time.time()
|
| 424 |
+
doc_tokens = split_documents(token_stream)
|
| 425 |
+
print(f" Found {len(doc_tokens):,} documents in {time.time()-t1:.1f}s")
|
| 426 |
+
del token_stream
|
| 427 |
+
|
| 428 |
+
# 3. Decode
|
| 429 |
+
t2 = time.time()
|
| 430 |
+
print(f" Decoding {len(doc_tokens):,} documents...")
|
| 431 |
+
raw_texts = []
|
| 432 |
+
for i, doc in enumerate(doc_tokens):
|
| 433 |
+
text = tokenizer.decode(doc.tolist(), skip_special_tokens=False)
|
| 434 |
+
raw_texts.append(text)
|
| 435 |
+
if (i + 1) % 20000 == 0 or i == len(doc_tokens) - 1:
|
| 436 |
+
print(f" Decoded {i+1:,}/{len(doc_tokens):,}")
|
| 437 |
+
del doc_tokens
|
| 438 |
+
print(f" Decoded in {time.time()-t2:.1f}s")
|
| 439 |
+
|
| 440 |
+
# 4. Compute BEFORE metrics (ALL documents)
|
| 441 |
+
print(f"\n Computing BEFORE quality metrics on ALL {len(raw_texts):,} documents...")
|
| 442 |
+
import random
|
| 443 |
+
random.seed(42)
|
| 444 |
+
sample_size = len(raw_texts)
|
| 445 |
+
before_metrics = compute_quality_metrics(raw_texts, "BEFORE")
|
| 446 |
+
|
| 447 |
+
# 5. Clean all documents (with tracking)
|
| 448 |
+
t3 = time.time()
|
| 449 |
+
print(f"\n Cleaning {len(raw_texts):,} documents...")
|
| 450 |
+
cleaned_texts = []
|
| 451 |
+
dropped_texts = [] # keep some dropped for the report
|
| 452 |
+
before_after_pairs = [] # keep pairs for sample diffs
|
| 453 |
+
dropped = 0
|
| 454 |
+
|
| 455 |
+
for i, text in enumerate(raw_texts):
|
| 456 |
+
result = clean_text_tracked(text)
|
| 457 |
+
if result is not None:
|
| 458 |
+
cleaned_texts.append(result)
|
| 459 |
+
# Keep some before/after pairs for diff display
|
| 460 |
+
if len(before_after_pairs) < 20 and text != result:
|
| 461 |
+
before_after_pairs.append((text, result))
|
| 462 |
+
else:
|
| 463 |
+
dropped += 1
|
| 464 |
+
if len(dropped_texts) < 10:
|
| 465 |
+
dropped_texts.append(text)
|
| 466 |
+
if (i + 1) % 20000 == 0 or i == len(raw_texts) - 1:
|
| 467 |
+
print(f" Processed {i+1:,}/{len(raw_texts):,} | kept={len(cleaned_texts):,} | dropped={dropped:,}")
|
| 468 |
+
|
| 469 |
+
del raw_texts
|
| 470 |
+
print(f" Cleaning done in {time.time()-t3:.1f}s")
|
| 471 |
+
print(f" Kept {len(cleaned_texts):,} | Dropped {dropped:,} ({dropped/max(dropped+len(cleaned_texts),1)*100:.1f}%)")
|
| 472 |
+
|
| 473 |
+
# 6. Compute AFTER metrics (ALL documents)
|
| 474 |
+
print(f"\n Computing AFTER quality metrics on ALL {len(cleaned_texts):,} documents...")
|
| 475 |
+
after_metrics = compute_quality_metrics(cleaned_texts, "AFTER")
|
| 476 |
+
|
| 477 |
+
# 7. Re-tokenize
|
| 478 |
+
t4 = time.time()
|
| 479 |
+
print(f"\n Re-tokenizing {len(cleaned_texts):,} documents...")
|
| 480 |
+
all_token_ids = []
|
| 481 |
+
total_new_tokens = 0
|
| 482 |
+
ENCODE_BATCH = 10000
|
| 483 |
+
for i in range(0, len(cleaned_texts), ENCODE_BATCH):
|
| 484 |
+
batch = cleaned_texts[i:i+ENCODE_BATCH]
|
| 485 |
+
encoded = tokenizer.encode_batch(batch, add_special_tokens=False)
|
| 486 |
+
for enc in encoded:
|
| 487 |
+
ids = enc.ids
|
| 488 |
+
all_token_ids.extend(ids)
|
| 489 |
+
all_token_ids.append(EOS_TOKEN_ID)
|
| 490 |
+
total_new_tokens += len(ids) + 1
|
| 491 |
+
done = min(i + ENCODE_BATCH, len(cleaned_texts))
|
| 492 |
+
if done == len(cleaned_texts):
|
| 493 |
+
print(f" Tokenized {done:,}/{len(cleaned_texts):,} ({total_new_tokens:,} tokens)")
|
| 494 |
+
del cleaned_texts
|
| 495 |
+
print(f" Tokenized in {time.time()-t4:.1f}s")
|
| 496 |
+
|
| 497 |
+
# 8. Write chunks
|
| 498 |
+
t5 = time.time()
|
| 499 |
+
print(f"\n Building & writing litdata chunks...")
|
| 500 |
+
new_stream = np.array(all_token_ids, dtype=DTYPE)
|
| 501 |
+
del all_token_ids
|
| 502 |
+
with open(input_dir / "index.json") as f:
|
| 503 |
+
config = json.load(f)["config"]
|
| 504 |
+
chunks = write_litdata_chunks(str(output_dir), new_stream, config)
|
| 505 |
+
total_in = sum(c["dim"] for c in json.load(open(input_dir / "index.json"))["chunks"])
|
| 506 |
+
total_out = sum(c["dim"] for c in chunks)
|
| 507 |
+
print(f" Written {len(chunks)} chunks in {time.time()-t5:.1f}s")
|
| 508 |
+
|
| 509 |
+
# ================================================================
|
| 510 |
+
# DETAILED COMPARISON REPORT
|
| 511 |
+
# ================================================================
|
| 512 |
+
elapsed = time.time() - t_start
|
| 513 |
+
|
| 514 |
+
report = []
|
| 515 |
+
report.append("")
|
| 516 |
+
report.append("=" * 75)
|
| 517 |
+
report.append(" LITDATA_ENGLISH RECLEANING - FULL COMPARISON REPORT")
|
| 518 |
+
report.append("=" * 75)
|
| 519 |
+
|
| 520 |
+
report.append(f"\n Processing time: {elapsed:.0f}s ({elapsed/60:.1f} min)")
|
| 521 |
+
|
| 522 |
+
report.append(f"\n TOKEN SUMMARY")
|
| 523 |
+
report.append(f" {'-'*60}")
|
| 524 |
+
report.append(f" Input tokens: {total_in:>15,}")
|
| 525 |
+
report.append(f" Output tokens: {total_out:>15,}")
|
| 526 |
+
report.append(f" Removed: {total_in - total_out:>15,} ({(total_in-total_out)/total_in*100:.2f}%)")
|
| 527 |
+
|
| 528 |
+
report.append(f"\n CLEANING OPERATIONS PERFORMED")
|
| 529 |
+
report.append(f" {'-'*60}")
|
| 530 |
+
report.append(f" Control chars removed: {stats.control_chars_removed:>10,}")
|
| 531 |
+
report.append(f" HTML entities fixed: {stats.html_entities_fixed:>10,}")
|
| 532 |
+
report.append(f" HTML tags removed: {stats.html_tags_removed:>10,}")
|
| 533 |
+
report.append(f" URLs removed: {stats.urls_removed:>10,}")
|
| 534 |
+
report.append(f" Emails removed: {stats.emails_removed:>10,}")
|
| 535 |
+
report.append(f" Code blocks removed: {stats.code_blocks_removed:>10,}")
|
| 536 |
+
report.append(f" Repeated lines fixed: {stats.repeated_lines_fixed:>10,}")
|
| 537 |
+
report.append(f" Repeated punctuation fixed: {stats.repeated_punct_fixed:>10,}")
|
| 538 |
+
report.append(f" Repeated words fixed: {stats.repeated_words_fixed:>10,}")
|
| 539 |
+
report.append(f" Punctuation fixes: {stats.punctuation_fixed:>10,}")
|
| 540 |
+
report.append(f" Smart quotes normalized: {stats.smart_quotes_fixed:>10,}")
|
| 541 |
+
report.append(f" Junk lines removed: {stats.junk_lines_removed:>10,}")
|
| 542 |
+
report.append(f" Lines capitalized: {stats.lines_capitalized:>10,}")
|
| 543 |
+
report.append(f" Duplicate paragraphs removed: {stats.duplicate_paragraphs_removed:>10,}")
|
| 544 |
+
report.append(f" Docs dropped (too short): {stats.docs_dropped_short:>10,}")
|
| 545 |
+
report.append(f" Docs dropped (non-English): {stats.docs_dropped_nonenglish:>10,}")
|
| 546 |
+
report.append(f" Total chars before: {stats.total_chars_before:>15,}")
|
| 547 |
+
report.append(f" Total chars after: {stats.total_chars_after:>15,}")
|
| 548 |
+
char_diff = stats.total_chars_before - stats.total_chars_after
|
| 549 |
+
report.append(f" Chars cleaned out: {char_diff:>15,} ({char_diff/max(stats.total_chars_before,1)*100:.2f}%)")
|
| 550 |
+
|
| 551 |
+
report.append(f"\n QUALITY METRICS COMPARISON (ALL {sample_size:,} docs)")
|
| 552 |
+
report.append(f" {'-'*60}")
|
| 553 |
+
# Print table header
|
| 554 |
+
report.append(f" {'METRIC':<30} {'BEFORE':>12} {'AFTER':>12} {'CHANGE':>12}")
|
| 555 |
+
report.append(f" {'='*66}")
|
| 556 |
+
metrics_rows = [
|
| 557 |
+
("Documents", "documents", "d"),
|
| 558 |
+
("Avg ASCII ratio", "avg_ascii_ratio", "f"),
|
| 559 |
+
("Avg alpha ratio", "avg_alpha_ratio", "f"),
|
| 560 |
+
("Avg word count", "avg_word_count", "f1"),
|
| 561 |
+
("Avg sentences/doc", "avg_sentences", "f1"),
|
| 562 |
+
("Avg unique word ratio", "avg_unique_word_ratio", "f"),
|
| 563 |
+
("URLs found", "total_urls", "d"),
|
| 564 |
+
("Docs with HTML", "docs_with_html", "d"),
|
| 565 |
+
]
|
| 566 |
+
for label, key, fmt in metrics_rows:
|
| 567 |
+
bv = before_metrics.get(key, 0)
|
| 568 |
+
av = after_metrics.get(key, 0)
|
| 569 |
+
diff = av - bv
|
| 570 |
+
sign = "+" if diff > 0 else ""
|
| 571 |
+
if fmt == "d":
|
| 572 |
+
report.append(f" {label:<30} {bv:>12,} {av:>12,} {sign}{diff:>11,}")
|
| 573 |
+
elif fmt == "f1":
|
| 574 |
+
report.append(f" {label:<30} {bv:>12.1f} {av:>12.1f} {sign}{diff:>11.1f}")
|
| 575 |
+
else:
|
| 576 |
+
report.append(f" {label:<30} {bv:>12.4f} {av:>12.4f} {sign}{diff:>11.4f}")
|
| 577 |
+
|
| 578 |
+
# Sample diffs
|
| 579 |
+
report.append(f"\n SAMPLE BEFORE/AFTER DIFFS (showing {min(5, len(before_after_pairs))} pairs)")
|
| 580 |
+
report.append(f" {'-'*60}")
|
| 581 |
+
for idx, (before, after) in enumerate(before_after_pairs[:5]):
|
| 582 |
+
report.append(f"\n --- Sample {idx+1} ---")
|
| 583 |
+
report.append(f" BEFORE ({len(before)} chars, {len(before.split())} words):")
|
| 584 |
+
preview_b = before[:300].replace('\n', ' | ')
|
| 585 |
+
report.append(f" \"{preview_b}...\"")
|
| 586 |
+
report.append(f" AFTER ({len(after)} chars, {len(after.split())} words):")
|
| 587 |
+
preview_a = after[:300].replace('\n', ' | ')
|
| 588 |
+
report.append(f" \"{preview_a}...\"")
|
| 589 |
+
|
| 590 |
+
# Dropped samples
|
| 591 |
+
if dropped_texts:
|
| 592 |
+
report.append(f"\n EXAMPLES OF DROPPED DOCUMENTS ({min(3, len(dropped_texts))} shown)")
|
| 593 |
+
report.append(f" {'-'*60}")
|
| 594 |
+
for idx, dt in enumerate(dropped_texts[:3]):
|
| 595 |
+
report.append(f" Dropped #{idx+1} ({len(dt)} chars): \"{dt[:150]}...\"")
|
| 596 |
+
|
| 597 |
+
report.append(f"\n {'='*60}")
|
| 598 |
+
report.append(f" VERDICT: litdata_english cleaned and ready!")
|
| 599 |
+
report.append(f" Output: {output_dir}")
|
| 600 |
+
report.append(f" {'='*60}")
|
| 601 |
+
|
| 602 |
+
# Print the full report
|
| 603 |
+
full_report = '\n'.join(report)
|
| 604 |
+
print(full_report)
|
| 605 |
+
|
| 606 |
+
# Save report
|
| 607 |
+
report_path = output_dir / "CLEANING_REPORT.txt"
|
| 608 |
+
with open(report_path, "w", encoding="utf-8") as f:
|
| 609 |
+
f.write(full_report)
|
| 610 |
+
print(f"\n Report saved to: {report_path}")
|
Base/scripts/reclean_litdata.py
ADDED
|
@@ -0,0 +1,444 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""
|
| 3 |
+
Reclean & normalize all pretraining data for optimal 100M model learning.
|
| 4 |
+
|
| 5 |
+
Reads litdata_3b and litdata_english, decodes all documents back to text,
|
| 6 |
+
applies comprehensive English cleaning/normalization, re-tokenizes, and
|
| 7 |
+
writes new litdata chunks.
|
| 8 |
+
"""
|
| 9 |
+
|
| 10 |
+
import json
|
| 11 |
+
import os
|
| 12 |
+
import re
|
| 13 |
+
import sys
|
| 14 |
+
import time
|
| 15 |
+
import unicodedata
|
| 16 |
+
from pathlib import Path
|
| 17 |
+
|
| 18 |
+
import numpy as np
|
| 19 |
+
from tokenizers import Tokenizer
|
| 20 |
+
|
| 21 |
+
ROOT = Path(__file__).resolve().parent.parent.parent
|
| 22 |
+
BLOCK_SIZE = 1025
|
| 23 |
+
DTYPE = np.int32
|
| 24 |
+
DTYPE_INDEX = 16
|
| 25 |
+
CHUNK_BYTES_TARGET = 64 * 1024 * 1024
|
| 26 |
+
EOS_TOKEN_ID = 0
|
| 27 |
+
|
| 28 |
+
# -- Load tokenizer -----------------------------------------------------------
|
| 29 |
+
print("Loading tokenizer...")
|
| 30 |
+
tokenizer = Tokenizer.from_file(
|
| 31 |
+
str(ROOT / "Base" / "checkpoints" / "EleutherAI" / "pythia-160m" / "tokenizer.json")
|
| 32 |
+
)
|
| 33 |
+
|
| 34 |
+
# ==============================================================================
|
| 35 |
+
# TEXT CLEANING PIPELINE
|
| 36 |
+
# ==============================================================================
|
| 37 |
+
|
| 38 |
+
# Null / control chars to strip
|
| 39 |
+
CONTROL_CHARS = [
|
| 40 |
+
"\x00", "\x01", "\x02", "\x03", "\x04", "\x05", "\x06", "\x07",
|
| 41 |
+
"\x08", "\x0b", "\x0c", "\x0e", "\x0f", "\x10", "\x11", "\x12",
|
| 42 |
+
"\x13", "\x14", "\x15", "\x16", "\x17", "\x18", "\x19", "\x1a",
|
| 43 |
+
"\x1b", "\x1c", "\x1d", "\x1e", "\x1f", "\x7f",
|
| 44 |
+
"\ufeff", "\ufffd",
|
| 45 |
+
]
|
| 46 |
+
|
| 47 |
+
# HTML entities
|
| 48 |
+
HTML_ENTITIES = [
|
| 49 |
+
("&", "&"), ("<", "<"), (">", ">"),
|
| 50 |
+
(""", '"'), ("'", "'"), ("'", "'"),
|
| 51 |
+
(" ", " "), ("—", " - "), ("–", "-"),
|
| 52 |
+
("…", "..."), ("«", '"'), ("»", '"'),
|
| 53 |
+
("•", "- "), ("·", " "), ("©", "(c)"),
|
| 54 |
+
("®", "(R)"), ("™", "(TM)"), ("°", " degrees"),
|
| 55 |
+
]
|
| 56 |
+
|
| 57 |
+
# URL/email/path patterns
|
| 58 |
+
RE_URL = re.compile(r'https?://\S+|www\.\S+', re.I)
|
| 59 |
+
RE_EMAIL = re.compile(r'\b[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}\b')
|
| 60 |
+
RE_FILE_PATH = re.compile(r'(?:[A-Z]:\\|/(?:home|usr|var|etc|opt)/)\S+')
|
| 61 |
+
|
| 62 |
+
# HTML tag leftovers
|
| 63 |
+
RE_HTML_TAG = re.compile(r'</?[a-zA-Z][a-zA-Z0-9]*(?:\s[^>]*)?\s*/?>')
|
| 64 |
+
RE_HTML_COMMENT = re.compile(r'<!--.*?-->', re.DOTALL)
|
| 65 |
+
|
| 66 |
+
# Code/programming artifacts
|
| 67 |
+
RE_CODE_BLOCK = re.compile(r'```[\s\S]*?```')
|
| 68 |
+
RE_IMPORT = re.compile(r'^(?:import |from \S+ import |#include |using namespace |require\()', re.M)
|
| 69 |
+
|
| 70 |
+
# Repeated content
|
| 71 |
+
RE_REPEATED_LINE = re.compile(r'^(.{20,})\n(?:\1\n?)+', re.M)
|
| 72 |
+
RE_REPEATED_PUNCT = re.compile(r'([!?.])\1{3,}')
|
| 73 |
+
RE_REPEATED_CHAR = re.compile(r'(.)\1{5,}')
|
| 74 |
+
RE_REPEATED_WORD = re.compile(r'\b(\w+)(?:\s+\1){2,}\b', re.I)
|
| 75 |
+
|
| 76 |
+
# Whitespace
|
| 77 |
+
RE_MULTI_NEWLINE = re.compile(r'\n{4,}')
|
| 78 |
+
RE_MULTI_SPACE = re.compile(r'[ \t]{2,}')
|
| 79 |
+
RE_TRAILING_SPACE = re.compile(r'[ \t]+$', re.M)
|
| 80 |
+
|
| 81 |
+
# Sentence fixing
|
| 82 |
+
RE_NO_SPACE_AFTER_PERIOD = re.compile(r'([.!?])([A-Z])')
|
| 83 |
+
RE_DOUBLE_PERIOD = re.compile(r'\.{2}(?!\.)') # .. but not ...
|
| 84 |
+
RE_SPACE_BEFORE_PUNCT = re.compile(r'\s+([.,;:!?])')
|
| 85 |
+
|
| 86 |
+
|
| 87 |
+
def clean_text(text):
|
| 88 |
+
"""Apply full cleaning pipeline to a single document text."""
|
| 89 |
+
if not text or len(text.strip()) < 30:
|
| 90 |
+
return None
|
| 91 |
+
|
| 92 |
+
# 1. Unicode normalization
|
| 93 |
+
text = unicodedata.normalize("NFKC", text)
|
| 94 |
+
|
| 95 |
+
# 2. Strip control characters
|
| 96 |
+
for ch in CONTROL_CHARS:
|
| 97 |
+
text = text.replace(ch, "")
|
| 98 |
+
|
| 99 |
+
# 3. Fix HTML entities
|
| 100 |
+
for old, new in HTML_ENTITIES:
|
| 101 |
+
text = text.replace(old, new)
|
| 102 |
+
|
| 103 |
+
# 4. Remove HTML tags and comments
|
| 104 |
+
text = RE_HTML_COMMENT.sub("", text)
|
| 105 |
+
text = RE_HTML_TAG.sub("", text)
|
| 106 |
+
|
| 107 |
+
# 5. Remove URLs, emails, file paths
|
| 108 |
+
text = RE_URL.sub("", text)
|
| 109 |
+
text = RE_EMAIL.sub("", text)
|
| 110 |
+
text = RE_FILE_PATH.sub("", text)
|
| 111 |
+
|
| 112 |
+
# 6. Remove code blocks
|
| 113 |
+
text = RE_CODE_BLOCK.sub("", text)
|
| 114 |
+
|
| 115 |
+
# 7. Fix repeated content
|
| 116 |
+
text = RE_REPEATED_LINE.sub(r'\1', text)
|
| 117 |
+
text = RE_REPEATED_PUNCT.sub(r'\1\1\1', text)
|
| 118 |
+
text = RE_REPEATED_CHAR.sub(r'\1\1\1', text)
|
| 119 |
+
text = RE_REPEATED_WORD.sub(r'\1', text)
|
| 120 |
+
|
| 121 |
+
# 8. Normalize whitespace
|
| 122 |
+
text = text.replace('\t', ' ')
|
| 123 |
+
text = RE_TRAILING_SPACE.sub('', text)
|
| 124 |
+
text = RE_MULTI_SPACE.sub(' ', text)
|
| 125 |
+
text = RE_MULTI_NEWLINE.sub('\n\n\n', text)
|
| 126 |
+
|
| 127 |
+
# 9. Fix punctuation
|
| 128 |
+
text = RE_DOUBLE_PERIOD.sub('.', text)
|
| 129 |
+
text = RE_NO_SPACE_AFTER_PERIOD.sub(r'\1 \2', text)
|
| 130 |
+
text = RE_SPACE_BEFORE_PUNCT.sub(r'\1', text)
|
| 131 |
+
|
| 132 |
+
# 10. Normalize smart quotes and special punctuation to ASCII
|
| 133 |
+
text = text.replace('\u2018', "'").replace('\u2019', "'")
|
| 134 |
+
text = text.replace('\u201c', '"').replace('\u201d', '"')
|
| 135 |
+
text = text.replace('\u2013', '-').replace('\u2014', ' - ')
|
| 136 |
+
text = text.replace('\u2026', '...')
|
| 137 |
+
text = text.replace('\u2022', '- ')
|
| 138 |
+
text = text.replace('\u00b7', ' ')
|
| 139 |
+
text = text.replace('\u00a0', ' ') # non-breaking space
|
| 140 |
+
|
| 141 |
+
# 11. Process line by line: capitalize sentence starts, remove junk lines
|
| 142 |
+
lines = text.split('\n')
|
| 143 |
+
clean_lines = []
|
| 144 |
+
for line in lines:
|
| 145 |
+
line = line.strip()
|
| 146 |
+
if not line:
|
| 147 |
+
clean_lines.append('')
|
| 148 |
+
continue
|
| 149 |
+
|
| 150 |
+
# Skip lines that are mostly non-alphabetic (tables, code, etc.)
|
| 151 |
+
if len(line) > 10:
|
| 152 |
+
alpha_count = sum(1 for c in line if c.isalpha())
|
| 153 |
+
if alpha_count / len(line) < 0.40:
|
| 154 |
+
continue
|
| 155 |
+
|
| 156 |
+
# Skip lines with too many special chars (tables, markup)
|
| 157 |
+
if line.count('|') > 3 or line.count('{') > 2 or line.count('}') > 2:
|
| 158 |
+
continue
|
| 159 |
+
|
| 160 |
+
# Skip lines that look like code imports
|
| 161 |
+
if RE_IMPORT.match(line):
|
| 162 |
+
continue
|
| 163 |
+
|
| 164 |
+
# Capitalize first letter of sentences
|
| 165 |
+
if line and line[0].isalpha() and line[0].islower():
|
| 166 |
+
if not clean_lines or clean_lines[-1] == '' or clean_lines[-1].rstrip().endswith(('.', '!', '?', ':')):
|
| 167 |
+
line = line[0].upper() + line[1:]
|
| 168 |
+
|
| 169 |
+
clean_lines.append(line)
|
| 170 |
+
|
| 171 |
+
text = '\n'.join(clean_lines)
|
| 172 |
+
|
| 173 |
+
# 12. Remove leading/trailing whitespace
|
| 174 |
+
text = text.strip()
|
| 175 |
+
|
| 176 |
+
# 13. Remove duplicate paragraphs
|
| 177 |
+
paragraphs = text.split('\n\n')
|
| 178 |
+
seen = set()
|
| 179 |
+
unique_paragraphs = []
|
| 180 |
+
for p in paragraphs:
|
| 181 |
+
p_stripped = p.strip()
|
| 182 |
+
if not p_stripped:
|
| 183 |
+
continue
|
| 184 |
+
p_key = ' '.join(p_stripped.lower().split())
|
| 185 |
+
if p_key not in seen:
|
| 186 |
+
seen.add(p_key)
|
| 187 |
+
unique_paragraphs.append(p_stripped)
|
| 188 |
+
text = '\n\n'.join(unique_paragraphs)
|
| 189 |
+
|
| 190 |
+
# 14. Final quality gate
|
| 191 |
+
text = text.strip()
|
| 192 |
+
if len(text) < 50:
|
| 193 |
+
return None
|
| 194 |
+
word_count = len(text.split())
|
| 195 |
+
if word_count < 10:
|
| 196 |
+
return None
|
| 197 |
+
# Must be mostly ASCII/English
|
| 198 |
+
ascii_count = sum(1 for c in text if ord(c) < 128)
|
| 199 |
+
if ascii_count / max(len(text), 1) < 0.85:
|
| 200 |
+
return None
|
| 201 |
+
|
| 202 |
+
return text
|
| 203 |
+
|
| 204 |
+
|
| 205 |
+
# ==============================================================================
|
| 206 |
+
# LITDATA I/O
|
| 207 |
+
# ==============================================================================
|
| 208 |
+
|
| 209 |
+
def read_all_tokens(litdata_dir):
|
| 210 |
+
"""Read all chunks and return the full flat token stream as numpy array."""
|
| 211 |
+
with open(litdata_dir / "index.json") as f:
|
| 212 |
+
index = json.load(f)
|
| 213 |
+
|
| 214 |
+
chunks = index["chunks"]
|
| 215 |
+
total_tokens = sum(c["dim"] for c in chunks)
|
| 216 |
+
print(f" Reading {len(chunks)} chunks ({total_tokens:,} tokens)...")
|
| 217 |
+
|
| 218 |
+
all_tokens = np.empty(total_tokens, dtype=DTYPE)
|
| 219 |
+
pos = 0
|
| 220 |
+
|
| 221 |
+
for i, chunk in enumerate(chunks):
|
| 222 |
+
chunk_path = litdata_dir / chunk["filename"]
|
| 223 |
+
n_blocks = chunk["chunk_size"]
|
| 224 |
+
header_ints = 1 + n_blocks + 1
|
| 225 |
+
header_bytes = header_ints * 4
|
| 226 |
+
|
| 227 |
+
with open(chunk_path, "rb") as f:
|
| 228 |
+
f.seek(header_bytes)
|
| 229 |
+
data = np.fromfile(f, dtype=DTYPE, count=chunk["dim"])
|
| 230 |
+
|
| 231 |
+
all_tokens[pos:pos + len(data)] = data
|
| 232 |
+
pos += len(data)
|
| 233 |
+
|
| 234 |
+
if (i + 1) % 50 == 0 or i == len(chunks) - 1:
|
| 235 |
+
print(f" Read {i+1}/{len(chunks)} chunks ({pos:,} tokens)")
|
| 236 |
+
|
| 237 |
+
return all_tokens[:pos]
|
| 238 |
+
|
| 239 |
+
|
| 240 |
+
def split_documents(token_stream):
|
| 241 |
+
"""Split token stream by EOS token (0) into individual documents."""
|
| 242 |
+
eos_positions = np.where(token_stream == EOS_TOKEN_ID)[0]
|
| 243 |
+
docs = []
|
| 244 |
+
start = 0
|
| 245 |
+
for eos_pos in eos_positions:
|
| 246 |
+
if eos_pos > start:
|
| 247 |
+
docs.append(token_stream[start:eos_pos])
|
| 248 |
+
start = eos_pos + 1
|
| 249 |
+
if start < len(token_stream):
|
| 250 |
+
docs.append(token_stream[start:])
|
| 251 |
+
return docs
|
| 252 |
+
|
| 253 |
+
|
| 254 |
+
def write_litdata_chunks(output_dir, token_stream, config):
|
| 255 |
+
"""Write token stream as litdata chunks, returns index metadata."""
|
| 256 |
+
os.makedirs(output_dir, exist_ok=True)
|
| 257 |
+
|
| 258 |
+
dtype_size = DTYPE().itemsize
|
| 259 |
+
tokens_per_chunk = CHUNK_BYTES_TARGET // dtype_size
|
| 260 |
+
tokens_per_chunk = (tokens_per_chunk // BLOCK_SIZE) * BLOCK_SIZE
|
| 261 |
+
|
| 262 |
+
chunks_metadata = []
|
| 263 |
+
pos = 0
|
| 264 |
+
chunk_idx = 0
|
| 265 |
+
|
| 266 |
+
while pos < len(token_stream):
|
| 267 |
+
remaining = len(token_stream) - pos
|
| 268 |
+
chunk_tokens = min(tokens_per_chunk, remaining)
|
| 269 |
+
num_blocks = chunk_tokens // BLOCK_SIZE
|
| 270 |
+
if num_blocks == 0:
|
| 271 |
+
break
|
| 272 |
+
actual_tokens = num_blocks * BLOCK_SIZE
|
| 273 |
+
|
| 274 |
+
chunk_data = token_stream[pos:pos + actual_tokens]
|
| 275 |
+
filename = f"chunk-0-{chunk_idx}.bin"
|
| 276 |
+
filepath = os.path.join(output_dir, filename)
|
| 277 |
+
|
| 278 |
+
# Header: [num_items(uint32)] + [offsets 0..num_blocks(uint32)]
|
| 279 |
+
header_num_items = np.array([num_blocks], dtype=np.uint32)
|
| 280 |
+
offsets = np.arange(num_blocks + 1, dtype=np.uint32) * (BLOCK_SIZE * dtype_size)
|
| 281 |
+
header = np.concatenate([header_num_items, offsets])
|
| 282 |
+
|
| 283 |
+
with open(filepath, "wb") as f:
|
| 284 |
+
header.tofile(f)
|
| 285 |
+
chunk_data.tofile(f)
|
| 286 |
+
|
| 287 |
+
meta = {
|
| 288 |
+
"chunk_bytes": int(header.nbytes + chunk_data.nbytes),
|
| 289 |
+
"chunk_size": num_blocks,
|
| 290 |
+
"dim": int(actual_tokens),
|
| 291 |
+
"filename": filename,
|
| 292 |
+
}
|
| 293 |
+
chunks_metadata.append(meta)
|
| 294 |
+
pos += actual_tokens
|
| 295 |
+
chunk_idx += 1
|
| 296 |
+
|
| 297 |
+
if chunk_idx % 25 == 0 or pos >= len(token_stream):
|
| 298 |
+
print(f" Written chunk {chunk_idx} ({pos:,}/{len(token_stream):,} tokens)")
|
| 299 |
+
|
| 300 |
+
# Write index.json
|
| 301 |
+
index = {
|
| 302 |
+
"chunks": chunks_metadata,
|
| 303 |
+
"config": config,
|
| 304 |
+
"updated_at": str(time.time()),
|
| 305 |
+
}
|
| 306 |
+
with open(os.path.join(output_dir, "index.json"), "w") as f:
|
| 307 |
+
json.dump(index, f, indent=2)
|
| 308 |
+
|
| 309 |
+
return chunks_metadata
|
| 310 |
+
|
| 311 |
+
|
| 312 |
+
# ==============================================================================
|
| 313 |
+
# MAIN PROCESSING
|
| 314 |
+
# ==============================================================================
|
| 315 |
+
|
| 316 |
+
def process_litdata(input_dir, output_dir, name):
|
| 317 |
+
print(f"\n{'='*65}")
|
| 318 |
+
print(f" PROCESSING: {name}")
|
| 319 |
+
print(f" Input: {input_dir}")
|
| 320 |
+
print(f" Output: {output_dir}")
|
| 321 |
+
print(f"{'='*65}")
|
| 322 |
+
|
| 323 |
+
# 1. Read all tokens
|
| 324 |
+
t0 = time.time()
|
| 325 |
+
token_stream = read_all_tokens(input_dir)
|
| 326 |
+
print(f" Read {len(token_stream):,} tokens in {time.time()-t0:.1f}s")
|
| 327 |
+
|
| 328 |
+
# 2. Split into documents
|
| 329 |
+
t1 = time.time()
|
| 330 |
+
doc_tokens = split_documents(token_stream)
|
| 331 |
+
print(f" Found {len(doc_tokens):,} documents in {time.time()-t1:.1f}s")
|
| 332 |
+
del token_stream
|
| 333 |
+
|
| 334 |
+
# 3. Decode all documents to text (batch for speed)
|
| 335 |
+
t2 = time.time()
|
| 336 |
+
print(f" Decoding documents back to text...")
|
| 337 |
+
raw_texts = []
|
| 338 |
+
BATCH = 5000
|
| 339 |
+
for i in range(0, len(doc_tokens), BATCH):
|
| 340 |
+
batch = doc_tokens[i:i+BATCH]
|
| 341 |
+
for doc in batch:
|
| 342 |
+
text = tokenizer.decode(doc.tolist(), skip_special_tokens=False)
|
| 343 |
+
raw_texts.append(text)
|
| 344 |
+
done = min(i + BATCH, len(doc_tokens))
|
| 345 |
+
if done % 100000 == 0 or done == len(doc_tokens):
|
| 346 |
+
print(f" Decoded {done:,}/{len(doc_tokens):,}")
|
| 347 |
+
del doc_tokens
|
| 348 |
+
print(f" Decoded in {time.time()-t2:.1f}s")
|
| 349 |
+
|
| 350 |
+
# 4. Clean each document
|
| 351 |
+
t3 = time.time()
|
| 352 |
+
print(f" Cleaning {len(raw_texts):,} documents...")
|
| 353 |
+
cleaned_texts = []
|
| 354 |
+
dropped = 0
|
| 355 |
+
for i, text in enumerate(raw_texts):
|
| 356 |
+
result = clean_text(text)
|
| 357 |
+
if result is not None:
|
| 358 |
+
cleaned_texts.append(result)
|
| 359 |
+
else:
|
| 360 |
+
dropped += 1
|
| 361 |
+
if (i + 1) % 200000 == 0 or i == len(raw_texts) - 1:
|
| 362 |
+
print(f" Processed {i+1:,}/{len(raw_texts):,} | kept={len(cleaned_texts):,} | dropped={dropped:,}")
|
| 363 |
+
del raw_texts
|
| 364 |
+
print(f" Cleaning done in {time.time()-t3:.1f}s")
|
| 365 |
+
print(f" Kept {len(cleaned_texts):,} docs | Dropped {dropped:,} ({dropped/(max(dropped+len(cleaned_texts),1))*100:.1f}%)")
|
| 366 |
+
|
| 367 |
+
# 5. Re-tokenize cleaned texts
|
| 368 |
+
t4 = time.time()
|
| 369 |
+
print(f" Re-tokenizing {len(cleaned_texts):,} documents...")
|
| 370 |
+
all_token_ids = []
|
| 371 |
+
total_new_tokens = 0
|
| 372 |
+
ENCODE_BATCH = 10000
|
| 373 |
+
for i in range(0, len(cleaned_texts), ENCODE_BATCH):
|
| 374 |
+
batch = cleaned_texts[i:i+ENCODE_BATCH]
|
| 375 |
+
encoded = tokenizer.encode_batch(batch, add_special_tokens=False)
|
| 376 |
+
for enc in encoded:
|
| 377 |
+
ids = enc.ids
|
| 378 |
+
all_token_ids.extend(ids)
|
| 379 |
+
all_token_ids.append(EOS_TOKEN_ID)
|
| 380 |
+
total_new_tokens += len(ids) + 1
|
| 381 |
+
done = min(i + ENCODE_BATCH, len(cleaned_texts))
|
| 382 |
+
if done % 200000 == 0 or done == len(cleaned_texts):
|
| 383 |
+
print(f" Tokenized {done:,}/{len(cleaned_texts):,} ({total_new_tokens:,} tokens)")
|
| 384 |
+
del cleaned_texts
|
| 385 |
+
print(f" Tokenized in {time.time()-t4:.1f}s")
|
| 386 |
+
print(f" New total: {total_new_tokens:,} tokens")
|
| 387 |
+
|
| 388 |
+
# 6. Convert to numpy and write chunks
|
| 389 |
+
t5 = time.time()
|
| 390 |
+
print(f" Building token stream array...")
|
| 391 |
+
new_stream = np.array(all_token_ids, dtype=DTYPE)
|
| 392 |
+
del all_token_ids
|
| 393 |
+
|
| 394 |
+
with open(input_dir / "index.json") as f:
|
| 395 |
+
config = json.load(f)["config"]
|
| 396 |
+
|
| 397 |
+
print(f" Writing litdata chunks...")
|
| 398 |
+
chunks = write_litdata_chunks(str(output_dir), new_stream, config)
|
| 399 |
+
print(f" Written {len(chunks)} chunks in {time.time()-t5:.1f}s")
|
| 400 |
+
|
| 401 |
+
total_in = sum(c["dim"] for c in json.load(open(input_dir / "index.json"))["chunks"])
|
| 402 |
+
total_out = sum(c["dim"] for c in chunks)
|
| 403 |
+
print(f"\n SUMMARY for {name}:")
|
| 404 |
+
print(f" Input tokens: {total_in:,}")
|
| 405 |
+
print(f" Output tokens: {total_out:,}")
|
| 406 |
+
print(f" Difference: {total_in - total_out:,} ({(total_in-total_out)/total_in*100:.2f}% removed)")
|
| 407 |
+
|
| 408 |
+
return total_in, total_out
|
| 409 |
+
|
| 410 |
+
|
| 411 |
+
if __name__ == "__main__":
|
| 412 |
+
t_start = time.time()
|
| 413 |
+
data_dir = ROOT / "Base" / "data"
|
| 414 |
+
|
| 415 |
+
# Process litdata_3b
|
| 416 |
+
orig_3b, clean_3b = process_litdata(
|
| 417 |
+
data_dir / "litdata_3b",
|
| 418 |
+
data_dir / "litdata_3b_clean",
|
| 419 |
+
"litdata_3b (General Knowledge)",
|
| 420 |
+
)
|
| 421 |
+
|
| 422 |
+
# Process litdata_english
|
| 423 |
+
orig_en, clean_en = process_litdata(
|
| 424 |
+
data_dir / "litdata_english",
|
| 425 |
+
data_dir / "litdata_english_clean",
|
| 426 |
+
"litdata_english (English Knowledge)",
|
| 427 |
+
)
|
| 428 |
+
|
| 429 |
+
# Final report
|
| 430 |
+
print(f"\n{'='*65}")
|
| 431 |
+
print(f" FINAL REPORT")
|
| 432 |
+
print(f"{'='*65}")
|
| 433 |
+
print(f" litdata_3b: {orig_3b:>15,} -> {clean_3b:>15,} tokens")
|
| 434 |
+
print(f" litdata_english: {orig_en:>15,} -> {clean_en:>15,} tokens")
|
| 435 |
+
print(f" ---------------------------------------------------------")
|
| 436 |
+
total_orig = orig_3b + orig_en
|
| 437 |
+
total_clean = clean_3b + clean_en
|
| 438 |
+
print(f" TOTAL: {total_orig:>15,} -> {total_clean:>15,} tokens")
|
| 439 |
+
print(f" Removed: {total_orig - total_clean:,} ({(total_orig-total_clean)/total_orig*100:.2f}%)")
|
| 440 |
+
print(f"\n Total time: {time.time()-t_start:.0f}s")
|
| 441 |
+
print(f"\n Clean data ready at:")
|
| 442 |
+
print(f" {data_dir / 'litdata_3b_clean'}")
|
| 443 |
+
print(f" {data_dir / 'litdata_english_clean'}")
|
| 444 |
+
print(f"\n Update your training configs to point to the _clean directories!")
|
Base/scripts/smart_cleanup_english.py
ADDED
|
@@ -0,0 +1,379 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# -*- coding: utf-8 -*-
|
| 2 |
+
"""
|
| 3 |
+
Smart cleanup of litdata_english_clean based on deep audit results.
|
| 4 |
+
|
| 5 |
+
Removes truly problematic docs while keeping legitimate educational content.
|
| 6 |
+
Strategy:
|
| 7 |
+
- REMOVE: docs with non-English script fragments (Cyrillic, CJK, Arabic, Devanagari)
|
| 8 |
+
- REMOVE: docs under 50 words
|
| 9 |
+
- REMOVE: very repetitive docs (unique word ratio < 0.20)
|
| 10 |
+
- REMOVE: docs with heavy residual code (5+ code patterns)
|
| 11 |
+
- KEEP: educational/historical content that mentions historical terms in context
|
| 12 |
+
(these are Wikipedia articles about WWII, civil rights, etc. - valuable learning)
|
| 13 |
+
- STRIP: clickbait phrases, subscribe prompts, cookie/login boilerplate FROM docs
|
| 14 |
+
(remove the noise parts, keep the content)
|
| 15 |
+
|
| 16 |
+
Rebuilds clean litdata chunks after filtering.
|
| 17 |
+
"""
|
| 18 |
+
|
| 19 |
+
import json
|
| 20 |
+
import os
|
| 21 |
+
import re
|
| 22 |
+
import time
|
| 23 |
+
import unicodedata
|
| 24 |
+
from pathlib import Path
|
| 25 |
+
from collections import Counter
|
| 26 |
+
|
| 27 |
+
import numpy as np
|
| 28 |
+
from tokenizers import Tokenizer
|
| 29 |
+
|
| 30 |
+
ROOT = Path(__file__).resolve().parent.parent.parent
|
| 31 |
+
BLOCK_SIZE = 1025
|
| 32 |
+
DTYPE = np.int32
|
| 33 |
+
CHUNK_BYTES_TARGET = 64 * 1024 * 1024
|
| 34 |
+
EOS_TOKEN_ID = 0
|
| 35 |
+
|
| 36 |
+
print("Loading tokenizer...")
|
| 37 |
+
tokenizer = Tokenizer.from_file(
|
| 38 |
+
str(ROOT / "Base" / "checkpoints" / "EleutherAI" / "pythia-160m" / "tokenizer.json")
|
| 39 |
+
)
|
| 40 |
+
|
| 41 |
+
# ==============================================================================
|
| 42 |
+
# PATTERNS
|
| 43 |
+
# ==============================================================================
|
| 44 |
+
|
| 45 |
+
# Non-English script detectors
|
| 46 |
+
RE_CJK = re.compile(r'[\u4e00-\u9fff\u3040-\u309f\u30a0-\u30ff\uac00-\ud7af]{3,}')
|
| 47 |
+
RE_ARABIC = re.compile(r'[\u0600-\u06ff]{5,}')
|
| 48 |
+
RE_CYRILLIC = re.compile(r'[\u0400-\u04ff]{5,}')
|
| 49 |
+
RE_DEVANAGARI = re.compile(r'[\u0900-\u097f]{5,}')
|
| 50 |
+
|
| 51 |
+
# Residual code
|
| 52 |
+
RE_RESIDUAL_CODE = re.compile(r'(function\s*\(|var\s+\w+\s*=|console\.log|document\.get|if\s*\(\s*\w+\s*[!=]==)', re.I)
|
| 53 |
+
|
| 54 |
+
# Boilerplate to STRIP from docs (not remove doc, just strip these lines)
|
| 55 |
+
RE_COOKIE_LINE = re.compile(r'^.*(?:cookie|cookies)\s+(?:policy|consent|notice|preferences|settings).*$', re.I | re.M)
|
| 56 |
+
RE_SUBSCRIBE_LINE = re.compile(r'^.*(?:subscribe|sign\s*up\s+(?:for|to)\s+(?:our|the)\s+newsletter|unsubscribe|opt[\s-]*out\s+of).*$', re.I | re.M)
|
| 57 |
+
RE_CLICKBAIT_LINE = re.compile(r'^.*(?:you\s+won\'?t\s+believe|click\s+here|read\s+more\s*\.{0,3}$|share\s+this\s+(?:article|post|story)|trending\s+now|sponsored\s+content|advertisement).*$', re.I | re.M)
|
| 58 |
+
RE_SOCIAL_LINE = re.compile(r'^.*(?:follow\s+us\s+on|share\s+on\s+(?:facebook|twitter|linkedin|instagram)|like\s+us\s+on|tweet\s+this).*$', re.I | re.M)
|
| 59 |
+
RE_NAV_LINE = re.compile(r'^.*(?:skip\s+to\s+(?:main\s+)?content|back\s+to\s+top|previous\s+article|next\s+article|related\s+(?:articles|posts)).*$', re.I | re.M)
|
| 60 |
+
RE_LOGIN_LINE = re.compile(r'^.*(?:log\s*in\s+to\s+(?:your|an)\s+account|create\s+(?:a\s+)?(?:free\s+)?account|forgot\s+(?:your\s+)?password|already\s+(?:a\s+)?member).*$', re.I | re.M)
|
| 61 |
+
RE_COMMENT_LINE = re.compile(r'^.*(?:leave\s+a\s+(?:comment|reply)|post\s+a\s+comment|\d+\s+comments?$|logged\s+in\s+as).*$', re.I | re.M)
|
| 62 |
+
|
| 63 |
+
# Lines that are just copyright
|
| 64 |
+
RE_COPYRIGHT_LINE = re.compile(r'^.*(?:all\s+rights\s+reserved|\(c\)\s*\d{4}|copyright\s+\d{4}).*$', re.I | re.M)
|
| 65 |
+
|
| 66 |
+
# Number-heavy lines (tables of just numbers)
|
| 67 |
+
RE_NUMBER_TABLE_LINE = re.compile(r'^[\d\s,.\-+/%$]+$', re.M)
|
| 68 |
+
|
| 69 |
+
# ==============================================================================
|
| 70 |
+
# LITDATA I/O
|
| 71 |
+
# ==============================================================================
|
| 72 |
+
|
| 73 |
+
def read_all_tokens(litdata_dir):
|
| 74 |
+
with open(litdata_dir / "index.json") as f:
|
| 75 |
+
index = json.load(f)
|
| 76 |
+
chunks = index["chunks"]
|
| 77 |
+
total_tokens = sum(c["dim"] for c in chunks)
|
| 78 |
+
print(f" Reading {len(chunks)} chunks ({total_tokens:,} tokens)...")
|
| 79 |
+
all_tokens = np.empty(total_tokens, dtype=DTYPE)
|
| 80 |
+
pos = 0
|
| 81 |
+
for i, chunk in enumerate(chunks):
|
| 82 |
+
chunk_path = litdata_dir / chunk["filename"]
|
| 83 |
+
n_blocks = chunk["chunk_size"]
|
| 84 |
+
header_ints = 1 + n_blocks + 1
|
| 85 |
+
header_bytes = header_ints * 4
|
| 86 |
+
with open(chunk_path, "rb") as f:
|
| 87 |
+
f.seek(header_bytes)
|
| 88 |
+
data = np.fromfile(f, dtype=DTYPE, count=chunk["dim"])
|
| 89 |
+
all_tokens[pos:pos + len(data)] = data
|
| 90 |
+
pos += len(data)
|
| 91 |
+
print(f" Read {len(chunks)} chunks ({pos:,} tokens)")
|
| 92 |
+
return all_tokens[:pos]
|
| 93 |
+
|
| 94 |
+
|
| 95 |
+
def split_documents(token_stream):
|
| 96 |
+
eos_positions = np.where(token_stream == EOS_TOKEN_ID)[0]
|
| 97 |
+
docs = []
|
| 98 |
+
start = 0
|
| 99 |
+
for eos_pos in eos_positions:
|
| 100 |
+
if eos_pos > start:
|
| 101 |
+
docs.append(token_stream[start:eos_pos])
|
| 102 |
+
start = eos_pos + 1
|
| 103 |
+
if start < len(token_stream):
|
| 104 |
+
docs.append(token_stream[start:])
|
| 105 |
+
return docs
|
| 106 |
+
|
| 107 |
+
|
| 108 |
+
def write_litdata_chunks(output_dir, token_stream, config):
|
| 109 |
+
os.makedirs(output_dir, exist_ok=True)
|
| 110 |
+
dtype_size = DTYPE().itemsize
|
| 111 |
+
tokens_per_chunk = CHUNK_BYTES_TARGET // dtype_size
|
| 112 |
+
tokens_per_chunk = (tokens_per_chunk // BLOCK_SIZE) * BLOCK_SIZE
|
| 113 |
+
chunks_metadata = []
|
| 114 |
+
pos = 0
|
| 115 |
+
chunk_idx = 0
|
| 116 |
+
while pos < len(token_stream):
|
| 117 |
+
remaining = len(token_stream) - pos
|
| 118 |
+
chunk_tokens = min(tokens_per_chunk, remaining)
|
| 119 |
+
num_blocks = chunk_tokens // BLOCK_SIZE
|
| 120 |
+
if num_blocks == 0:
|
| 121 |
+
break
|
| 122 |
+
actual_tokens = num_blocks * BLOCK_SIZE
|
| 123 |
+
chunk_data = token_stream[pos:pos + actual_tokens]
|
| 124 |
+
filename = f"chunk-0-{chunk_idx}.bin"
|
| 125 |
+
filepath = os.path.join(output_dir, filename)
|
| 126 |
+
header_num_items = np.array([num_blocks], dtype=np.uint32)
|
| 127 |
+
offsets = np.arange(num_blocks + 1, dtype=np.uint32) * (BLOCK_SIZE * dtype_size)
|
| 128 |
+
header = np.concatenate([header_num_items, offsets])
|
| 129 |
+
with open(filepath, "wb") as f:
|
| 130 |
+
header.tofile(f)
|
| 131 |
+
chunk_data.tofile(f)
|
| 132 |
+
meta = {
|
| 133 |
+
"chunk_bytes": int(header.nbytes + chunk_data.nbytes),
|
| 134 |
+
"chunk_size": num_blocks,
|
| 135 |
+
"dim": int(actual_tokens),
|
| 136 |
+
"filename": filename,
|
| 137 |
+
}
|
| 138 |
+
chunks_metadata.append(meta)
|
| 139 |
+
pos += actual_tokens
|
| 140 |
+
chunk_idx += 1
|
| 141 |
+
print(f" Written chunk {chunk_idx} ({pos:,}/{len(token_stream):,} tokens)")
|
| 142 |
+
index = {"chunks": chunks_metadata, "config": config, "updated_at": str(time.time())}
|
| 143 |
+
with open(os.path.join(output_dir, "index.json"), "w") as f:
|
| 144 |
+
json.dump(index, f, indent=2)
|
| 145 |
+
return chunks_metadata
|
| 146 |
+
|
| 147 |
+
|
| 148 |
+
# ==============================================================================
|
| 149 |
+
# SMART CLEANUP
|
| 150 |
+
# ==============================================================================
|
| 151 |
+
|
| 152 |
+
def should_remove(text):
|
| 153 |
+
"""Returns (remove: bool, reason: str or None)."""
|
| 154 |
+
words = text.split()
|
| 155 |
+
word_count = len(words)
|
| 156 |
+
|
| 157 |
+
# Too short
|
| 158 |
+
if word_count < 50:
|
| 159 |
+
return True, f"too short ({word_count} words)"
|
| 160 |
+
|
| 161 |
+
# Non-English scripts
|
| 162 |
+
scripts_found = []
|
| 163 |
+
if RE_CJK.search(text):
|
| 164 |
+
scripts_found.append("CJK")
|
| 165 |
+
if RE_ARABIC.search(text):
|
| 166 |
+
scripts_found.append("Arabic")
|
| 167 |
+
if RE_CYRILLIC.search(text):
|
| 168 |
+
scripts_found.append("Cyrillic")
|
| 169 |
+
if RE_DEVANAGARI.search(text):
|
| 170 |
+
scripts_found.append("Devanagari")
|
| 171 |
+
if scripts_found:
|
| 172 |
+
return True, f"non-English scripts: {', '.join(scripts_found)}"
|
| 173 |
+
|
| 174 |
+
# Very repetitive
|
| 175 |
+
if word_count > 50:
|
| 176 |
+
unique_ratio = len(set(w.lower() for w in words)) / word_count
|
| 177 |
+
if unique_ratio < 0.20:
|
| 178 |
+
return True, f"very repetitive (unique ratio: {unique_ratio:.3f})"
|
| 179 |
+
|
| 180 |
+
# Heavy residual code (5+ code patterns)
|
| 181 |
+
code_matches = RE_RESIDUAL_CODE.findall(text)
|
| 182 |
+
if len(code_matches) >= 5:
|
| 183 |
+
return True, f"residual code ({len(code_matches)} matches)"
|
| 184 |
+
|
| 185 |
+
return False, None
|
| 186 |
+
|
| 187 |
+
|
| 188 |
+
def strip_boilerplate(text):
|
| 189 |
+
"""Strip boilerplate lines from a document, keeping the educational content."""
|
| 190 |
+
original_len = len(text)
|
| 191 |
+
|
| 192 |
+
# Strip specific boilerplate patterns (line by line removal)
|
| 193 |
+
for pattern in [
|
| 194 |
+
RE_COOKIE_LINE, RE_SUBSCRIBE_LINE, RE_CLICKBAIT_LINE,
|
| 195 |
+
RE_SOCIAL_LINE, RE_NAV_LINE, RE_LOGIN_LINE,
|
| 196 |
+
RE_COMMENT_LINE, RE_COPYRIGHT_LINE,
|
| 197 |
+
]:
|
| 198 |
+
text = pattern.sub('', text)
|
| 199 |
+
|
| 200 |
+
# Strip pure number table lines (> 80% digits/spaces)
|
| 201 |
+
lines = text.split('\n')
|
| 202 |
+
clean_lines = []
|
| 203 |
+
stripped_number_lines = 0
|
| 204 |
+
for line in lines:
|
| 205 |
+
stripped = line.strip()
|
| 206 |
+
if stripped and len(stripped) > 10:
|
| 207 |
+
digit_count = sum(1 for c in stripped if c.isdigit() or c in ' ,.\t-+/%$')
|
| 208 |
+
if digit_count / len(stripped) > 0.80:
|
| 209 |
+
stripped_number_lines += 1
|
| 210 |
+
continue
|
| 211 |
+
clean_lines.append(line)
|
| 212 |
+
text = '\n'.join(clean_lines)
|
| 213 |
+
|
| 214 |
+
# Clean up resulting whitespace
|
| 215 |
+
text = re.sub(r'\n{3,}', '\n\n', text)
|
| 216 |
+
text = text.strip()
|
| 217 |
+
|
| 218 |
+
chars_stripped = original_len - len(text)
|
| 219 |
+
return text, chars_stripped, stripped_number_lines
|
| 220 |
+
|
| 221 |
+
|
| 222 |
+
# ==============================================================================
|
| 223 |
+
# MAIN
|
| 224 |
+
# ==============================================================================
|
| 225 |
+
|
| 226 |
+
def main():
|
| 227 |
+
input_dir = ROOT / "Base" / "data" / "litdata_english_clean"
|
| 228 |
+
output_dir = ROOT / "Base" / "data" / "litdata_english_clean"
|
| 229 |
+
|
| 230 |
+
print(f"\n{'='*75}")
|
| 231 |
+
print(f" SMART CLEANUP: litdata_english_clean")
|
| 232 |
+
print(f" Input/Output: {input_dir}")
|
| 233 |
+
print(f"{'='*75}")
|
| 234 |
+
|
| 235 |
+
# 1. Read and decode
|
| 236 |
+
t0 = time.time()
|
| 237 |
+
token_stream = read_all_tokens(input_dir)
|
| 238 |
+
doc_tokens = split_documents(token_stream)
|
| 239 |
+
total_input = len(doc_tokens)
|
| 240 |
+
print(f" Found {total_input:,} documents")
|
| 241 |
+
del token_stream
|
| 242 |
+
|
| 243 |
+
print(f" Decoding ALL {total_input:,} documents...")
|
| 244 |
+
texts = []
|
| 245 |
+
t1 = time.time()
|
| 246 |
+
for i, toks in enumerate(doc_tokens):
|
| 247 |
+
text = tokenizer.decode(toks.tolist(), skip_special_tokens=False)
|
| 248 |
+
texts.append(text)
|
| 249 |
+
if (i + 1) % 20000 == 0 or i == total_input - 1:
|
| 250 |
+
print(f" Decoded {i+1:,}/{total_input:,}")
|
| 251 |
+
del doc_tokens
|
| 252 |
+
print(f" Decoded in {time.time()-t1:.1f}s")
|
| 253 |
+
|
| 254 |
+
# 2. Filter and strip
|
| 255 |
+
print(f"\n Processing {total_input:,} documents...")
|
| 256 |
+
t2 = time.time()
|
| 257 |
+
|
| 258 |
+
kept_texts = []
|
| 259 |
+
removed_reasons = Counter()
|
| 260 |
+
total_boilerplate_chars = 0
|
| 261 |
+
total_number_lines_stripped = 0
|
| 262 |
+
removed_examples = []
|
| 263 |
+
|
| 264 |
+
for i, text in enumerate(texts):
|
| 265 |
+
# Check if doc should be removed entirely
|
| 266 |
+
remove, reason = should_remove(text)
|
| 267 |
+
if remove:
|
| 268 |
+
removed_reasons[reason.split('(')[0].strip().split(':')[0].strip()] += 1
|
| 269 |
+
if len(removed_examples) < 20:
|
| 270 |
+
removed_examples.append((i, reason, text[:300]))
|
| 271 |
+
continue
|
| 272 |
+
|
| 273 |
+
# Strip boilerplate from kept docs
|
| 274 |
+
cleaned, chars_stripped, num_lines = strip_boilerplate(text)
|
| 275 |
+
total_boilerplate_chars += chars_stripped
|
| 276 |
+
total_number_lines_stripped += num_lines
|
| 277 |
+
|
| 278 |
+
# Final check: did stripping make it too short?
|
| 279 |
+
if len(cleaned.split()) < 50:
|
| 280 |
+
removed_reasons["stripped too short"] += 1
|
| 281 |
+
continue
|
| 282 |
+
|
| 283 |
+
kept_texts.append(cleaned)
|
| 284 |
+
|
| 285 |
+
if (i + 1) % 10000 == 0 or i == total_input - 1:
|
| 286 |
+
print(f" Processed {i+1:,}/{total_input:,} | kept={len(kept_texts):,} | removed={i+1-len(kept_texts):,}")
|
| 287 |
+
|
| 288 |
+
total_removed = total_input - len(kept_texts)
|
| 289 |
+
print(f" Processing done in {time.time()-t2:.1f}s")
|
| 290 |
+
print(f" Kept: {len(kept_texts):,} | Removed: {total_removed:,} ({total_removed/total_input*100:.2f}%)")
|
| 291 |
+
|
| 292 |
+
# 3. Re-tokenize
|
| 293 |
+
t3 = time.time()
|
| 294 |
+
print(f"\n Re-tokenizing {len(kept_texts):,} documents...")
|
| 295 |
+
all_token_ids = []
|
| 296 |
+
total_new_tokens = 0
|
| 297 |
+
for i, text in enumerate(kept_texts):
|
| 298 |
+
enc = tokenizer.encode(text)
|
| 299 |
+
ids = enc.ids
|
| 300 |
+
all_token_ids.extend(ids)
|
| 301 |
+
all_token_ids.append(EOS_TOKEN_ID)
|
| 302 |
+
total_new_tokens += len(ids) + 1
|
| 303 |
+
if (i + 1) % 20000 == 0 or i == len(kept_texts) - 1:
|
| 304 |
+
print(f" Tokenized {i+1:,}/{len(kept_texts):,} ({total_new_tokens:,} tokens)")
|
| 305 |
+
print(f" Tokenized in {time.time()-t3:.1f}s")
|
| 306 |
+
|
| 307 |
+
# 4. Rebuild chunks
|
| 308 |
+
t4 = time.time()
|
| 309 |
+
print(f"\n Rebuilding litdata chunks...")
|
| 310 |
+
token_array = np.array(all_token_ids, dtype=DTYPE)
|
| 311 |
+
del all_token_ids
|
| 312 |
+
|
| 313 |
+
import gc
|
| 314 |
+
gc.collect()
|
| 315 |
+
|
| 316 |
+
os.makedirs(str(output_dir), exist_ok=True)
|
| 317 |
+
config = {
|
| 318 |
+
"block_size": BLOCK_SIZE,
|
| 319 |
+
"vocab_size": tokenizer.get_vocab_size(),
|
| 320 |
+
}
|
| 321 |
+
chunks = write_litdata_chunks(str(output_dir), token_array, config)
|
| 322 |
+
del token_array
|
| 323 |
+
print(f" Written {len(chunks)} chunks in {time.time()-t4:.1f}s")
|
| 324 |
+
|
| 325 |
+
# 5. Summary report
|
| 326 |
+
total_time = time.time() - t0
|
| 327 |
+
|
| 328 |
+
report = []
|
| 329 |
+
report.append(f"\n{'='*75}")
|
| 330 |
+
report.append(f" SMART CLEANUP REPORT - litdata_english_clean")
|
| 331 |
+
report.append(f"{'='*75}")
|
| 332 |
+
report.append(f"\n Processing time: {total_time:.0f}s ({total_time/60:.1f} min)")
|
| 333 |
+
report.append(f"\n DOCUMENT COUNTS")
|
| 334 |
+
report.append(f" {'-'*50}")
|
| 335 |
+
report.append(f" Input documents: {total_input:>10,}")
|
| 336 |
+
report.append(f" Output documents: {len(kept_texts):>10,}")
|
| 337 |
+
report.append(f" Removed: {total_removed:>10,} ({total_removed/total_input*100:.2f}%)")
|
| 338 |
+
|
| 339 |
+
report.append(f"\n TOKEN COUNTS")
|
| 340 |
+
report.append(f" {'-'*50}")
|
| 341 |
+
report.append(f" Input tokens: {53446575:>15,}")
|
| 342 |
+
report.append(f" Output tokens: {total_new_tokens:>15,}")
|
| 343 |
+
report.append(f" Tokens after align: {sum(c['dim'] for c in chunks):>15,}")
|
| 344 |
+
|
| 345 |
+
report.append(f"\n REMOVAL REASONS")
|
| 346 |
+
report.append(f" {'-'*50}")
|
| 347 |
+
for reason, count in sorted(removed_reasons.items(), key=lambda x: -x[1]):
|
| 348 |
+
report.append(f" {reason:<35} {count:>6,}")
|
| 349 |
+
|
| 350 |
+
report.append(f"\n BOILERPLATE STRIPPED (from kept docs)")
|
| 351 |
+
report.append(f" {'-'*50}")
|
| 352 |
+
report.append(f" Total boilerplate chars stripped: {total_boilerplate_chars:>10,}")
|
| 353 |
+
report.append(f" Number-heavy lines removed: {total_number_lines_stripped:>10,}")
|
| 354 |
+
|
| 355 |
+
report.append(f"\n REMOVED DOCUMENT EXAMPLES (first 15)")
|
| 356 |
+
report.append(f" {'-'*50}")
|
| 357 |
+
for idx, reason, preview in removed_examples[:15]:
|
| 358 |
+
preview_clean = preview.replace('\n', ' ')[:200]
|
| 359 |
+
report.append(f"\n Doc #{idx} - {reason}")
|
| 360 |
+
report.append(f" \"{preview_clean}...\"")
|
| 361 |
+
|
| 362 |
+
report.append(f"\n VERDICT")
|
| 363 |
+
report.append(f" {'-'*50}")
|
| 364 |
+
report.append(f" Removed non-English fragments, very short docs, repetitive content,")
|
| 365 |
+
report.append(f" and heavy code. Stripped boilerplate from all remaining docs.")
|
| 366 |
+
report.append(f" Dataset is now optimized for English language learning!")
|
| 367 |
+
report.append(f"\n{'='*75}")
|
| 368 |
+
|
| 369 |
+
full_report = '\n'.join(report)
|
| 370 |
+
print(full_report)
|
| 371 |
+
|
| 372 |
+
report_path = output_dir / "SMART_CLEANUP_REPORT.txt"
|
| 373 |
+
with open(report_path, "w", encoding="utf-8") as f:
|
| 374 |
+
f.write(full_report)
|
| 375 |
+
print(f"\n Report saved to: {report_path}")
|
| 376 |
+
|
| 377 |
+
|
| 378 |
+
if __name__ == "__main__":
|
| 379 |
+
main()
|
Base/scripts/validate_model.py
ADDED
|
@@ -0,0 +1,137 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Post-training validation script.
|
| 3 |
+
Tests the trained model by generating text from prompts.
|
| 4 |
+
|
| 5 |
+
Usage:
|
| 6 |
+
python Base/scripts/validate_model.py --checkpoint_dir Base/out/pretrain/custom-100m-10m-test/final
|
| 7 |
+
python Base/scripts/validate_model.py --checkpoint_dir Base/out/pretrain/custom-100m-3b/final
|
| 8 |
+
"""
|
| 9 |
+
|
| 10 |
+
import argparse
|
| 11 |
+
from pathlib import Path
|
| 12 |
+
|
| 13 |
+
import torch
|
| 14 |
+
from litgpt import Tokenizer
|
| 15 |
+
from litgpt.config import Config
|
| 16 |
+
from litgpt.model import GPT
|
| 17 |
+
|
| 18 |
+
|
| 19 |
+
def load_model(checkpoint_dir: str):
|
| 20 |
+
"""Load a pretrained model from a litgpt checkpoint directory."""
|
| 21 |
+
checkpoint_dir = Path(checkpoint_dir)
|
| 22 |
+
|
| 23 |
+
if not checkpoint_dir.exists():
|
| 24 |
+
raise FileNotFoundError(f"Checkpoint directory not found: {checkpoint_dir}")
|
| 25 |
+
|
| 26 |
+
config = Config.from_checkpoint(checkpoint_dir)
|
| 27 |
+
model_path = checkpoint_dir / "lit_model.pth"
|
| 28 |
+
|
| 29 |
+
if not model_path.exists():
|
| 30 |
+
raise FileNotFoundError(f"Model file not found: {model_path}")
|
| 31 |
+
|
| 32 |
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| 33 |
+
print(f"Using device: {device}")
|
| 34 |
+
|
| 35 |
+
with torch.device("meta"):
|
| 36 |
+
model = GPT(config)
|
| 37 |
+
|
| 38 |
+
checkpoint = torch.load(str(model_path), map_location="cpu", weights_only=False)
|
| 39 |
+
|
| 40 |
+
# Handle both direct state_dict and wrapped checkpoint formats
|
| 41 |
+
if "model" in checkpoint:
|
| 42 |
+
state_dict = checkpoint["model"]
|
| 43 |
+
else:
|
| 44 |
+
state_dict = checkpoint
|
| 45 |
+
|
| 46 |
+
model.load_state_dict(state_dict, assign=True)
|
| 47 |
+
model = model.to(device)
|
| 48 |
+
model.eval()
|
| 49 |
+
|
| 50 |
+
return model, config, device
|
| 51 |
+
|
| 52 |
+
|
| 53 |
+
def generate_text(model, tokenizer, device, prompt, max_new_tokens=200, temperature=0.8, top_k=50):
|
| 54 |
+
"""Generate text from a prompt."""
|
| 55 |
+
input_ids = tokenizer.encode(prompt, device=device).unsqueeze(0)
|
| 56 |
+
|
| 57 |
+
with torch.no_grad():
|
| 58 |
+
for _ in range(max_new_tokens):
|
| 59 |
+
# Crop to block_size if needed
|
| 60 |
+
idx_cond = input_ids[:, -model.max_seq_length:]
|
| 61 |
+
logits = model(idx_cond)
|
| 62 |
+
logits = logits[:, -1, :] / temperature
|
| 63 |
+
|
| 64 |
+
# Top-k filtering
|
| 65 |
+
if top_k is not None:
|
| 66 |
+
v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
|
| 67 |
+
logits[logits < v[:, [-1]]] = float("-inf")
|
| 68 |
+
|
| 69 |
+
probs = torch.softmax(logits, dim=-1)
|
| 70 |
+
next_id = torch.multinomial(probs, num_samples=1)
|
| 71 |
+
input_ids = torch.cat([input_ids, next_id], dim=1)
|
| 72 |
+
|
| 73 |
+
# Stop on EOS
|
| 74 |
+
if next_id.item() == tokenizer.eos_id:
|
| 75 |
+
break
|
| 76 |
+
|
| 77 |
+
return tokenizer.decode(input_ids[0])
|
| 78 |
+
|
| 79 |
+
|
| 80 |
+
def main():
|
| 81 |
+
parser = argparse.ArgumentParser(description="Validate a pretrained LitGPT model")
|
| 82 |
+
parser.add_argument(
|
| 83 |
+
"--checkpoint_dir",
|
| 84 |
+
type=str,
|
| 85 |
+
required=True,
|
| 86 |
+
help="Path to the checkpoint directory (e.g., Base/out/pretrain/custom-100m-10m-test/final)",
|
| 87 |
+
)
|
| 88 |
+
parser.add_argument(
|
| 89 |
+
"--tokenizer_dir",
|
| 90 |
+
type=str,
|
| 91 |
+
default="Base/checkpoints/EleutherAI/pythia-160m",
|
| 92 |
+
help="Path to the tokenizer directory",
|
| 93 |
+
)
|
| 94 |
+
parser.add_argument(
|
| 95 |
+
"--max_new_tokens",
|
| 96 |
+
type=int,
|
| 97 |
+
default=200,
|
| 98 |
+
help="Maximum tokens to generate",
|
| 99 |
+
)
|
| 100 |
+
args = parser.parse_args()
|
| 101 |
+
|
| 102 |
+
print(f"Loading model from: {args.checkpoint_dir}")
|
| 103 |
+
model, config, device = load_model(args.checkpoint_dir)
|
| 104 |
+
|
| 105 |
+
print(f"\nModel config: {config.name}")
|
| 106 |
+
print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")
|
| 107 |
+
print(f"Block size: {config.block_size}")
|
| 108 |
+
|
| 109 |
+
tokenizer = Tokenizer(Path(args.tokenizer_dir))
|
| 110 |
+
|
| 111 |
+
test_prompts = [
|
| 112 |
+
"The future of artificial intelligence",
|
| 113 |
+
"In a groundbreaking study, researchers found that",
|
| 114 |
+
"The most important thing about education is",
|
| 115 |
+
"Once upon a time, in a land far away,",
|
| 116 |
+
]
|
| 117 |
+
|
| 118 |
+
print(f"\n{'='*60}")
|
| 119 |
+
print("GENERATION TESTS")
|
| 120 |
+
print(f"{'='*60}")
|
| 121 |
+
|
| 122 |
+
for i, prompt in enumerate(test_prompts, 1):
|
| 123 |
+
print(f"\n--- Prompt {i}: \"{prompt}\" ---")
|
| 124 |
+
output = generate_text(
|
| 125 |
+
model, tokenizer, device, prompt,
|
| 126 |
+
max_new_tokens=args.max_new_tokens,
|
| 127 |
+
)
|
| 128 |
+
print(f"Generated:\n{output}\n")
|
| 129 |
+
|
| 130 |
+
print(f"{'='*60}")
|
| 131 |
+
print("Validation complete!")
|
| 132 |
+
print("Note: A freshly pretrained model will produce semi-coherent text.")
|
| 133 |
+
print("Quality improves with more training data and compute.")
|
| 134 |
+
|
| 135 |
+
|
| 136 |
+
if __name__ == "__main__":
|
| 137 |
+
main()
|
README.md
CHANGED
|
@@ -1,2 +1,119 @@
|
|
| 1 |
-
# LUNA
|
| 2 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# LUNA - 100M Parameter LLM from Scratch
|
| 2 |
+
|
| 3 |
+
Custom ~100M parameter GPT model (Pythia-like architecture) pretrained on 4.5B tokens of clean English text.
|
| 4 |
+
|
| 5 |
+
## Quick Start (RunPod / Cloud GPU)
|
| 6 |
+
|
| 7 |
+
### 1. Clone & Install (one command)
|
| 8 |
+
|
| 9 |
+
```bash
|
| 10 |
+
git clone https://huggingface.co/spaces/ASTERIZER/LUNA /workspace/LUNA && \
|
| 11 |
+
cd /workspace/LUNA && \
|
| 12 |
+
pip install -q -r requirements.txt
|
| 13 |
+
```
|
| 14 |
+
|
| 15 |
+
### 2. Get Dataset + Train (one command)
|
| 16 |
+
|
| 17 |
+
**From HuggingFace (recommended):**
|
| 18 |
+
```bash
|
| 19 |
+
bash setup_and_train.sh huggingface ASTERIZER/LUNA-pretrain-data
|
| 20 |
+
```
|
| 21 |
+
|
| 22 |
+
**From Google Drive:**
|
| 23 |
+
```bash
|
| 24 |
+
bash setup_and_train.sh gdrive YOUR_GDRIVE_FOLDER_ID
|
| 25 |
+
```
|
| 26 |
+
|
| 27 |
+
**Smoke test (10M tokens only):**
|
| 28 |
+
```bash
|
| 29 |
+
bash setup_and_train.sh huggingface ASTERIZER/LUNA-pretrain-data 10000000
|
| 30 |
+
```
|
| 31 |
+
|
| 32 |
+
That's it. The script auto-detects your GPU, VRAM, RAM, CPU cores and configures everything for maximum utilization.
|
| 33 |
+
|
| 34 |
+
---
|
| 35 |
+
|
| 36 |
+
## How It Works
|
| 37 |
+
|
| 38 |
+
### Auto vs Manual Config
|
| 39 |
+
|
| 40 |
+
All hyperparameters live in `train_config.yaml`:
|
| 41 |
+
|
| 42 |
+
```yaml
|
| 43 |
+
auto_config: true # auto-detect everything from hardware
|
| 44 |
+
auto_config: false # use exact values below, no overrides
|
| 45 |
+
```
|
| 46 |
+
|
| 47 |
+
When `auto_config: true` (default), the trainer:
|
| 48 |
+
- **Probes VRAM** via binary search to find max micro_batch_size (82% safety)
|
| 49 |
+
- **Sets grad_accum** to hit the target global_batch_size
|
| 50 |
+
- **Picks precision** (bf16 on Ampere+, fp16 otherwise)
|
| 51 |
+
- **Scales workers** to half your CPU cores, capped by RAM
|
| 52 |
+
- **Enables torch.compile** if Triton is available (Linux)
|
| 53 |
+
|
| 54 |
+
When `auto_config: false`, every value in the YAML is used exactly as-is.
|
| 55 |
+
|
| 56 |
+
### CLI Overrides
|
| 57 |
+
|
| 58 |
+
Any config value can be overridden from the command line:
|
| 59 |
+
|
| 60 |
+
```bash
|
| 61 |
+
python train.py --config train_config.yaml --data_path /data/litdata --max_tokens 100000000
|
| 62 |
+
```
|
| 63 |
+
|
| 64 |
+
Priority: CLI args > train_config.yaml > auto-detection
|
| 65 |
+
|
| 66 |
+
---
|
| 67 |
+
|
| 68 |
+
## Dataset
|
| 69 |
+
|
| 70 |
+
- **4,515,286,950 tokens** (4.5B) in 270 binary chunks
|
| 71 |
+
- Sources: Wikipedia, FineWeb-Edu, OpenWebText (deduplicated, cleaned)
|
| 72 |
+
- Format: LitData binary (int32, block_size=1025, TokensLoader)
|
| 73 |
+
- Tokenizer: EleutherAI/pythia-160m (50,254 vocab)
|
| 74 |
+
|
| 75 |
+
## Model Architecture
|
| 76 |
+
|
| 77 |
+
| Parameter | Value |
|
| 78 |
+
|-----------|-------|
|
| 79 |
+
| Layers | 10 |
|
| 80 |
+
| Hidden dim | 768 |
|
| 81 |
+
| Attention heads | 12 |
|
| 82 |
+
| Vocab size | 50,304 (padded) |
|
| 83 |
+
| Context length | 1,024 |
|
| 84 |
+
| Total params | ~109M (70M unique, tied embeddings) |
|
| 85 |
+
| Rotary % | 25% |
|
| 86 |
+
|
| 87 |
+
## File Structure
|
| 88 |
+
|
| 89 |
+
```
|
| 90 |
+
LUNA/
|
| 91 |
+
train.py # Main training script (config-driven, auto-detects hardware)
|
| 92 |
+
train_config.yaml # All hyperparameters (auto_config: true/false)
|
| 93 |
+
fetch_data.py # Downloads dataset from HuggingFace / GDrive
|
| 94 |
+
setup_and_train.sh # One-command cloud entrypoint
|
| 95 |
+
benchmark_runpod.py # Local benchmark + RunPod cost calculator
|
| 96 |
+
requirements.txt # Python dependencies
|
| 97 |
+
Base/
|
| 98 |
+
checkpoints/EleutherAI/pythia-160m/ # Tokenizer files
|
| 99 |
+
configs/ # Legacy litgpt YAML configs (reference only)
|
| 100 |
+
scripts/ # Data preprocessing scripts
|
| 101 |
+
```
|
| 102 |
+
|
| 103 |
+
## Estimated Training Times (RunPod)
|
| 104 |
+
|
| 105 |
+
| GPU | $/hr | tok/s | Hours | Cost USD | Cost INR |
|
| 106 |
+
|-----|------|-------|-------|----------|----------|
|
| 107 |
+
| RTX A5000 | $0.16 | ~6,400 | ~196h | ~$31 | ~2,700 |
|
| 108 |
+
| RTX 3090 | $0.22 | ~7,600 | ~165h | ~$36 | ~3,100 |
|
| 109 |
+
| RTX 4090 | $0.34 | ~10,000 | ~125h | ~$42 | ~3,600 |
|
| 110 |
+
| RTX 5090 | $0.69 | ~16,000 | ~78h | ~$54 | ~4,600 |
|
| 111 |
+
| H100 NVL | $2.59 | ~43,000 | ~29h | ~$75 | ~6,400 |
|
| 112 |
+
|
| 113 |
+
## Resume Training
|
| 114 |
+
|
| 115 |
+
Training auto-saves `latest.pt` every save_interval steps. If interrupted, just re-run the same command -- it picks up where it left off.
|
| 116 |
+
|
| 117 |
+
## License
|
| 118 |
+
|
| 119 |
+
Private / ASTERIZER 2026
|
benchmark_runpod.py
ADDED
|
@@ -0,0 +1,412 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
LUNA 100M - Local Benchmark + RunPod Cost Calculator
|
| 3 |
+
=====================================================
|
| 4 |
+
Uses PyTorch SDPA (Flash Attention) for realistic training throughput.
|
| 5 |
+
Matches the exact LUNA model architecture and training config.
|
| 6 |
+
"""
|
| 7 |
+
|
| 8 |
+
import os
|
| 9 |
+
import sys
|
| 10 |
+
import time
|
| 11 |
+
import math
|
| 12 |
+
import json
|
| 13 |
+
import gc
|
| 14 |
+
import torch
|
| 15 |
+
import torch.nn as nn
|
| 16 |
+
import torch.nn.functional as F
|
| 17 |
+
from torch.amp import autocast, GradScaler
|
| 18 |
+
|
| 19 |
+
# βββ Model Architecture (matches your config exactly) βββββββββββββββββββββββββ
|
| 20 |
+
|
| 21 |
+
class RotaryEmbedding(nn.Module):
|
| 22 |
+
def __init__(self, dim, max_seq_len=1024):
|
| 23 |
+
super().__init__()
|
| 24 |
+
inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
|
| 25 |
+
self.register_buffer("inv_freq", inv_freq)
|
| 26 |
+
t = torch.arange(max_seq_len).float()
|
| 27 |
+
freqs = torch.einsum("i,j->ij", t, inv_freq)
|
| 28 |
+
emb = torch.cat([freqs, freqs], dim=-1)
|
| 29 |
+
self.register_buffer("cos_cached", emb.cos())
|
| 30 |
+
self.register_buffer("sin_cached", emb.sin())
|
| 31 |
+
|
| 32 |
+
def forward(self, seq_len):
|
| 33 |
+
return self.cos_cached[:seq_len], self.sin_cached[:seq_len]
|
| 34 |
+
|
| 35 |
+
def rotate_half(x):
|
| 36 |
+
x1, x2 = x.chunk(2, dim=-1)
|
| 37 |
+
return torch.cat([-x2, x1], dim=-1)
|
| 38 |
+
|
| 39 |
+
def apply_rotary(x, cos, sin):
|
| 40 |
+
cos = cos.unsqueeze(0).unsqueeze(0)
|
| 41 |
+
sin = sin.unsqueeze(0).unsqueeze(0)
|
| 42 |
+
return x * cos + rotate_half(x) * sin
|
| 43 |
+
|
| 44 |
+
class CausalSelfAttention(nn.Module):
|
| 45 |
+
def __init__(self, n_embd, n_head, block_size, rotary_pct=0.25):
|
| 46 |
+
super().__init__()
|
| 47 |
+
self.n_head = n_head
|
| 48 |
+
self.head_dim = n_embd // n_head
|
| 49 |
+
self.rotary_dim = int(self.head_dim * rotary_pct)
|
| 50 |
+
self.c_attn = nn.Linear(n_embd, 3 * n_embd, bias=True)
|
| 51 |
+
self.c_proj = nn.Linear(n_embd, n_embd, bias=True)
|
| 52 |
+
self.rotary = RotaryEmbedding(self.rotary_dim, block_size)
|
| 53 |
+
|
| 54 |
+
def forward(self, x):
|
| 55 |
+
B, T, C = x.size()
|
| 56 |
+
qkv = self.c_attn(x).reshape(B, T, 3, self.n_head, self.head_dim).permute(2, 0, 3, 1, 4)
|
| 57 |
+
q, k, v = qkv.unbind(0)
|
| 58 |
+
|
| 59 |
+
cos, sin = self.rotary(T)
|
| 60 |
+
q_rot = apply_rotary(q[..., :self.rotary_dim], cos, sin)
|
| 61 |
+
k_rot = apply_rotary(k[..., :self.rotary_dim], cos, sin)
|
| 62 |
+
q = torch.cat([q_rot, q[..., self.rotary_dim:]], dim=-1)
|
| 63 |
+
k = torch.cat([k_rot, k[..., self.rotary_dim:]], dim=-1)
|
| 64 |
+
|
| 65 |
+
# SDPA = Flash Attention / Memory Efficient Attention
|
| 66 |
+
y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
|
| 67 |
+
y = y.transpose(1, 2).contiguous().view(B, T, C)
|
| 68 |
+
return self.c_proj(y)
|
| 69 |
+
|
| 70 |
+
class MLP(nn.Module):
|
| 71 |
+
def __init__(self, n_embd):
|
| 72 |
+
super().__init__()
|
| 73 |
+
self.c_fc = nn.Linear(n_embd, 4 * n_embd, bias=True)
|
| 74 |
+
self.gelu = nn.GELU()
|
| 75 |
+
self.c_proj = nn.Linear(4 * n_embd, n_embd, bias=True)
|
| 76 |
+
def forward(self, x):
|
| 77 |
+
return self.c_proj(self.gelu(self.c_fc(x)))
|
| 78 |
+
|
| 79 |
+
class Block(nn.Module):
|
| 80 |
+
def __init__(self, n_embd, n_head, block_size):
|
| 81 |
+
super().__init__()
|
| 82 |
+
self.ln_1 = nn.LayerNorm(n_embd)
|
| 83 |
+
self.attn = CausalSelfAttention(n_embd, n_head, block_size)
|
| 84 |
+
self.ln_2 = nn.LayerNorm(n_embd)
|
| 85 |
+
self.mlp = MLP(n_embd)
|
| 86 |
+
def forward(self, x):
|
| 87 |
+
x = x + self.attn(self.ln_1(x))
|
| 88 |
+
x = x + self.mlp(self.ln_2(x))
|
| 89 |
+
return x
|
| 90 |
+
|
| 91 |
+
class LUNAModel(nn.Module):
|
| 92 |
+
def __init__(self, vocab_size=50254, block_size=1024, n_layer=10, n_embd=768, n_head=12):
|
| 93 |
+
super().__init__()
|
| 94 |
+
self.block_size = block_size
|
| 95 |
+
self.wte = nn.Embedding(vocab_size, n_embd)
|
| 96 |
+
self.blocks = nn.ModuleList([Block(n_embd, n_head, block_size) for _ in range(n_layer)])
|
| 97 |
+
self.ln_f = nn.LayerNorm(n_embd)
|
| 98 |
+
self.lm_head = nn.Linear(n_embd, vocab_size, bias=False)
|
| 99 |
+
self.lm_head.weight = self.wte.weight
|
| 100 |
+
self.apply(self._init_weights)
|
| 101 |
+
|
| 102 |
+
def _init_weights(self, module):
|
| 103 |
+
if isinstance(module, (nn.Linear, nn.Embedding)):
|
| 104 |
+
module.weight.data.normal_(mean=0.0, std=0.02)
|
| 105 |
+
if isinstance(module, nn.Linear) and module.bias is not None:
|
| 106 |
+
module.bias.data.zero_()
|
| 107 |
+
|
| 108 |
+
def forward(self, idx, targets=None):
|
| 109 |
+
B, T = idx.size()
|
| 110 |
+
x = self.wte(idx)
|
| 111 |
+
for block in self.blocks:
|
| 112 |
+
x = block(x)
|
| 113 |
+
x = self.ln_f(x)
|
| 114 |
+
logits = self.lm_head(x)
|
| 115 |
+
loss = None
|
| 116 |
+
if targets is not None:
|
| 117 |
+
loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
|
| 118 |
+
return logits, loss
|
| 119 |
+
|
| 120 |
+
# βββ Config ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 121 |
+
|
| 122 |
+
DATASET_TOTAL_TOKENS = 4_515_286_950 # Verified from index.json: 270 chunks, sum of all dims
|
| 123 |
+
BLOCK_SIZE = 1024
|
| 124 |
+
VOCAB_SIZE = 50254
|
| 125 |
+
N_LAYER = 10
|
| 126 |
+
N_EMBD = 768
|
| 127 |
+
N_HEAD = 12
|
| 128 |
+
GLOBAL_BATCH_SIZE = 120
|
| 129 |
+
MAX_SEQ_LENGTH = 1024
|
| 130 |
+
|
| 131 |
+
WARMUP_STEPS = 3
|
| 132 |
+
BENCHMARK_STEPS = 4
|
| 133 |
+
|
| 134 |
+
# RunPod GPUs (Community Cloud pricing April 2026)
|
| 135 |
+
RUNPOD_GPUS = [
|
| 136 |
+
# (name, $/hr, VRAM_GB, bf16_TF_nonsparse, mem_bw_GBs, arch)
|
| 137 |
+
("RTX A5000", 0.16, 24, 65, 768, "Ampere"),
|
| 138 |
+
("RTX 3090", 0.22, 24, 71, 936, "Ampere"),
|
| 139 |
+
("RTX A6000", 0.33, 48, 77, 768, "Ampere"),
|
| 140 |
+
("RTX 4090", 0.34, 24, 165, 1008, "Ada"),
|
| 141 |
+
("A40", 0.35, 48, 75, 696, "Ampere"),
|
| 142 |
+
("L4", 0.44, 24, 121, 300, "Ada"),
|
| 143 |
+
("RTX 5090", 0.69, 32, 210, 1792, "Blackwell"),
|
| 144 |
+
("L40", 0.69, 48, 181, 864, "Ada"),
|
| 145 |
+
("RTX 6000 Ada", 0.74, 48, 181, 960, "Ada"),
|
| 146 |
+
("L40S", 0.79, 48, 183, 864, "Ada"),
|
| 147 |
+
("A100 PCIe 80GB", 1.19, 80, 312, 2039, "Ampere"),
|
| 148 |
+
("A100 SXM 80GB", 1.39, 80, 312, 2039, "Ampere"),
|
| 149 |
+
("RTX Pro 6000", 1.69, 96, 260, 1280, "Blackwell"),
|
| 150 |
+
("H100 PCIe", 1.99, 80, 756, 2039, "Hopper"),
|
| 151 |
+
("H100 NVL", 2.59, 94, 835, 3938, "Hopper"),
|
| 152 |
+
("H100 SXM", 2.69, 80, 990, 3352, "Hopper"),
|
| 153 |
+
("H200", 3.59,141, 990, 4800, "Hopper"),
|
| 154 |
+
]
|
| 155 |
+
|
| 156 |
+
USD_TO_INR = 86.0
|
| 157 |
+
|
| 158 |
+
|
| 159 |
+
def find_max_micro_batch(model, device, seq_len=1024, start=32):
|
| 160 |
+
"""Binary search for max micro_batch_size, with 0.65 safety factor."""
|
| 161 |
+
model.train()
|
| 162 |
+
lo, hi, best = 1, start, 1
|
| 163 |
+
opt_tmp = torch.optim.AdamW(model.parameters(), lr=1e-4)
|
| 164 |
+
|
| 165 |
+
while hi >= lo:
|
| 166 |
+
mid = (lo + hi) // 2
|
| 167 |
+
try:
|
| 168 |
+
torch.cuda.empty_cache()
|
| 169 |
+
torch.cuda.reset_peak_memory_stats()
|
| 170 |
+
opt_tmp.zero_grad(set_to_none=True)
|
| 171 |
+
x = torch.randint(0, VOCAB_SIZE, (mid, seq_len), device=device)
|
| 172 |
+
t = torch.randint(0, VOCAB_SIZE, (mid, seq_len), device=device)
|
| 173 |
+
with autocast(device_type='cuda', dtype=torch.bfloat16):
|
| 174 |
+
_, loss = model(x, t)
|
| 175 |
+
loss.backward()
|
| 176 |
+
opt_tmp.step()
|
| 177 |
+
opt_tmp.zero_grad(set_to_none=True)
|
| 178 |
+
best = mid
|
| 179 |
+
lo = mid + 1
|
| 180 |
+
del x, t, loss
|
| 181 |
+
torch.cuda.empty_cache()
|
| 182 |
+
except (torch.cuda.OutOfMemoryError, RuntimeError):
|
| 183 |
+
try:
|
| 184 |
+
del x, t, loss
|
| 185 |
+
except:
|
| 186 |
+
pass
|
| 187 |
+
torch.cuda.empty_cache()
|
| 188 |
+
opt_tmp.zero_grad(set_to_none=True)
|
| 189 |
+
hi = mid - 1
|
| 190 |
+
|
| 191 |
+
safe = max(1, int(best * 0.65))
|
| 192 |
+
del opt_tmp
|
| 193 |
+
torch.cuda.empty_cache()
|
| 194 |
+
gc.collect()
|
| 195 |
+
return safe
|
| 196 |
+
|
| 197 |
+
|
| 198 |
+
def run_benchmark():
|
| 199 |
+
device = torch.device("cuda")
|
| 200 |
+
torch.backends.cuda.matmul.allow_tf32 = True
|
| 201 |
+
torch.backends.cudnn.allow_tf32 = True
|
| 202 |
+
|
| 203 |
+
print("=" * 72)
|
| 204 |
+
print(" LUNA 100M - TRAINING BENCHMARK & RUNPOD COST CALCULATOR")
|
| 205 |
+
print("=" * 72)
|
| 206 |
+
|
| 207 |
+
gpu_name = torch.cuda.get_device_name(0)
|
| 208 |
+
gpu_mem = torch.cuda.get_device_properties(0).total_memory / 1024**3
|
| 209 |
+
print(f"\n Local GPU: {gpu_name}")
|
| 210 |
+
print(f" VRAM: {gpu_mem:.1f} GB")
|
| 211 |
+
print(f" PyTorch: {torch.__version__}, CUDA: {torch.version.cuda}")
|
| 212 |
+
|
| 213 |
+
print(f"\n Creating LUNA-100M (SDPA/Flash Attention)...")
|
| 214 |
+
model = LUNAModel(VOCAB_SIZE, BLOCK_SIZE, N_LAYER, N_EMBD, N_HEAD).to(device)
|
| 215 |
+
|
| 216 |
+
total_params = sum(p.numel() for p in model.parameters())
|
| 217 |
+
# Tied embeddings: wte(50254*768) = 38,595,072 counted once
|
| 218 |
+
unique_params = total_params - model.wte.weight.numel()
|
| 219 |
+
print(f" Parameters: {total_params:,} total, {unique_params:,} unique")
|
| 220 |
+
|
| 221 |
+
print(f"\n Probing max micro_batch_size...")
|
| 222 |
+
max_mbs = find_max_micro_batch(model, device, MAX_SEQ_LENGTH, start=40)
|
| 223 |
+
print(f" Safe micro_batch_size: {max_mbs}")
|
| 224 |
+
|
| 225 |
+
grad_accum = max(1, GLOBAL_BATCH_SIZE // max_mbs)
|
| 226 |
+
effective_batch = max_mbs * grad_accum
|
| 227 |
+
tokens_per_step = effective_batch * MAX_SEQ_LENGTH
|
| 228 |
+
print(f" grad_accum={grad_accum}, effective_batch={effective_batch}")
|
| 229 |
+
print(f" Tokens/step: {tokens_per_step:,}")
|
| 230 |
+
|
| 231 |
+
optimizer = torch.optim.AdamW(
|
| 232 |
+
model.parameters(), lr=6e-4, weight_decay=0.1,
|
| 233 |
+
betas=(0.9, 0.95), eps=1e-8
|
| 234 |
+
)
|
| 235 |
+
scaler = GradScaler()
|
| 236 |
+
|
| 237 |
+
total_steps = WARMUP_STEPS + BENCHMARK_STEPS
|
| 238 |
+
print(f"\n Running {WARMUP_STEPS} warmup + {BENCHMARK_STEPS} benchmark steps...")
|
| 239 |
+
|
| 240 |
+
model.train()
|
| 241 |
+
step_times = []
|
| 242 |
+
torch.cuda.synchronize()
|
| 243 |
+
torch.cuda.reset_peak_memory_stats()
|
| 244 |
+
|
| 245 |
+
for step in range(total_steps):
|
| 246 |
+
t0 = time.perf_counter()
|
| 247 |
+
optimizer.zero_grad(set_to_none=True)
|
| 248 |
+
step_loss = 0.0
|
| 249 |
+
|
| 250 |
+
for _ in range(grad_accum):
|
| 251 |
+
x = torch.randint(0, VOCAB_SIZE, (max_mbs, MAX_SEQ_LENGTH), device=device)
|
| 252 |
+
tgt = torch.randint(0, VOCAB_SIZE, (max_mbs, MAX_SEQ_LENGTH), device=device)
|
| 253 |
+
with autocast(device_type='cuda', dtype=torch.bfloat16):
|
| 254 |
+
_, loss = model(x, tgt)
|
| 255 |
+
loss = loss / grad_accum
|
| 256 |
+
scaler.scale(loss).backward()
|
| 257 |
+
step_loss += loss.item()
|
| 258 |
+
del x, tgt, loss
|
| 259 |
+
|
| 260 |
+
scaler.unscale_(optimizer)
|
| 261 |
+
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
|
| 262 |
+
scaler.step(optimizer)
|
| 263 |
+
scaler.update()
|
| 264 |
+
torch.cuda.synchronize()
|
| 265 |
+
dt = time.perf_counter() - t0
|
| 266 |
+
|
| 267 |
+
if step >= WARMUP_STEPS:
|
| 268 |
+
step_times.append(dt)
|
| 269 |
+
|
| 270 |
+
tps = tokens_per_step / dt
|
| 271 |
+
phase = "WARM" if step < WARMUP_STEPS else "BENCH"
|
| 272 |
+
print(f" [{phase}] Step {step:3d} | Loss {step_loss:.4f} | {dt:.2f}s | {tps:,.0f} tok/s")
|
| 273 |
+
|
| 274 |
+
# βββ Results ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 275 |
+
peak_vram_gb = torch.cuda.max_memory_allocated() / 1024**3
|
| 276 |
+
|
| 277 |
+
avg_time = sum(step_times) / len(step_times)
|
| 278 |
+
med_time = sorted(step_times)[len(step_times) // 2]
|
| 279 |
+
avg_tps = tokens_per_step / avg_time
|
| 280 |
+
med_tps = tokens_per_step / med_time
|
| 281 |
+
peak_tps = tokens_per_step / min(step_times)
|
| 282 |
+
|
| 283 |
+
flops_per_token = 6 * unique_params
|
| 284 |
+
achieved_tf = (avg_tps * flops_per_token) / 1e12
|
| 285 |
+
LOCAL_TF = 88.0 # RTX 4060 Ti BF16 non-sparse
|
| 286 |
+
LOCAL_BW = 288.0 # GB/s
|
| 287 |
+
mfu = achieved_tf / LOCAL_TF
|
| 288 |
+
|
| 289 |
+
print("\n" + "=" * 72)
|
| 290 |
+
print(" LOCAL BENCHMARK RESULTS")
|
| 291 |
+
print("=" * 72)
|
| 292 |
+
print(f" GPU: {gpu_name} ({gpu_mem:.1f} GB)")
|
| 293 |
+
print(f" Peak VRAM: {peak_vram_gb:.2f} GB ({peak_vram_gb/gpu_mem*100:.0f}%)")
|
| 294 |
+
print(f" Batch: micro={max_mbs}, accum={grad_accum}, global={effective_batch}")
|
| 295 |
+
print(f" Step time: avg={avg_time:.3f}s, median={med_time:.3f}s")
|
| 296 |
+
print(f" Tokens/sec: avg={avg_tps:,.0f}, median={med_tps:,.0f}, peak={peak_tps:,.0f}")
|
| 297 |
+
print(f" TFLOPS: {achieved_tf:.2f}, MFU: {mfu*100:.1f}%")
|
| 298 |
+
|
| 299 |
+
n_steps = math.ceil(DATASET_TOTAL_TOKENS / tokens_per_step)
|
| 300 |
+
local_hrs = (n_steps * avg_time) / 3600
|
| 301 |
+
|
| 302 |
+
print(f" Dataset: 4,515,286,950 tokens | Steps needed: {n_steps:,}")
|
| 303 |
+
print(f" Local training: {local_hrs:.1f} hrs ({local_hrs/24:.1f} days)")
|
| 304 |
+
|
| 305 |
+
# βββ RunPod Estimates βββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 306 |
+
print("\n" + "=" * 72)
|
| 307 |
+
print(" RUNPOD GPU COMPARISON (Community Cloud, INR/86/USD)")
|
| 308 |
+
print("=" * 72)
|
| 309 |
+
|
| 310 |
+
results = []
|
| 311 |
+
for name, price, vram, bf16, bw, arch in RUNPOD_GPUS:
|
| 312 |
+
# Memory estimation for each GPU
|
| 313 |
+
fixed_gb = (unique_params * (2 + 8 + 2)) / 1024**3 + 0.5
|
| 314 |
+
avail_gb = vram - fixed_gb
|
| 315 |
+
act_per_sample = max(0.05, (peak_vram_gb - fixed_gb) / max(max_mbs, 1))
|
| 316 |
+
est_mbs = max(1, min(128, int(avail_gb / act_per_sample)))
|
| 317 |
+
est_ga = max(1, GLOBAL_BATCH_SIZE // est_mbs)
|
| 318 |
+
est_tps_step = est_mbs * est_ga * MAX_SEQ_LENGTH
|
| 319 |
+
|
| 320 |
+
# Scaling: 50% compute + 50% bandwidth (validated for small transformers)
|
| 321 |
+
speedup = 0.50 * (bf16 / LOCAL_TF) + 0.50 * (bw / LOCAL_BW)
|
| 322 |
+
est_tps = avg_tps * speedup * 0.90 # 0.90 cloud overhead
|
| 323 |
+
|
| 324 |
+
est_steps = math.ceil(DATASET_TOTAL_TOKENS / est_tps_step)
|
| 325 |
+
est_sec = (est_tps_step / est_tps) * est_steps
|
| 326 |
+
est_hrs = est_sec / 3600
|
| 327 |
+
cost_usd = est_hrs * price
|
| 328 |
+
cost_inr = cost_usd * USD_TO_INR
|
| 329 |
+
|
| 330 |
+
results.append({
|
| 331 |
+
"gpu": name, "price": price, "vram": vram, "bf16": bf16,
|
| 332 |
+
"bw": bw, "arch": arch, "mbs": est_mbs, "ga": est_ga,
|
| 333 |
+
"tps": round(est_tps), "hours": round(est_hrs, 1),
|
| 334 |
+
"usd": round(cost_usd, 2), "inr": round(cost_inr),
|
| 335 |
+
})
|
| 336 |
+
|
| 337 |
+
results.sort(key=lambda r: r["inr"])
|
| 338 |
+
|
| 339 |
+
print(f"\n {'#':<3} {'GPU':<18} {'$/hr':>5} {'VRAM':>5} {'tok/s':>10} "
|
| 340 |
+
f"{'Hours':>7} {'$ USD':>8} {'INR':>10}")
|
| 341 |
+
print(" " + "β" * 72)
|
| 342 |
+
for i, r in enumerate(results):
|
| 343 |
+
s = " *" if i < 3 else ""
|
| 344 |
+
print(f" {i+1:<3} {r['gpu']:<18} {r['price']:>5.2f} {r['vram']:>4}G "
|
| 345 |
+
f"{r['tps']:>10,} {r['hours']:>7.1f} {r['usd']:>8.2f} {r['inr']:>10,}{s}")
|
| 346 |
+
|
| 347 |
+
# Top 5 detailed
|
| 348 |
+
print("\n" + "=" * 72)
|
| 349 |
+
print(" TOP 5 CHEAPEST - DETAILS")
|
| 350 |
+
print("=" * 72)
|
| 351 |
+
for i, r in enumerate(results[:5]):
|
| 352 |
+
sx = r["tps"] / avg_tps if avg_tps > 0 else 0
|
| 353 |
+
print(f"\n #{i+1}: {r['gpu']} ({r['arch']})")
|
| 354 |
+
print(f" βββ ${r['price']:.2f}/hr | {r['vram']}GB VRAM | {r['bf16']} TF | {r['bw']} GB/s")
|
| 355 |
+
print(f" βββ micro_batch: {r['mbs']}, grad_accum: {r['ga']}")
|
| 356 |
+
print(f" βββ {r['tps']:,} tok/s ({sx:.2f}Γ local)")
|
| 357 |
+
print(f" βββ {r['hours']:.1f} hrs ({r['hours']/24:.1f} days)")
|
| 358 |
+
print(f" +-- ${r['usd']:.2f} = INR {r['inr']:,}")
|
| 359 |
+
|
| 360 |
+
# Local reference
|
| 361 |
+
print("\n" + "=" * 72)
|
| 362 |
+
print(" YOUR LOCAL GPU")
|
| 363 |
+
print("=" * 72)
|
| 364 |
+
print(f" RTX 4060 Ti 16GB: {avg_tps:,.0f} tok/s")
|
| 365 |
+
print(f" Training: {local_hrs:.1f} hrs ({local_hrs/24:.1f} days)")
|
| 366 |
+
print(f" Electricity: ~INR {local_hrs * 0.16 * 8:,.0f} (160W x Rs8/kWh)")
|
| 367 |
+
|
| 368 |
+
# Recommendation
|
| 369 |
+
best = results[0]
|
| 370 |
+
print("\n" + "=" * 72)
|
| 371 |
+
print(" * RECOMMENDATION")
|
| 372 |
+
print("=" * 72)
|
| 373 |
+
print(f" Most affordable: {best['gpu']} @ ${best['price']:.2f}/hr")
|
| 374 |
+
print(f" Time: {best['hours']:.1f} hrs ({best['hours']/24:.1f} days)")
|
| 375 |
+
print(f" Cost: INR {best['inr']:,} (${best['usd']:.2f})")
|
| 376 |
+
print(f" Speed: {best['tps']/avg_tps:.1f}Γ local" if avg_tps > 0 else "")
|
| 377 |
+
|
| 378 |
+
fast = [r for r in results if r["hours"] < max(8, local_hrs * 0.15)]
|
| 379 |
+
if fast:
|
| 380 |
+
fb = min(fast, key=lambda r: r["inr"])
|
| 381 |
+
if fb["gpu"] != best["gpu"]:
|
| 382 |
+
print(f"\n Fastest affordable: {fb['gpu']} @ ${fb['price']:.2f}/hr")
|
| 383 |
+
print(f" Time: {fb['hours']:.1f} hrs | Cost: INR {fb['inr']:,}")
|
| 384 |
+
|
| 385 |
+
print("\n" + "=" * 72)
|
| 386 |
+
|
| 387 |
+
# Save JSON
|
| 388 |
+
out = os.path.join(os.path.dirname(os.path.abspath(__file__)), "runpod_cost_estimate.json")
|
| 389 |
+
with open(out, "w") as f:
|
| 390 |
+
json.dump({
|
| 391 |
+
"benchmark": {
|
| 392 |
+
"gpu": gpu_name, "vram_gb": round(gpu_mem, 1),
|
| 393 |
+
"peak_vram_gb": round(peak_vram_gb, 2),
|
| 394 |
+
"micro_batch": max_mbs, "grad_accum": grad_accum,
|
| 395 |
+
"tokens_per_step": tokens_per_step,
|
| 396 |
+
"avg_tok_per_sec": round(avg_tps),
|
| 397 |
+
"median_tok_per_sec": round(med_tps),
|
| 398 |
+
"achieved_tflops": round(achieved_tf, 2),
|
| 399 |
+
"mfu_pct": round(mfu*100, 1),
|
| 400 |
+
},
|
| 401 |
+
"dataset": {"tokens": DATASET_TOTAL_TOKENS, "chunks": 270},
|
| 402 |
+
"model": {"total_params": total_params, "unique_params": unique_params},
|
| 403 |
+
"local_hours": round(local_hrs, 1),
|
| 404 |
+
"runpod": results,
|
| 405 |
+
"usd_to_inr": USD_TO_INR,
|
| 406 |
+
}, f, indent=2)
|
| 407 |
+
print(f" Saved: {out}")
|
| 408 |
+
print("=" * 72)
|
| 409 |
+
|
| 410 |
+
|
| 411 |
+
if __name__ == "__main__":
|
| 412 |
+
run_benchmark()
|
fetch_data.py
ADDED
|
@@ -0,0 +1,162 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
LUNA β Dataset Fetcher
|
| 3 |
+
======================
|
| 4 |
+
Downloads the tokenized litdata dataset from either:
|
| 5 |
+
- HuggingFace Hub (recommended, free, fast)
|
| 6 |
+
- Google Drive (direct link, requires gdown)
|
| 7 |
+
|
| 8 |
+
Usage:
|
| 9 |
+
python fetch_data.py --source huggingface --hf_repo YourName/LUNA-pretrain-data --out_dir /workspace/data
|
| 10 |
+
python fetch_data.py --source gdrive --gdrive_id <FILE_OR_FOLDER_ID> --out_dir /workspace/data
|
| 11 |
+
python fetch_data.py --source local --local_path Base/data/litdata_pretrain_final --out_dir /workspace/data
|
| 12 |
+
|
| 13 |
+
After running, pass --data_path /workspace/data/litdata_pretrain_final to train.py
|
| 14 |
+
"""
|
| 15 |
+
|
| 16 |
+
import os
|
| 17 |
+
import sys
|
| 18 |
+
import json
|
| 19 |
+
import shutil
|
| 20 |
+
import argparse
|
| 21 |
+
from pathlib import Path
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
# βββ HuggingFace Download βββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 25 |
+
|
| 26 |
+
def download_huggingface(repo_id: str, out_dir: Path, hf_token: str = None):
|
| 27 |
+
try:
|
| 28 |
+
from huggingface_hub import snapshot_download, hf_hub_download
|
| 29 |
+
except ImportError:
|
| 30 |
+
print(" Installing huggingface_hub...")
|
| 31 |
+
os.system(f"{sys.executable} -m pip install -q huggingface_hub")
|
| 32 |
+
from huggingface_hub import snapshot_download
|
| 33 |
+
|
| 34 |
+
print(f" Downloading from HuggingFace: {repo_id}")
|
| 35 |
+
out_dir.mkdir(parents=True, exist_ok=True)
|
| 36 |
+
snapshot_download(
|
| 37 |
+
repo_id=repo_id,
|
| 38 |
+
repo_type="dataset",
|
| 39 |
+
local_dir=str(out_dir),
|
| 40 |
+
token=hf_token,
|
| 41 |
+
ignore_patterns=["*.md", ".gitattributes"],
|
| 42 |
+
)
|
| 43 |
+
print(f" Downloaded to: {out_dir}")
|
| 44 |
+
_verify(out_dir)
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
# βββ Google Drive Download ββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 48 |
+
|
| 49 |
+
def download_gdrive(gdrive_id: str, out_dir: Path):
|
| 50 |
+
try:
|
| 51 |
+
import gdown
|
| 52 |
+
except ImportError:
|
| 53 |
+
print(" Installing gdown...")
|
| 54 |
+
os.system(f"{sys.executable} -m pip install -q gdown")
|
| 55 |
+
import gdown
|
| 56 |
+
|
| 57 |
+
out_dir.mkdir(parents=True, exist_ok=True)
|
| 58 |
+
# Try as folder first, then single file
|
| 59 |
+
url = f"https://drive.google.com/drive/folders/{gdrive_id}"
|
| 60 |
+
print(f" Attempting GDrive folder download: {gdrive_id}")
|
| 61 |
+
try:
|
| 62 |
+
gdown.download_folder(url=url, output=str(out_dir), quiet=False, use_cookies=False)
|
| 63 |
+
except Exception as e:
|
| 64 |
+
print(f" Folder download failed ({e}), trying single file...")
|
| 65 |
+
url = f"https://drive.google.com/uc?id={gdrive_id}"
|
| 66 |
+
dest = out_dir / "data.zip"
|
| 67 |
+
gdown.download(url, str(dest), quiet=False)
|
| 68 |
+
if dest.suffix == ".zip":
|
| 69 |
+
print(" Extracting zip...")
|
| 70 |
+
import zipfile
|
| 71 |
+
with zipfile.ZipFile(dest) as z:
|
| 72 |
+
z.extractall(out_dir)
|
| 73 |
+
dest.unlink()
|
| 74 |
+
print(f" Downloaded to: {out_dir}")
|
| 75 |
+
_verify(out_dir)
|
| 76 |
+
|
| 77 |
+
|
| 78 |
+
# βββ Local Copy βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 79 |
+
|
| 80 |
+
def copy_local(local_path: str, out_dir: Path):
|
| 81 |
+
src = Path(local_path)
|
| 82 |
+
if not src.exists():
|
| 83 |
+
raise FileNotFoundError(f"Local path not found: {src}")
|
| 84 |
+
if out_dir.resolve() == src.resolve():
|
| 85 |
+
print(f" Source == destination, no copy needed.")
|
| 86 |
+
_verify(out_dir)
|
| 87 |
+
return
|
| 88 |
+
print(f" Copying {src} β {out_dir}")
|
| 89 |
+
if out_dir.exists():
|
| 90 |
+
shutil.rmtree(out_dir)
|
| 91 |
+
shutil.copytree(src, out_dir)
|
| 92 |
+
print(f" Copied to: {out_dir}")
|
| 93 |
+
_verify(out_dir)
|
| 94 |
+
|
| 95 |
+
|
| 96 |
+
# βββ Verify βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 97 |
+
|
| 98 |
+
def _verify(data_dir: Path):
|
| 99 |
+
idx_path = data_dir / "index.json"
|
| 100 |
+
if not idx_path.exists():
|
| 101 |
+
# Search one level deeper
|
| 102 |
+
found = list(data_dir.glob("**/index.json"))
|
| 103 |
+
if found:
|
| 104 |
+
print(f" Note: index.json found at {found[0]}, not root. Check your --out_dir.")
|
| 105 |
+
else:
|
| 106 |
+
print(f" WARNING: index.json NOT found in {data_dir}")
|
| 107 |
+
return
|
| 108 |
+
|
| 109 |
+
with open(idx_path) as f:
|
| 110 |
+
idx = json.load(f)
|
| 111 |
+
chunks = idx.get("chunks", [])
|
| 112 |
+
total_tokens = sum(c.get("dim", 0) for c in chunks)
|
| 113 |
+
present = sum(1 for c in chunks if (data_dir / c["filename"]).exists())
|
| 114 |
+
missing = len(chunks) - present
|
| 115 |
+
|
| 116 |
+
print(f"\n Dataset verified:")
|
| 117 |
+
print(f" Chunks declared : {len(chunks)}")
|
| 118 |
+
print(f" Chunks on disk : {present}")
|
| 119 |
+
print(f" Missing chunks : {missing}")
|
| 120 |
+
print(f" Total tokens : {total_tokens:,}")
|
| 121 |
+
if missing > 0:
|
| 122 |
+
print(f" WARNING: {missing} chunk(s) missing β training will error on those blocks!")
|
| 123 |
+
else:
|
| 124 |
+
print(f" All chunks present. Ready to train.")
|
| 125 |
+
|
| 126 |
+
|
| 127 |
+
# βββ Args βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 128 |
+
|
| 129 |
+
def parse_args():
|
| 130 |
+
p = argparse.ArgumentParser(description="LUNA dataset fetcher")
|
| 131 |
+
p.add_argument("--source", choices=["huggingface", "gdrive", "local"], required=True)
|
| 132 |
+
p.add_argument("--out_dir", type=str, default="/workspace/data/litdata_pretrain_final",
|
| 133 |
+
help="Where to save the dataset")
|
| 134 |
+
p.add_argument("--hf_repo", type=str, default="",
|
| 135 |
+
help="HuggingFace dataset repo ID (e.g. YourName/LUNA-pretrain-data)")
|
| 136 |
+
p.add_argument("--hf_token", type=str, default=os.environ.get("HF_TOKEN", ""),
|
| 137 |
+
help="HuggingFace token (or set HF_TOKEN env var)")
|
| 138 |
+
p.add_argument("--gdrive_id", type=str, default="",
|
| 139 |
+
help="Google Drive file/folder ID")
|
| 140 |
+
p.add_argument("--local_path", type=str, default="Base/data/litdata_pretrain_final",
|
| 141 |
+
help="Local path to the dataset (for local source)")
|
| 142 |
+
return p.parse_args()
|
| 143 |
+
|
| 144 |
+
|
| 145 |
+
if __name__ == "__main__":
|
| 146 |
+
args = parse_args()
|
| 147 |
+
out = Path(args.out_dir)
|
| 148 |
+
|
| 149 |
+
if args.source == "huggingface":
|
| 150 |
+
if not args.hf_repo:
|
| 151 |
+
print("ERROR: --hf_repo required for HuggingFace source")
|
| 152 |
+
sys.exit(1)
|
| 153 |
+
download_huggingface(args.hf_repo, out, hf_token=args.hf_token or None)
|
| 154 |
+
|
| 155 |
+
elif args.source == "gdrive":
|
| 156 |
+
if not args.gdrive_id:
|
| 157 |
+
print("ERROR: --gdrive_id required for GDrive source")
|
| 158 |
+
sys.exit(1)
|
| 159 |
+
download_gdrive(args.gdrive_id, out)
|
| 160 |
+
|
| 161 |
+
elif args.source == "local":
|
| 162 |
+
copy_local(args.local_path, out)
|
quantisations/convert_to_gguf.py
ADDED
|
@@ -0,0 +1,195 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Convert a litgpt GPT-NeoX checkpoint to GGUF format (F32) for LM Studio.
|
| 3 |
+
|
| 4 |
+
Usage:
|
| 5 |
+
python quantisations/convert_to_gguf.py
|
| 6 |
+
python quantisations/convert_to_gguf.py --checkpoint_dir Base/out/finetune/custom-100m-english-instruct/step-008000
|
| 7 |
+
"""
|
| 8 |
+
|
| 9 |
+
import argparse
|
| 10 |
+
import json
|
| 11 |
+
import torch
|
| 12 |
+
import numpy as np
|
| 13 |
+
import yaml
|
| 14 |
+
from pathlib import Path
|
| 15 |
+
from gguf import GGUFWriter
|
| 16 |
+
|
| 17 |
+
|
| 18 |
+
def main():
|
| 19 |
+
parser = argparse.ArgumentParser(description="Convert litgpt checkpoint to GGUF")
|
| 20 |
+
parser.add_argument("--checkpoint_dir", type=str,
|
| 21 |
+
default="Base/out/finetune/custom-100m-english-instruct/final")
|
| 22 |
+
parser.add_argument("--tokenizer_dir", type=str,
|
| 23 |
+
default="Base/checkpoints/EleutherAI/pythia-160m")
|
| 24 |
+
parser.add_argument("--output", type=str, default=None,
|
| 25 |
+
help="Output GGUF path (default: quantisations/LUNA-100M-F32.gguf)")
|
| 26 |
+
parser.add_argument("--model_name", type=str, default="LUNA-100M")
|
| 27 |
+
args = parser.parse_args()
|
| 28 |
+
|
| 29 |
+
ckpt_dir = Path(args.checkpoint_dir)
|
| 30 |
+
tok_dir = Path(args.tokenizer_dir)
|
| 31 |
+
|
| 32 |
+
if args.output is None:
|
| 33 |
+
args.output = f"quantisations/{args.model_name}-F32.gguf"
|
| 34 |
+
output_path = Path(args.output)
|
| 35 |
+
output_path.parent.mkdir(parents=True, exist_ok=True)
|
| 36 |
+
|
| 37 |
+
# ββ Load checkpoint ββββββββββββββββββββββββββββββββββββββββββββββ
|
| 38 |
+
print(f"Loading checkpoint from {ckpt_dir}...")
|
| 39 |
+
sd = torch.load(ckpt_dir / "lit_model.pth", map_location="cpu", weights_only=False)
|
| 40 |
+
if "model" in sd:
|
| 41 |
+
sd = sd["model"]
|
| 42 |
+
|
| 43 |
+
# ββ Read model config ββββββββββββββββββββββββββββββββββββββββββββ
|
| 44 |
+
with open(ckpt_dir / "model_config.yaml") as f:
|
| 45 |
+
cfg = yaml.safe_load(f)
|
| 46 |
+
|
| 47 |
+
n_layer = cfg["n_layer"]
|
| 48 |
+
n_embd = cfg["n_embd"]
|
| 49 |
+
n_head = cfg["n_head"]
|
| 50 |
+
n_query_groups = cfg.get("n_query_groups", n_head)
|
| 51 |
+
head_size = cfg.get("head_size", n_embd // n_head)
|
| 52 |
+
intermediate_size = cfg["intermediate_size"]
|
| 53 |
+
block_size = cfg["block_size"]
|
| 54 |
+
vocab_size = cfg["vocab_size"]
|
| 55 |
+
padded_vocab_size = cfg["padded_vocab_size"]
|
| 56 |
+
norm_eps = float(cfg.get("norm_eps", 1e-5))
|
| 57 |
+
rotary_pct = float(cfg.get("rotary_percentage", 0.25))
|
| 58 |
+
rope_dim = int(head_size * rotary_pct)
|
| 59 |
+
rope_base = float(cfg.get("rope_base", 10000))
|
| 60 |
+
parallel_residual = cfg.get("parallel_residual", True)
|
| 61 |
+
|
| 62 |
+
total_params = sum(v.numel() for v in sd.values())
|
| 63 |
+
print(f"Model: {n_layer}L, {n_embd}d, {n_head}h, vocab={vocab_size} "
|
| 64 |
+
f"(padded {padded_vocab_size}), {total_params:,} params")
|
| 65 |
+
|
| 66 |
+
# ββ Create GGUF writer βββββββββββββββββββββββββββββββββββββββββββ
|
| 67 |
+
print(f"Writing GGUF to {output_path}...")
|
| 68 |
+
writer = GGUFWriter(str(output_path), arch="gptneox")
|
| 69 |
+
|
| 70 |
+
# General metadata
|
| 71 |
+
writer.add_name(args.model_name)
|
| 72 |
+
writer.add_description(f"{args.model_name} by Asterizer β GPT-NeoX {total_params // 1_000_000}M")
|
| 73 |
+
|
| 74 |
+
# Architecture parameters
|
| 75 |
+
writer.add_context_length(block_size)
|
| 76 |
+
writer.add_embedding_length(n_embd)
|
| 77 |
+
writer.add_block_count(n_layer)
|
| 78 |
+
writer.add_feed_forward_length(intermediate_size)
|
| 79 |
+
writer.add_head_count(n_head)
|
| 80 |
+
writer.add_head_count_kv(n_query_groups)
|
| 81 |
+
writer.add_layer_norm_eps(norm_eps)
|
| 82 |
+
writer.add_rope_dimension_count(rope_dim)
|
| 83 |
+
writer.add_rope_freq_base(rope_base)
|
| 84 |
+
writer.add_parallel_residual(parallel_residual)
|
| 85 |
+
writer.add_file_type(0) # ALL_F32
|
| 86 |
+
|
| 87 |
+
# ββ Tokenizer ββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 88 |
+
print("Processing tokenizer...")
|
| 89 |
+
with open(tok_dir / "tokenizer.json", "r", encoding="utf-8") as f:
|
| 90 |
+
tok_data = json.load(f)
|
| 91 |
+
|
| 92 |
+
vocab = tok_data["model"]["vocab"]
|
| 93 |
+
merges = tok_data["model"]["merges"]
|
| 94 |
+
|
| 95 |
+
# Identify special tokens
|
| 96 |
+
special_ids = set()
|
| 97 |
+
if "added_tokens" in tok_data:
|
| 98 |
+
for at in tok_data["added_tokens"]:
|
| 99 |
+
if at.get("special", False):
|
| 100 |
+
special_ids.add(at["id"])
|
| 101 |
+
|
| 102 |
+
# Build sorted token list
|
| 103 |
+
sorted_tokens = sorted(vocab.items(), key=lambda x: x[1])
|
| 104 |
+
|
| 105 |
+
tokens = []
|
| 106 |
+
token_types = []
|
| 107 |
+
for token_str, token_id in sorted_tokens:
|
| 108 |
+
tokens.append(token_str.encode("utf-8"))
|
| 109 |
+
token_types.append(3 if token_id in special_ids else 1)
|
| 110 |
+
|
| 111 |
+
# Pad to padded_vocab_size
|
| 112 |
+
while len(tokens) < padded_vocab_size:
|
| 113 |
+
tokens.append(f"[PAD{len(tokens)}]".encode("utf-8"))
|
| 114 |
+
token_types.append(3) # CONTROL
|
| 115 |
+
|
| 116 |
+
writer.add_tokenizer_model("gpt2")
|
| 117 |
+
writer.add_token_list(tokens)
|
| 118 |
+
writer.add_token_types(token_types)
|
| 119 |
+
writer.add_token_merges([m.encode("utf-8") for m in merges])
|
| 120 |
+
writer.add_bos_token_id(0)
|
| 121 |
+
writer.add_eos_token_id(0)
|
| 122 |
+
|
| 123 |
+
print(f" Tokens: {len(sorted_tokens)} real + {padded_vocab_size - len(sorted_tokens)} padding")
|
| 124 |
+
print(f" Merges: {len(merges)}")
|
| 125 |
+
|
| 126 |
+
# ββ Add tensors ββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 127 |
+
# litgpt GPT-NeoX β GGUF GPT-NeoX tensor name mapping:
|
| 128 |
+
# transformer.wte.weight β token_embd.weight
|
| 129 |
+
# transformer.ln_f.{weight,bias} β output_norm.{weight,bias}
|
| 130 |
+
# lm_head.weight β output.weight
|
| 131 |
+
# transformer.h.{i}.norm_1.* β blk.{i}.attn_norm.*
|
| 132 |
+
# transformer.h.{i}.norm_2.* β blk.{i}.ffn_norm.*
|
| 133 |
+
# transformer.h.{i}.attn.qkv.* β blk.{i}.attn_qkv.* (interleaved per-head)
|
| 134 |
+
# transformer.h.{i}.attn.proj.* β blk.{i}.attn_output.*
|
| 135 |
+
# transformer.h.{i}.mlp.fc.* β blk.{i}.ffn_up.*
|
| 136 |
+
# transformer.h.{i}.mlp.proj.* β blk.{i}.ffn_down.*
|
| 137 |
+
|
| 138 |
+
print("Adding tensors...")
|
| 139 |
+
tensor_count = 0
|
| 140 |
+
|
| 141 |
+
def add(gguf_name, litgpt_key):
|
| 142 |
+
nonlocal tensor_count
|
| 143 |
+
tensor = sd[litgpt_key].float().numpy()
|
| 144 |
+
writer.add_tensor(gguf_name, tensor)
|
| 145 |
+
tensor_count += 1
|
| 146 |
+
|
| 147 |
+
# Global tensors
|
| 148 |
+
add("token_embd.weight", "transformer.wte.weight")
|
| 149 |
+
add("output_norm.weight", "transformer.ln_f.weight")
|
| 150 |
+
add("output_norm.bias", "transformer.ln_f.bias")
|
| 151 |
+
add("output.weight", "lm_head.weight")
|
| 152 |
+
|
| 153 |
+
# Per-layer tensors
|
| 154 |
+
for i in range(n_layer):
|
| 155 |
+
pfx = f"transformer.h.{i}"
|
| 156 |
+
blk = f"blk.{i}"
|
| 157 |
+
print(f" Layer {i + 1}/{n_layer}")
|
| 158 |
+
|
| 159 |
+
# Layer norms
|
| 160 |
+
add(f"{blk}.attn_norm.weight", f"{pfx}.norm_1.weight")
|
| 161 |
+
add(f"{blk}.attn_norm.bias", f"{pfx}.norm_1.bias")
|
| 162 |
+
add(f"{blk}.ffn_norm.weight", f"{pfx}.norm_2.weight")
|
| 163 |
+
add(f"{blk}.ffn_norm.bias", f"{pfx}.norm_2.bias")
|
| 164 |
+
|
| 165 |
+
# Attention QKV (litgpt uses same interleaved-per-head layout as HF GPT-NeoX)
|
| 166 |
+
add(f"{blk}.attn_qkv.weight", f"{pfx}.attn.qkv.weight")
|
| 167 |
+
add(f"{blk}.attn_qkv.bias", f"{pfx}.attn.qkv.bias")
|
| 168 |
+
|
| 169 |
+
# Attention output projection
|
| 170 |
+
add(f"{blk}.attn_output.weight", f"{pfx}.attn.proj.weight")
|
| 171 |
+
add(f"{blk}.attn_output.bias", f"{pfx}.attn.proj.bias")
|
| 172 |
+
|
| 173 |
+
# FFN
|
| 174 |
+
add(f"{blk}.ffn_up.weight", f"{pfx}.mlp.fc.weight")
|
| 175 |
+
add(f"{blk}.ffn_up.bias", f"{pfx}.mlp.fc.bias")
|
| 176 |
+
add(f"{blk}.ffn_down.weight", f"{pfx}.mlp.proj.weight")
|
| 177 |
+
add(f"{blk}.ffn_down.bias", f"{pfx}.mlp.proj.bias")
|
| 178 |
+
|
| 179 |
+
# ββ Write file βββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 180 |
+
writer.write_header_to_file()
|
| 181 |
+
writer.write_kv_data_to_file()
|
| 182 |
+
writer.write_tensors_to_file()
|
| 183 |
+
writer.close()
|
| 184 |
+
|
| 185 |
+
size_mb = output_path.stat().st_size / (1024 * 1024)
|
| 186 |
+
print(f"\nDone!")
|
| 187 |
+
print(f" Output: {output_path}")
|
| 188 |
+
print(f" Tensors: {tensor_count}")
|
| 189 |
+
print(f" File size: {size_mb:.1f} MB")
|
| 190 |
+
print(f"\nLoad in LM Studio or llama.cpp:")
|
| 191 |
+
print(f" llama-cli -m {output_path} -p \"Hello\"")
|
| 192 |
+
|
| 193 |
+
|
| 194 |
+
if __name__ == "__main__":
|
| 195 |
+
main()
|
requirements.txt
ADDED
|
@@ -0,0 +1,6 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
torch>=2.3.0
|
| 2 |
+
psutil>=5.9
|
| 3 |
+
pyyaml>=6.0
|
| 4 |
+
huggingface_hub>=0.22
|
| 5 |
+
gdown>=5.1
|
| 6 |
+
tensorboard>=2.16
|
setup_and_train.sh
ADDED
|
@@ -0,0 +1,94 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env bash
|
| 2 |
+
# =============================================================================
|
| 3 |
+
# LUNA 100M β Cloud Setup & Train Entrypoint
|
| 4 |
+
# Runs on RunPod, Vast.ai, Lambda Labs, or any Linux GPU pod.
|
| 5 |
+
#
|
| 6 |
+
# USAGE (after cloning repo):
|
| 7 |
+
# bash setup_and_train.sh [gdrive|huggingface] [SOURCE_ID] [MAX_TOKENS]
|
| 8 |
+
#
|
| 9 |
+
# EXAMPLES:
|
| 10 |
+
# # Full dataset from Google Drive folder:
|
| 11 |
+
# bash setup_and_train.sh gdrive 1AbCdEfGhIjKlMnOpQrStUvWx
|
| 12 |
+
#
|
| 13 |
+
# # Full dataset from HuggingFace:
|
| 14 |
+
# bash setup_and_train.sh huggingface YourName/LUNA-pretrain-data
|
| 15 |
+
#
|
| 16 |
+
# # Quick smoke test (10M tokens only):
|
| 17 |
+
# bash setup_and_train.sh gdrive 1AbCdEfGhIjKlMnOpQrStUvWx 10000000
|
| 18 |
+
#
|
| 19 |
+
# # Dataset already on disk:
|
| 20 |
+
# bash setup_and_train.sh local /workspace/data/litdata_pretrain_final
|
| 21 |
+
# =============================================================================
|
| 22 |
+
|
| 23 |
+
set -e
|
| 24 |
+
|
| 25 |
+
DATA_SOURCE="${1:-local}"
|
| 26 |
+
DATA_ID="${2:-Base/data/litdata_pretrain_final}"
|
| 27 |
+
MAX_TOKENS="${3:-4515286950}"
|
| 28 |
+
DATA_DIR="/workspace/data/litdata_pretrain_final"
|
| 29 |
+
OUT_DIR="/workspace/out/pretrain/luna-100m"
|
| 30 |
+
|
| 31 |
+
echo "=========================================="
|
| 32 |
+
echo " LUNA 100M β Cloud Setup"
|
| 33 |
+
echo " Source : $DATA_SOURCE"
|
| 34 |
+
echo " ID/Path : $DATA_ID"
|
| 35 |
+
echo " Tokens : $MAX_TOKENS"
|
| 36 |
+
echo "=========================================="
|
| 37 |
+
|
| 38 |
+
# ββ 1. Python packages ββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 39 |
+
echo ""
|
| 40 |
+
echo "[1/4] Installing dependencies..."
|
| 41 |
+
|
| 42 |
+
pip install -q --upgrade pip
|
| 43 |
+
pip install -q \
|
| 44 |
+
torch torchvision \
|
| 45 |
+
psutil \
|
| 46 |
+
huggingface_hub \
|
| 47 |
+
gdown \
|
| 48 |
+
tensorboard \
|
| 49 |
+
litgpt 2>/dev/null || true
|
| 50 |
+
|
| 51 |
+
echo " Done."
|
| 52 |
+
|
| 53 |
+
# ββ 2. Download dataset βββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 54 |
+
echo ""
|
| 55 |
+
echo "[2/4] Fetching dataset..."
|
| 56 |
+
|
| 57 |
+
if [ "$DATA_SOURCE" = "gdrive" ]; then
|
| 58 |
+
python fetch_data.py --source gdrive --gdrive_id "$DATA_ID" --out_dir "$DATA_DIR"
|
| 59 |
+
elif [ "$DATA_SOURCE" = "huggingface" ]; then
|
| 60 |
+
HF_TOKEN="${HF_TOKEN:-}"
|
| 61 |
+
python fetch_data.py --source huggingface --hf_repo "$DATA_ID" --out_dir "$DATA_DIR" --hf_token "$HF_TOKEN"
|
| 62 |
+
elif [ "$DATA_SOURCE" = "local" ]; then
|
| 63 |
+
python fetch_data.py --source local --local_path "$DATA_ID" --out_dir "$DATA_DIR"
|
| 64 |
+
else
|
| 65 |
+
echo "Unknown source: $DATA_SOURCE (use: gdrive | huggingface | local)"
|
| 66 |
+
exit 1
|
| 67 |
+
fi
|
| 68 |
+
|
| 69 |
+
# ββ 3. System + batch size probe ββββββββββββββββββββββββββββββββββββββββββββββ
|
| 70 |
+
echo ""
|
| 71 |
+
echo "[3/4] System probe (auto-detects VRAM, RAM, CPU)..."
|
| 72 |
+
python -c "
|
| 73 |
+
import torch, psutil, os
|
| 74 |
+
props = torch.cuda.get_device_properties(0) if torch.cuda.is_available() else None
|
| 75 |
+
print(f' GPU : {props.name if props else \"CPU only\"} ({props.total_memory/1024**3:.1f} GB)' if props else ' GPU: None')
|
| 76 |
+
print(f' RAM : {psutil.virtual_memory().total/1024**3:.1f} GB')
|
| 77 |
+
print(f' CPUs : {os.cpu_count()}')
|
| 78 |
+
"
|
| 79 |
+
|
| 80 |
+
# ββ 4. Train ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 81 |
+
echo ""
|
| 82 |
+
echo "[4/4] Starting training (auto_config reads from train_config.yaml)..."
|
| 83 |
+
echo ""
|
| 84 |
+
|
| 85 |
+
python train.py \
|
| 86 |
+
--config train_config.yaml \
|
| 87 |
+
--data_path "$DATA_DIR" \
|
| 88 |
+
--out_dir "$OUT_DIR" \
|
| 89 |
+
--max_tokens "$MAX_TOKENS"
|
| 90 |
+
|
| 91 |
+
echo ""
|
| 92 |
+
echo "=========================================="
|
| 93 |
+
echo " Training complete! Output: $OUT_DIR"
|
| 94 |
+
echo "=========================================="
|
train.py
ADDED
|
@@ -0,0 +1,608 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
LUNA 100M β Config-Driven Dynamic Training Script
|
| 3 |
+
==================================================
|
| 4 |
+
Reads train_config.yaml for all hyperparameters.
|
| 5 |
+
|
| 6 |
+
auto_config: true -> hardware probed; batch/lr/workers set automatically
|
| 7 |
+
auto_config: false -> every value in config used exactly as-is
|
| 8 |
+
|
| 9 |
+
Usage:
|
| 10 |
+
python train.py # uses train_config.yaml defaults
|
| 11 |
+
python train.py --config train_config.yaml # explicit config path
|
| 12 |
+
python train.py --data_path /mnt/data/litdata_final # override data path only
|
| 13 |
+
python train.py --max_tokens 10000000 # short smoke-test run
|
| 14 |
+
"""
|
| 15 |
+
|
| 16 |
+
import os
|
| 17 |
+
import gc
|
| 18 |
+
import sys
|
| 19 |
+
import math
|
| 20 |
+
import time
|
| 21 |
+
import json
|
| 22 |
+
import argparse
|
| 23 |
+
import yaml
|
| 24 |
+
import psutil
|
| 25 |
+
import torch
|
| 26 |
+
import torch.nn as nn
|
| 27 |
+
import torch.nn.functional as F
|
| 28 |
+
from torch.amp import autocast, GradScaler
|
| 29 |
+
from pathlib import Path
|
| 30 |
+
|
| 31 |
+
|
| 32 |
+
# βββ Model ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 33 |
+
|
| 34 |
+
class RotaryEmbedding(nn.Module):
|
| 35 |
+
def __init__(self, dim, max_seq_len=1024):
|
| 36 |
+
super().__init__()
|
| 37 |
+
inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
|
| 38 |
+
self.register_buffer("inv_freq", inv_freq)
|
| 39 |
+
t = torch.arange(max_seq_len).float()
|
| 40 |
+
freqs = torch.einsum("i,j->ij", t, inv_freq)
|
| 41 |
+
emb = torch.cat([freqs, freqs], dim=-1)
|
| 42 |
+
self.register_buffer("cos_cached", emb.cos())
|
| 43 |
+
self.register_buffer("sin_cached", emb.sin())
|
| 44 |
+
|
| 45 |
+
def forward(self, seq_len):
|
| 46 |
+
return self.cos_cached[:seq_len], self.sin_cached[:seq_len]
|
| 47 |
+
|
| 48 |
+
|
| 49 |
+
def rotate_half(x):
|
| 50 |
+
x1, x2 = x.chunk(2, dim=-1)
|
| 51 |
+
return torch.cat([-x2, x1], dim=-1)
|
| 52 |
+
|
| 53 |
+
|
| 54 |
+
def apply_rotary(x, cos, sin):
|
| 55 |
+
c = cos.unsqueeze(0).unsqueeze(0)
|
| 56 |
+
s = sin.unsqueeze(0).unsqueeze(0)
|
| 57 |
+
return x * c + rotate_half(x) * s
|
| 58 |
+
|
| 59 |
+
|
| 60 |
+
class CausalSelfAttention(nn.Module):
|
| 61 |
+
def __init__(self, n_embd, n_head, block_size, rotary_pct=0.25):
|
| 62 |
+
super().__init__()
|
| 63 |
+
self.n_head = n_head
|
| 64 |
+
self.head_dim = n_embd // n_head
|
| 65 |
+
self.rot_dim = int(self.head_dim * rotary_pct)
|
| 66 |
+
self.c_attn = nn.Linear(n_embd, 3 * n_embd, bias=True)
|
| 67 |
+
self.c_proj = nn.Linear(n_embd, n_embd, bias=True)
|
| 68 |
+
self.rotary = RotaryEmbedding(self.rot_dim, block_size)
|
| 69 |
+
|
| 70 |
+
def forward(self, x):
|
| 71 |
+
B, T, C = x.size()
|
| 72 |
+
qkv = self.c_attn(x).reshape(B, T, 3, self.n_head, self.head_dim).permute(2, 0, 3, 1, 4)
|
| 73 |
+
q, k, v = qkv.unbind(0)
|
| 74 |
+
cos, sin = self.rotary(T)
|
| 75 |
+
q = torch.cat([apply_rotary(q[..., :self.rot_dim], cos, sin), q[..., self.rot_dim:]], dim=-1)
|
| 76 |
+
k = torch.cat([apply_rotary(k[..., :self.rot_dim], cos, sin), k[..., self.rot_dim:]], dim=-1)
|
| 77 |
+
y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
|
| 78 |
+
return self.c_proj(y.transpose(1, 2).contiguous().view(B, T, C))
|
| 79 |
+
|
| 80 |
+
|
| 81 |
+
class MLP(nn.Module):
|
| 82 |
+
def __init__(self, n_embd):
|
| 83 |
+
super().__init__()
|
| 84 |
+
self.fc = nn.Linear(n_embd, 4 * n_embd, bias=True)
|
| 85 |
+
self.gelu = nn.GELU()
|
| 86 |
+
self.proj = nn.Linear(4 * n_embd, n_embd, bias=True)
|
| 87 |
+
|
| 88 |
+
def forward(self, x):
|
| 89 |
+
return self.proj(self.gelu(self.fc(x)))
|
| 90 |
+
|
| 91 |
+
|
| 92 |
+
class Block(nn.Module):
|
| 93 |
+
def __init__(self, n_embd, n_head, block_size):
|
| 94 |
+
super().__init__()
|
| 95 |
+
self.ln1 = nn.LayerNorm(n_embd)
|
| 96 |
+
self.attn = CausalSelfAttention(n_embd, n_head, block_size)
|
| 97 |
+
self.ln2 = nn.LayerNorm(n_embd)
|
| 98 |
+
self.mlp = MLP(n_embd)
|
| 99 |
+
|
| 100 |
+
def forward(self, x):
|
| 101 |
+
x = x + self.attn(self.ln1(x))
|
| 102 |
+
x = x + self.mlp(self.ln2(x))
|
| 103 |
+
return x
|
| 104 |
+
|
| 105 |
+
|
| 106 |
+
class LUNAModel(nn.Module):
|
| 107 |
+
def __init__(self, vocab_size, block_size, n_layer, n_embd, n_head):
|
| 108 |
+
super().__init__()
|
| 109 |
+
self.wte = nn.Embedding(vocab_size, n_embd)
|
| 110 |
+
self.blocks = nn.ModuleList([Block(n_embd, n_head, block_size) for _ in range(n_layer)])
|
| 111 |
+
self.ln_f = nn.LayerNorm(n_embd)
|
| 112 |
+
self.lm_head = nn.Linear(n_embd, vocab_size, bias=False)
|
| 113 |
+
self.lm_head.weight = self.wte.weight # tie
|
| 114 |
+
self.apply(self._init_weights)
|
| 115 |
+
|
| 116 |
+
def _init_weights(self, m):
|
| 117 |
+
if isinstance(m, (nn.Linear, nn.Embedding)):
|
| 118 |
+
m.weight.data.normal_(mean=0.0, std=0.02)
|
| 119 |
+
if isinstance(m, nn.Linear) and m.bias is not None:
|
| 120 |
+
m.bias.data.zero_()
|
| 121 |
+
|
| 122 |
+
def forward(self, idx, targets=None):
|
| 123 |
+
x = self.wte(idx)
|
| 124 |
+
for block in self.blocks:
|
| 125 |
+
x = block(x)
|
| 126 |
+
x = self.ln_f(x)
|
| 127 |
+
logits = self.lm_head(x)
|
| 128 |
+
loss = None
|
| 129 |
+
if targets is not None:
|
| 130 |
+
loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
|
| 131 |
+
return logits, loss
|
| 132 |
+
|
| 133 |
+
@property
|
| 134 |
+
def num_params(self):
|
| 135 |
+
return sum(p.numel() for p in self.parameters()) - self.wte.weight.numel()
|
| 136 |
+
|
| 137 |
+
|
| 138 |
+
# βββ Dataset βββββββββββββββββββββββοΏ½οΏ½οΏ½ββββββββββββββββββββββββββββββββββββββββββ
|
| 139 |
+
|
| 140 |
+
class LitDataDataset(torch.utils.data.Dataset):
|
| 141 |
+
def __init__(self, data_path: str, block_size: int = 1024):
|
| 142 |
+
import struct, numpy as np
|
| 143 |
+
self.block_size = block_size
|
| 144 |
+
self.data_path = Path(data_path)
|
| 145 |
+
with open(self.data_path / "index.json") as f:
|
| 146 |
+
idx = json.load(f)
|
| 147 |
+
self.chunks_meta = idx["chunks"]
|
| 148 |
+
self._cum_blocks = []
|
| 149 |
+
total = 0
|
| 150 |
+
for c in self.chunks_meta:
|
| 151 |
+
n = c["dim"] // (block_size + 1)
|
| 152 |
+
total += n
|
| 153 |
+
self._cum_blocks.append(total)
|
| 154 |
+
self.total_blocks = total
|
| 155 |
+
self._chunk_cache = {}
|
| 156 |
+
|
| 157 |
+
def _load_chunk(self, chunk_idx: int):
|
| 158 |
+
if chunk_idx in self._chunk_cache:
|
| 159 |
+
return self._chunk_cache[chunk_idx]
|
| 160 |
+
import struct, numpy as np
|
| 161 |
+
meta = self.chunks_meta[chunk_idx]
|
| 162 |
+
with open(self.data_path / meta["filename"], "rb") as f:
|
| 163 |
+
raw = f.read()
|
| 164 |
+
num_items = struct.unpack_from("<I", raw, 0)[0]
|
| 165 |
+
header_bytes = (num_items + 2) * 4
|
| 166 |
+
tokens = torch.from_numpy(np.frombuffer(raw[header_bytes:], dtype=np.int32).copy())
|
| 167 |
+
if len(self._chunk_cache) >= 4:
|
| 168 |
+
del self._chunk_cache[next(iter(self._chunk_cache))]
|
| 169 |
+
self._chunk_cache[chunk_idx] = tokens
|
| 170 |
+
return tokens
|
| 171 |
+
|
| 172 |
+
def __len__(self):
|
| 173 |
+
return self.total_blocks
|
| 174 |
+
|
| 175 |
+
def __getitem__(self, idx):
|
| 176 |
+
chunk_idx = 0
|
| 177 |
+
for i, cum in enumerate(self._cum_blocks):
|
| 178 |
+
if idx < cum:
|
| 179 |
+
chunk_idx = i
|
| 180 |
+
break
|
| 181 |
+
prev = self._cum_blocks[chunk_idx - 1] if chunk_idx > 0 else 0
|
| 182 |
+
tokens = self._load_chunk(chunk_idx)
|
| 183 |
+
s = (idx - prev) * (self.block_size + 1)
|
| 184 |
+
e = s + self.block_size + 1
|
| 185 |
+
chunk = tokens[s:e]
|
| 186 |
+
if len(chunk) < self.block_size + 1:
|
| 187 |
+
pad = torch.zeros(self.block_size + 1, dtype=torch.int32)
|
| 188 |
+
pad[:len(chunk)] = chunk
|
| 189 |
+
chunk = pad
|
| 190 |
+
chunk = chunk.long()
|
| 191 |
+
return chunk[:self.block_size], chunk[1:self.block_size + 1]
|
| 192 |
+
|
| 193 |
+
|
| 194 |
+
# βββ Hardware Detection ββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 195 |
+
|
| 196 |
+
def probe_hardware():
|
| 197 |
+
info = {
|
| 198 |
+
"cpu_cores": os.cpu_count() or 4,
|
| 199 |
+
"ram_gb": psutil.virtual_memory().total / 1024**3,
|
| 200 |
+
}
|
| 201 |
+
if torch.cuda.is_available():
|
| 202 |
+
props = torch.cuda.get_device_properties(0)
|
| 203 |
+
info.update({
|
| 204 |
+
"device": "cuda",
|
| 205 |
+
"gpu_name": props.name,
|
| 206 |
+
"vram_gb": props.total_memory / 1024**3,
|
| 207 |
+
"sm_major": props.major,
|
| 208 |
+
})
|
| 209 |
+
if props.major >= 8:
|
| 210 |
+
torch.backends.cuda.matmul.allow_tf32 = True
|
| 211 |
+
torch.backends.cudnn.allow_tf32 = True
|
| 212 |
+
info["precision"] = "bf16"
|
| 213 |
+
info["dtype"] = torch.bfloat16
|
| 214 |
+
else:
|
| 215 |
+
info["precision"] = "fp16"
|
| 216 |
+
info["dtype"] = torch.float16
|
| 217 |
+
else:
|
| 218 |
+
info.update({
|
| 219 |
+
"device": "cpu",
|
| 220 |
+
"gpu_name": "CPU",
|
| 221 |
+
"vram_gb": 0,
|
| 222 |
+
"sm_major": 0,
|
| 223 |
+
"precision": "fp32",
|
| 224 |
+
"dtype": torch.float32,
|
| 225 |
+
})
|
| 226 |
+
return info
|
| 227 |
+
|
| 228 |
+
|
| 229 |
+
def probe_max_batch(model, device, dtype, seq_len, vocab_size, max_search=256):
|
| 230 |
+
"""Binary search: max micro_batch that survives fwd+bwd+optim. Safety: Γ0.82."""
|
| 231 |
+
tmp_opt = torch.optim.AdamW(model.parameters(), lr=1e-4)
|
| 232 |
+
lo, hi, best = 1, max_search, 1
|
| 233 |
+
while lo <= hi:
|
| 234 |
+
mid = (lo + hi) // 2
|
| 235 |
+
try:
|
| 236 |
+
torch.cuda.empty_cache(); gc.collect()
|
| 237 |
+
tmp_opt.zero_grad(set_to_none=True)
|
| 238 |
+
x = torch.randint(0, vocab_size, (mid, seq_len), device=device)
|
| 239 |
+
t = torch.randint(0, vocab_size, (mid, seq_len), device=device)
|
| 240 |
+
with autocast(device_type="cuda", dtype=dtype):
|
| 241 |
+
_, loss = model(x, t)
|
| 242 |
+
loss.backward()
|
| 243 |
+
tmp_opt.step()
|
| 244 |
+
tmp_opt.zero_grad(set_to_none=True)
|
| 245 |
+
best = mid; lo = mid + 1
|
| 246 |
+
del x, t, loss; torch.cuda.empty_cache()
|
| 247 |
+
except (torch.cuda.OutOfMemoryError, RuntimeError):
|
| 248 |
+
try: del x, t, loss
|
| 249 |
+
except: pass
|
| 250 |
+
torch.cuda.empty_cache()
|
| 251 |
+
tmp_opt.zero_grad(set_to_none=True)
|
| 252 |
+
hi = mid - 1
|
| 253 |
+
del tmp_opt; torch.cuda.empty_cache(); gc.collect()
|
| 254 |
+
return max(1, int(best * 0.82))
|
| 255 |
+
|
| 256 |
+
|
| 257 |
+
# βββ LR Schedule ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 258 |
+
|
| 259 |
+
def cosine_lr(step, warmup, total, lr_max, lr_min):
|
| 260 |
+
if step < warmup:
|
| 261 |
+
return lr_max * (step + 1) / warmup
|
| 262 |
+
p = (step - warmup) / max(1, total - warmup)
|
| 263 |
+
return lr_min + 0.5 * (1 + math.cos(math.pi * p)) * (lr_max - lr_min)
|
| 264 |
+
|
| 265 |
+
|
| 266 |
+
# βββ Config Loading βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 267 |
+
|
| 268 |
+
def load_config(config_path: str) -> dict:
|
| 269 |
+
"""Load YAML config and return flat namespace dict."""
|
| 270 |
+
with open(config_path, encoding="utf-8") as f:
|
| 271 |
+
raw = yaml.safe_load(f)
|
| 272 |
+
|
| 273 |
+
cfg = {
|
| 274 |
+
# top-level
|
| 275 |
+
"auto_config": raw.get("auto_config", True),
|
| 276 |
+
"data_path": raw.get("data_path", "Base/data/litdata_pretrain_final"),
|
| 277 |
+
"out_dir": raw.get("out_dir", "out/pretrain/luna-100m"),
|
| 278 |
+
"tokenizer_dir": raw.get("tokenizer_dir", "Base/checkpoints/EleutherAI/pythia-160m"),
|
| 279 |
+
# model
|
| 280 |
+
"vocab_size": raw["model"]["vocab_size"],
|
| 281 |
+
"seq_len": raw["model"]["seq_len"],
|
| 282 |
+
"n_layer": raw["model"]["n_layer"],
|
| 283 |
+
"n_embd": raw["model"]["n_embd"],
|
| 284 |
+
"n_head": raw["model"]["n_head"],
|
| 285 |
+
# train
|
| 286 |
+
"max_tokens": raw["train"]["max_tokens"],
|
| 287 |
+
"lr_warmup_steps":raw["train"]["lr_warmup_steps"],
|
| 288 |
+
"save_interval": raw["train"]["save_interval"],
|
| 289 |
+
"log_interval": raw["train"]["log_interval"],
|
| 290 |
+
"max_norm": raw["train"]["max_norm"],
|
| 291 |
+
# optimizer
|
| 292 |
+
"lr": raw["optimizer"]["lr"],
|
| 293 |
+
"min_lr": raw["optimizer"]["min_lr"],
|
| 294 |
+
"weight_decay": raw["optimizer"]["weight_decay"],
|
| 295 |
+
"betas": tuple(raw["optimizer"]["betas"]),
|
| 296 |
+
"eps": raw["optimizer"]["eps"],
|
| 297 |
+
# batch
|
| 298 |
+
"global_batch": raw["batch"]["global_batch"],
|
| 299 |
+
"micro_batch": raw["batch"]["micro_batch"],
|
| 300 |
+
"grad_accum": raw["batch"]["grad_accum"],
|
| 301 |
+
# dataloader
|
| 302 |
+
"num_workers": raw["dataloader"]["num_workers"],
|
| 303 |
+
"pin_memory": raw["dataloader"]["pin_memory"],
|
| 304 |
+
# hardware
|
| 305 |
+
"precision": raw["hardware"]["precision"],
|
| 306 |
+
"compile": raw["hardware"]["compile"],
|
| 307 |
+
}
|
| 308 |
+
return cfg
|
| 309 |
+
|
| 310 |
+
|
| 311 |
+
def apply_cli_overrides(cfg: dict, cli_args: argparse.Namespace) -> dict:
|
| 312 |
+
"""CLI args override config values (only if explicitly provided)."""
|
| 313 |
+
for key, val in vars(cli_args).items():
|
| 314 |
+
if key == "config":
|
| 315 |
+
continue
|
| 316 |
+
if val is not None: # argparse default=None means "not provided"
|
| 317 |
+
cfg[key] = val
|
| 318 |
+
return cfg
|
| 319 |
+
|
| 320 |
+
|
| 321 |
+
def resolve_auto(cfg: dict, hw: dict) -> dict:
|
| 322 |
+
"""
|
| 323 |
+
When auto_config=True: override batch, workers, lr-warmup, pin_memory,
|
| 324 |
+
precision from real hardware. Never touches model arch or max_tokens.
|
| 325 |
+
Returns updated cfg plus injected hw info.
|
| 326 |
+
"""
|
| 327 |
+
if not cfg["auto_config"]:
|
| 328 |
+
print(" [CONFIG] auto_config=false -- using manual values as-is")
|
| 329 |
+
cfg.update({"_hw": hw})
|
| 330 |
+
return cfg
|
| 331 |
+
|
| 332 |
+
print(" [CONFIG] auto_config=true -- tuning settings to this hardware")
|
| 333 |
+
|
| 334 |
+
# Precision
|
| 335 |
+
cfg["precision"] = hw["precision"]
|
| 336 |
+
cfg["_dtype"] = hw["dtype"]
|
| 337 |
+
|
| 338 |
+
# Workers
|
| 339 |
+
auto_workers = hw["cpu_cores"] // 2
|
| 340 |
+
# Cap by RAM: each worker caches up to 4 chunks Γ ~67MB
|
| 341 |
+
max_by_ram = max(0, int(hw["ram_gb"] * 0.25 * 1024 / 268))
|
| 342 |
+
cfg["num_workers"] = min(auto_workers, max_by_ram, hw["cpu_cores"])
|
| 343 |
+
if cfg["num_workers"] == -1:
|
| 344 |
+
cfg["num_workers"] = 0
|
| 345 |
+
|
| 346 |
+
# Pin memory
|
| 347 |
+
cfg["pin_memory"] = hw["ram_gb"] > 16 and hw["device"] == "cuda"
|
| 348 |
+
|
| 349 |
+
# LR warmup: 5% of total steps (will be computed again in train())
|
| 350 |
+
cfg["_auto_warmup"] = True # flag: recompute once total_steps is known
|
| 351 |
+
|
| 352 |
+
# LR scaling: sqrt(global_batch / 120) relative to base lr
|
| 353 |
+
base_global = 120
|
| 354 |
+
cfg["lr"] = cfg["lr"] * math.sqrt(cfg["global_batch"] / base_global)
|
| 355 |
+
cfg["min_lr"] = cfg["min_lr"] * math.sqrt(cfg["global_batch"] / base_global)
|
| 356 |
+
|
| 357 |
+
cfg["_hw"] = hw
|
| 358 |
+
return cfg
|
| 359 |
+
|
| 360 |
+
|
| 361 |
+
# βββ Training βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 362 |
+
|
| 363 |
+
SEP = "=" * 72
|
| 364 |
+
|
| 365 |
+
def train(cfg: dict):
|
| 366 |
+
hw = cfg["_hw"]
|
| 367 |
+
device = torch.device(hw["device"])
|
| 368 |
+
|
| 369 |
+
# Pick precision dtype
|
| 370 |
+
if cfg["auto_config"]:
|
| 371 |
+
dtype = hw.get("dtype", torch.float32)
|
| 372 |
+
else:
|
| 373 |
+
dtype = {"bf16": torch.bfloat16, "fp16": torch.float16,
|
| 374 |
+
"fp32": torch.float32}.get(cfg["precision"], torch.float32)
|
| 375 |
+
|
| 376 |
+
print(SEP)
|
| 377 |
+
print(" LUNA 100M - Training")
|
| 378 |
+
print(SEP)
|
| 379 |
+
mode = "AUTO" if cfg["auto_config"] else "MANUAL"
|
| 380 |
+
print(f" Config mode : {mode}")
|
| 381 |
+
print(f" GPU : {hw['gpu_name']} ({hw['vram_gb']:.1f} GB)")
|
| 382 |
+
print(f" RAM : {hw['ram_gb']:.1f} GB CPU: {hw['cpu_cores']} cores")
|
| 383 |
+
print(f" Precision : {cfg['precision']} dtype={dtype}")
|
| 384 |
+
print(f" Workers : {cfg['num_workers']} pin_memory={cfg['pin_memory']}")
|
| 385 |
+
|
| 386 |
+
# ββ Model ββββββββββββββββββββββββββββββββββββββββββοΏ½οΏ½οΏ½ββββββββββββββββββββββ
|
| 387 |
+
print(f"\n Building LUNA-100M...")
|
| 388 |
+
model = LUNAModel(
|
| 389 |
+
vocab_size=cfg["vocab_size"],
|
| 390 |
+
block_size=cfg["seq_len"],
|
| 391 |
+
n_layer=cfg["n_layer"],
|
| 392 |
+
n_embd=cfg["n_embd"],
|
| 393 |
+
n_head=cfg["n_head"],
|
| 394 |
+
).to(device)
|
| 395 |
+
|
| 396 |
+
if cfg["compile"] and device.type == "cuda":
|
| 397 |
+
try:
|
| 398 |
+
import triton # noqa
|
| 399 |
+
model = torch.compile(model, mode="reduce-overhead")
|
| 400 |
+
print(" torch.compile: enabled (reduce-overhead)")
|
| 401 |
+
except (ImportError, Exception):
|
| 402 |
+
print(" torch.compile: skipped (Triton unavailable - Linux/cloud only)")
|
| 403 |
+
|
| 404 |
+
print(f" Parameters: {model.num_params:,} (unique)")
|
| 405 |
+
|
| 406 |
+
# ββ Batch sizing ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 407 |
+
if cfg["auto_config"] and device.type == "cuda":
|
| 408 |
+
print(f"\n Probing max micro_batch_size (VRAM search)...")
|
| 409 |
+
probe_m = LUNAModel(
|
| 410 |
+
vocab_size=cfg["vocab_size"], block_size=cfg["seq_len"],
|
| 411 |
+
n_layer=cfg["n_layer"], n_embd=cfg["n_embd"], n_head=cfg["n_head"],
|
| 412 |
+
).to(device)
|
| 413 |
+
max_mbs = probe_max_batch(
|
| 414 |
+
probe_m, device, dtype, cfg["seq_len"], cfg["vocab_size"]
|
| 415 |
+
)
|
| 416 |
+
del probe_m; torch.cuda.empty_cache(); gc.collect()
|
| 417 |
+
# grad_accum to hit global_batch
|
| 418 |
+
grad_accum = max(1, math.ceil(cfg["global_batch"] / max_mbs))
|
| 419 |
+
effective_batch = max_mbs * grad_accum
|
| 420 |
+
print(f" AUTO -> micro_batch={max_mbs}, grad_accum={grad_accum}, "
|
| 421 |
+
f"effective_batch={effective_batch}")
|
| 422 |
+
else:
|
| 423 |
+
max_mbs = cfg["micro_batch"]
|
| 424 |
+
grad_accum = cfg["grad_accum"]
|
| 425 |
+
effective_batch = max_mbs * grad_accum
|
| 426 |
+
print(f"\n MANUAL -> micro_batch={max_mbs}, grad_accum={grad_accum}, "
|
| 427 |
+
f"effective_batch={effective_batch}")
|
| 428 |
+
|
| 429 |
+
tokens_per_step = effective_batch * cfg["seq_len"]
|
| 430 |
+
print(f" Tokens/step : {tokens_per_step:,}")
|
| 431 |
+
|
| 432 |
+
# ββ Dataset βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 433 |
+
print(f"\n Dataset: {cfg['data_path']}")
|
| 434 |
+
dataset = LitDataDataset(cfg["data_path"], block_size=cfg["seq_len"])
|
| 435 |
+
print(f" Blocks : {len(dataset):,} ({len(dataset) * cfg['seq_len']:,} tokens)")
|
| 436 |
+
|
| 437 |
+
loader = torch.utils.data.DataLoader(
|
| 438 |
+
dataset,
|
| 439 |
+
batch_size=max_mbs,
|
| 440 |
+
shuffle=True,
|
| 441 |
+
num_workers=cfg["num_workers"],
|
| 442 |
+
pin_memory=cfg["pin_memory"],
|
| 443 |
+
drop_last=True,
|
| 444 |
+
prefetch_factor=4 if cfg["num_workers"] > 0 else None,
|
| 445 |
+
persistent_workers=cfg["num_workers"] > 0,
|
| 446 |
+
)
|
| 447 |
+
|
| 448 |
+
# ββ Optimiser βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 449 |
+
fused_ok = device.type == "cuda" and hasattr(torch.optim, "AdamW")
|
| 450 |
+
try:
|
| 451 |
+
optimizer = torch.optim.AdamW(
|
| 452 |
+
model.parameters(),
|
| 453 |
+
lr=cfg["lr"], weight_decay=cfg["weight_decay"],
|
| 454 |
+
betas=cfg["betas"], eps=cfg["eps"],
|
| 455 |
+
fused=True,
|
| 456 |
+
)
|
| 457 |
+
except TypeError:
|
| 458 |
+
optimizer = torch.optim.AdamW(
|
| 459 |
+
model.parameters(),
|
| 460 |
+
lr=cfg["lr"], weight_decay=cfg["weight_decay"],
|
| 461 |
+
betas=cfg["betas"], eps=cfg["eps"],
|
| 462 |
+
)
|
| 463 |
+
|
| 464 |
+
use_scaler = dtype == torch.float16
|
| 465 |
+
scaler = GradScaler(enabled=use_scaler)
|
| 466 |
+
|
| 467 |
+
# ββ Schedule ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 468 |
+
total_steps = max(1, cfg["max_tokens"] // tokens_per_step)
|
| 469 |
+
if cfg["auto_config"] and cfg.get("_auto_warmup"):
|
| 470 |
+
warmup_steps = max(50, min(500, total_steps // 20))
|
| 471 |
+
else:
|
| 472 |
+
warmup_steps = min(cfg["lr_warmup_steps"], total_steps)
|
| 473 |
+
|
| 474 |
+
out_dir = Path(cfg["out_dir"])
|
| 475 |
+
out_dir.mkdir(parents=True, exist_ok=True)
|
| 476 |
+
|
| 477 |
+
print(f"\n max_tokens : {cfg['max_tokens']:,}")
|
| 478 |
+
print(f" total_steps : {total_steps:,}")
|
| 479 |
+
print(f" warmup_steps : {warmup_steps}")
|
| 480 |
+
print(f" lr : {cfg['lr']:.2e} -> {cfg['min_lr']:.2e}")
|
| 481 |
+
print(f" save every : {cfg['save_interval']} steps")
|
| 482 |
+
print(f" out_dir : {out_dir}")
|
| 483 |
+
print(SEP)
|
| 484 |
+
|
| 485 |
+
# ββ Resume ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 486 |
+
start_step = 0
|
| 487 |
+
ckpt_path = out_dir / "latest.pt"
|
| 488 |
+
if ckpt_path.exists():
|
| 489 |
+
print(f"\n Resuming from {ckpt_path}...")
|
| 490 |
+
ckpt = torch.load(ckpt_path, map_location=device, weights_only=True)
|
| 491 |
+
model.load_state_dict(ckpt["model"])
|
| 492 |
+
optimizer.load_state_dict(ckpt["optimizer"])
|
| 493 |
+
start_step = ckpt["step"]
|
| 494 |
+
print(f" Resumed at step {start_step}")
|
| 495 |
+
|
| 496 |
+
# ββ Loop ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 497 |
+
model.train()
|
| 498 |
+
data_iter = iter(loader)
|
| 499 |
+
|
| 500 |
+
def get_batch():
|
| 501 |
+
nonlocal data_iter
|
| 502 |
+
try:
|
| 503 |
+
return next(data_iter)
|
| 504 |
+
except StopIteration:
|
| 505 |
+
data_iter = iter(loader)
|
| 506 |
+
return next(data_iter)
|
| 507 |
+
|
| 508 |
+
run_t0 = time.perf_counter()
|
| 509 |
+
tokens_seen = start_step * tokens_per_step
|
| 510 |
+
step = start_step
|
| 511 |
+
|
| 512 |
+
print(f"\n Starting training (step {start_step} -> {total_steps})...")
|
| 513 |
+
|
| 514 |
+
while step < total_steps:
|
| 515 |
+
t0 = time.perf_counter()
|
| 516 |
+
lr_now = cosine_lr(step, warmup_steps, total_steps, cfg["lr"], cfg["min_lr"])
|
| 517 |
+
for pg in optimizer.param_groups:
|
| 518 |
+
pg["lr"] = lr_now
|
| 519 |
+
|
| 520 |
+
optimizer.zero_grad(set_to_none=True)
|
| 521 |
+
total_loss = 0.0
|
| 522 |
+
|
| 523 |
+
for _ in range(grad_accum):
|
| 524 |
+
x, t = get_batch()
|
| 525 |
+
x = x.to(device, non_blocking=True)
|
| 526 |
+
t = t.to(device, non_blocking=True)
|
| 527 |
+
with autocast(device_type=device.type, dtype=dtype, enabled=(device.type == "cuda")):
|
| 528 |
+
_, loss = model(x, t)
|
| 529 |
+
loss = loss / grad_accum
|
| 530 |
+
scaler.scale(loss).backward()
|
| 531 |
+
total_loss += loss.item()
|
| 532 |
+
|
| 533 |
+
scaler.unscale_(optimizer)
|
| 534 |
+
torch.nn.utils.clip_grad_norm_(model.parameters(), cfg["max_norm"])
|
| 535 |
+
scaler.step(optimizer)
|
| 536 |
+
scaler.update()
|
| 537 |
+
|
| 538 |
+
if device.type == "cuda":
|
| 539 |
+
torch.cuda.synchronize()
|
| 540 |
+
|
| 541 |
+
dt = time.perf_counter() - t0
|
| 542 |
+
step += 1
|
| 543 |
+
tokens_seen += tokens_per_step
|
| 544 |
+
|
| 545 |
+
if step % cfg["log_interval"] == 0 or step <= 2:
|
| 546 |
+
tps = tokens_per_step / dt
|
| 547 |
+
steps_left = total_steps - step
|
| 548 |
+
eta_h = steps_left * dt / 3600
|
| 549 |
+
vram = torch.cuda.memory_allocated() / 1024**3 if device.type == "cuda" else 0
|
| 550 |
+
print(f" step {step:6d}/{total_steps} | loss {total_loss:.4f} | "
|
| 551 |
+
f"lr {lr_now:.2e} | {tps:,.0f} tok/s | VRAM {vram:.1f}GB | ETA {eta_h:.1f}h")
|
| 552 |
+
|
| 553 |
+
if step % cfg["save_interval"] == 0 or step == total_steps:
|
| 554 |
+
raw = model._orig_mod if hasattr(model, "_orig_mod") else model
|
| 555 |
+
step_dir = out_dir / f"step-{step:08d}"
|
| 556 |
+
step_dir.mkdir(parents=True, exist_ok=True)
|
| 557 |
+
torch.save(raw.state_dict(), step_dir / "lit_model.pth")
|
| 558 |
+
torch.save({"step": step, "model": raw.state_dict(),
|
| 559 |
+
"optimizer": optimizer.state_dict(),
|
| 560 |
+
"tokens_seen": tokens_seen},
|
| 561 |
+
out_dir / "latest.pt")
|
| 562 |
+
print(f" Saved -> {step_dir}")
|
| 563 |
+
|
| 564 |
+
# ββ Final βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 565 |
+
final_dir = out_dir / "final"
|
| 566 |
+
final_dir.mkdir(parents=True, exist_ok=True)
|
| 567 |
+
raw = model._orig_mod if hasattr(model, "_orig_mod") else model
|
| 568 |
+
torch.save(raw.state_dict(), final_dir / "lit_model.pth")
|
| 569 |
+
|
| 570 |
+
import shutil
|
| 571 |
+
tok_src = Path(cfg["tokenizer_dir"])
|
| 572 |
+
if tok_src.exists():
|
| 573 |
+
shutil.copytree(tok_src, final_dir / "tokenizer", dirs_exist_ok=True)
|
| 574 |
+
|
| 575 |
+
total_h = (time.perf_counter() - run_t0) / 3600
|
| 576 |
+
print(SEP)
|
| 577 |
+
print(f" Done! {total_h:.2f} h -> {final_dir}")
|
| 578 |
+
print(SEP)
|
| 579 |
+
|
| 580 |
+
|
| 581 |
+
# βββ Entry point ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 582 |
+
|
| 583 |
+
def parse_args():
|
| 584 |
+
p = argparse.ArgumentParser(description="LUNA 100M Trainer")
|
| 585 |
+
p.add_argument("--config", type=str, default="train_config.yaml",
|
| 586 |
+
help="Path to train_config.yaml")
|
| 587 |
+
# CLI overrides (all optional - omit to use config value)
|
| 588 |
+
p.add_argument("--data_path", type=str, default=None)
|
| 589 |
+
p.add_argument("--out_dir", type=str, default=None)
|
| 590 |
+
p.add_argument("--max_tokens", type=int, default=None)
|
| 591 |
+
p.add_argument("--micro_batch", type=int, default=None)
|
| 592 |
+
p.add_argument("--global_batch",type=int, default=None)
|
| 593 |
+
p.add_argument("--lr", type=float, default=None)
|
| 594 |
+
p.add_argument("--num_workers", type=int, default=None)
|
| 595 |
+
p.add_argument("--save_interval",type=int, default=None)
|
| 596 |
+
p.add_argument("--log_interval",type=int, default=None)
|
| 597 |
+
p.add_argument("--auto_config", type=lambda x: x.lower() in ("1","true","yes"),
|
| 598 |
+
default=None, help="Override auto_config (true/false)")
|
| 599 |
+
return p.parse_args()
|
| 600 |
+
|
| 601 |
+
|
| 602 |
+
if __name__ == "__main__":
|
| 603 |
+
args = parse_args()
|
| 604 |
+
cfg = load_config(args.config)
|
| 605 |
+
cfg = apply_cli_overrides(cfg, args)
|
| 606 |
+
hw = probe_hardware()
|
| 607 |
+
cfg = resolve_auto(cfg, hw)
|
| 608 |
+
train(cfg)
|
train_config.yaml
ADDED
|
@@ -0,0 +1,67 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 2 |
+
# LUNA 100M β Training Configuration
|
| 3 |
+
# Single source of truth for all hyperparameters.
|
| 4 |
+
#
|
| 5 |
+
# auto_config: true β All batch/LR/worker settings are auto-detected from
|
| 6 |
+
# available VRAM / RAM / CPU at runtime. Your values below
|
| 7 |
+
# are used as FALLBACKS only if detection fails.
|
| 8 |
+
#
|
| 9 |
+
# auto_config: false β Every value below is used as-is. Nothing is overridden.
|
| 10 |
+
# Use this when you've already benchmarked and want
|
| 11 |
+
# repeatable, fixed runs.
|
| 12 |
+
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 13 |
+
|
| 14 |
+
auto_config: true # β flip to false to lock everything below
|
| 15 |
+
|
| 16 |
+
# ββ Data ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 17 |
+
data_path: "Base/data/litdata_pretrain_final" # local default; overridden by --data_path
|
| 18 |
+
out_dir: "out/pretrain/luna-100m"
|
| 19 |
+
tokenizer_dir: "Base/checkpoints/EleutherAI/pythia-160m"
|
| 20 |
+
|
| 21 |
+
# ββ Model (fixed for LUNA-100M β do not change) βββββββββββββββββββββββββββββββ
|
| 22 |
+
model:
|
| 23 |
+
vocab_size: 50304 # ceil(50277/128)*128 β pythia tokenizer with EOS padding
|
| 24 |
+
seq_len: 1024
|
| 25 |
+
n_layer: 10
|
| 26 |
+
n_embd: 768
|
| 27 |
+
n_head: 12
|
| 28 |
+
|
| 29 |
+
# ββ Training budget βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 30 |
+
train:
|
| 31 |
+
max_tokens: 4515286950 # full dataset (verified from index.json, 270 chunks)
|
| 32 |
+
lr_warmup_steps: 500 # [AUTO] scaled to 5% of total_steps if auto_config
|
| 33 |
+
save_interval: 1000 # save checkpoint every N optimizer steps
|
| 34 |
+
log_interval: 10 # print log every N steps
|
| 35 |
+
max_norm: 1.0 # gradient clip norm
|
| 36 |
+
|
| 37 |
+
# ββ Optimiser βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 38 |
+
optimizer:
|
| 39 |
+
lr: 0.0006 # 6e-4 [AUTO] scaled by sqrt(global_batch/120) if auto_config
|
| 40 |
+
min_lr: 0.00006 # 6e-5
|
| 41 |
+
weight_decay: 0.1
|
| 42 |
+
betas: [0.9, 0.95]
|
| 43 |
+
eps: 1.0e-8
|
| 44 |
+
|
| 45 |
+
# ββ Batch sizing ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 46 |
+
# When auto_config: true β micro_batch and grad_accum are ignored; the script
|
| 47 |
+
# probes VRAM and fills it to ~82% saturation, then
|
| 48 |
+
# computes grad_accum to hit global_batch.
|
| 49 |
+
# When auto_config: false β micro_batch Γ grad_accum must equal global_batch.
|
| 50 |
+
batch:
|
| 51 |
+
global_batch: 120 # target total samples per optimizer step
|
| 52 |
+
micro_batch: 12 # [MANUAL] samples per GPU forward pass (ignored when auto)
|
| 53 |
+
grad_accum: 10 # [MANUAL] accumulation steps (ignored when auto)
|
| 54 |
+
|
| 55 |
+
# ββ DataLoader ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 56 |
+
# When auto_config: true β num_workers auto = cpu_cores // 2, capped by RAM
|
| 57 |
+
# When auto_config: false β num_workers used as-is
|
| 58 |
+
dataloader:
|
| 59 |
+
num_workers: -1 # -1 = auto; set to 0 to disable multiprocessing
|
| 60 |
+
pin_memory: true # [AUTO] disabled if RAM < 16GB
|
| 61 |
+
|
| 62 |
+
# ββ Hardware / precision ββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 63 |
+
# When auto_config: true β precision detected from GPU compute capability
|
| 64 |
+
# When auto_config: false β use the value below
|
| 65 |
+
hardware:
|
| 66 |
+
precision: "bf16" # bf16 | fp16 | fp32
|
| 67 |
+
compile: true # torch.compile (requires Triton β Linux/cloud only)
|