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

Files changed (40) hide show
  1. .gitattributes +0 -2
  2. .gitignore +60 -0
  3. Base/checkpoints/EleutherAI/pythia-160m/config.json +24 -0
  4. Base/checkpoints/EleutherAI/pythia-160m/tokenizer.json +0 -0
  5. Base/checkpoints/EleutherAI/pythia-160m/tokenizer_config.json +9 -0
  6. Base/configs/finetune_100m_english_instruct.yaml +55 -0
  7. Base/configs/finetune_100m_instruct.yaml +56 -0
  8. Base/configs/pretrain_100m_10m.yaml +49 -0
  9. Base/configs/pretrain_100m_3b.yaml +50 -0
  10. Base/configs/pretrain_100m_3b_full.yaml +53 -0
  11. Base/configs/pretrain_100m_english.yaml +58 -0
  12. Base/scripts/audit_5b_pretrain.py +494 -0
  13. Base/scripts/audit_english_clean.py +707 -0
  14. Base/scripts/audit_litdata.py +475 -0
  15. Base/scripts/build_100m_english.py +667 -0
  16. Base/scripts/build_5b_pretrain.py +761 -0
  17. Base/scripts/build_english_corpus.py +349 -0
  18. Base/scripts/build_instruct_dataset.py +769 -0
  19. Base/scripts/chat.py +68 -0
  20. Base/scripts/clean_and_merge_pretrain.py +497 -0
  21. Base/scripts/consolidate_datasets.py +194 -0
  22. Base/scripts/consolidate_litdata.py +187 -0
  23. Base/scripts/dedup_5b_pretrain.py +298 -0
  24. Base/scripts/filter_datasets.py +151 -0
  25. Base/scripts/merge_and_build.py +701 -0
  26. Base/scripts/prepare_finetune_data.py +88 -0
  27. Base/scripts/prepare_litdata.py +275 -0
  28. Base/scripts/reclean_3b.py +593 -0
  29. Base/scripts/reclean_english.py +610 -0
  30. Base/scripts/reclean_litdata.py +444 -0
  31. Base/scripts/smart_cleanup_english.py +379 -0
  32. Base/scripts/validate_model.py +137 -0
  33. README.md +119 -2
  34. benchmark_runpod.py +412 -0
  35. fetch_data.py +162 -0
  36. quantisations/convert_to_gguf.py +195 -0
  37. requirements.txt +6 -0
  38. setup_and_train.sh +94 -0
  39. train.py +608 -0
  40. 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
+ ("&amp;", "&"), ("&lt;", "<"), ("&gt;", ">"),
124
+ ("&quot;", '"'), ("&#39;", "'"), ("&apos;", "'"),
125
+ ("&nbsp;", " "), ("&mdash;", " - "), ("&ndash;", "-"),
126
+ ("&hellip;", "..."), ("&laquo;", '"'), ("&raquo;", '"'),
127
+ ("&bull;", "- "), ("&middot;", " "), ("&copy;", "(c)"),
128
+ ("&reg;", "(R)"), ("&trade;", "(TM)"), ("&deg;", " 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
+ ("&amp;", "&"), ("&lt;", "<"), ("&gt;", ">"),
80
+ ("&quot;", '"'), ("&#39;", "'"), ("&apos;", "'"),
81
+ ("&nbsp;", " "), ("&mdash;", " - "), ("&ndash;", "-"),
82
+ ("&hellip;", "..."), ("&laquo;", '"'), ("&raquo;", '"'),
83
+ ("&bull;", "- "), ("&middot;", " "), ("&copy;", "(c)"),
84
+ ("&reg;", "(R)"), ("&trade;", "(TM)"), ("&deg;", " 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
+ ("&amp;", "&"), ("&lt;", "<"), ("&gt;", ">"),
71
+ ("&quot;", '"'), ("&#39;", "'"), ("&apos;", "'"),
72
+ ("&nbsp;", " "), ("&mdash;", " - "), ("&ndash;", "-"),
73
+ ("&hellip;", "..."), ("&laquo;", '"'), ("&raquo;", '"'),
74
+ ("&bull;", "- "), ("&middot;", " "), ("&copy;", "(c)"),
75
+ ("&reg;", "(R)"), ("&trade;", "(TM)"), ("&deg;", " 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
+ ("&amp;", "&"), ("&lt;", "<"), ("&gt;", ">"),
122
+ ("&quot;", '"'), ("&#39;", "'"), ("&apos;", "'"),
123
+ ("&nbsp;", " "), ("&mdash;", " - "), ("&ndash;", "-"),
124
+ ("&hellip;", "..."), ("&laquo;", '"'), ("&raquo;", '"'),
125
+ ("&bull;", "- "), ("&middot;", " "), ("&copy;", "(c)"),
126
+ ("&reg;", "(R)"), ("&trade;", "(TM)"), ("&deg;", " 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
+ ("&amp;", "&"), ("&lt;", "<"), ("&gt;", ">"),
63
+ ("&quot;", '"'), ("&#39;", "'"), ("&apos;", "'"),
64
+ ("&nbsp;", " "), ("&mdash;", " - "), ("&ndash;", "-"),
65
+ ("&hellip;", "..."), ("&laquo;", '"'), ("&raquo;", '"'),
66
+ ("&bull;", "- "), ("&middot;", " "), ("&copy;", "(c)"),
67
+ ("&reg;", "(R)"), ("&trade;", "(TM)"), ("&deg;", " 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
+ ("&amp;", "&"), ("&lt;", "<"), ("&gt;", ">"),
46
+ ("&quot;", '"'), ("&#39;", "'"), ("&apos;", "'"),
47
+ ("&nbsp;", " "), ("&mdash;", " - "), ("&ndash;", "-"),
48
+ ("&hellip;", "..."), ("&laquo;", '"'), ("&raquo;", '"'),
49
+ ("&bull;", "- "), ("&middot;", " "), ("&copy;", "(c)"),
50
+ ("&reg;", "(R)"), ("&trade;", "(TM)"), ("&deg;", " 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
+ ("&amp;", "&"), ("&lt;", "<"), ("&gt;", ">"),
50
+ ("&quot;", '"'), ("&#39;", "'"), ("&apos;", "'"),
51
+ ("&nbsp;", " "), ("&mdash;", " - "), ("&ndash;", "-"),
52
+ ("&hellip;", "..."), ("&laquo;", '"'), ("&raquo;", '"'),
53
+ ("&bull;", "- "), ("&middot;", " "), ("&copy;", "(c)"),
54
+ ("&reg;", "(R)"), ("&trade;", "(TM)"), ("&deg;", " 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
- Model for Desk Assitant
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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)