File size: 6,400 Bytes
7f1c419 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 | """
ViuAI-1M Training Script.
End-to-end training pipeline:
1. Load or prepare data
2. Train BPE tokenizer on the corpus
3. Create datasets and dataloaders
4. Initialize model
5. Train with the ViuAI trainer
6. Save the final model and tokenizer
Usage:
python scripts/train.py
python scripts/train.py --epochs 20 --batch_size 16 --lr 1e-4
"""
import argparse
import os
import sys
import json
# Add project root to path
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, project_root)
import torch
from torch.utils.data import DataLoader
from viuai.model.config import ViuAIConfig
from viuai.model.viuai_model import ViuAIForCausalLM
from viuai.tokenizer.bpe_tokenizer import ViuAITokenizer
from viuai.training.dataset import ViuAIDataset, ChatDataCollator
from viuai.training.trainer import ViuAITrainer
def parse_args():
parser = argparse.ArgumentParser(description="Train ViuAI-1M")
parser.add_argument("--config", type=str, default="configs/viuai_1m.json",
help="Path to model config JSON")
parser.add_argument("--train_data", type=str, default="data/train.jsonl",
help="Path to training data JSONL")
parser.add_argument("--eval_data", type=str, default="data/eval.jsonl",
help="Path to evaluation data JSONL")
parser.add_argument("--tokenizer_corpus", type=str, default="data/tokenizer_corpus.txt",
help="Path to tokenizer training corpus")
parser.add_argument("--output_dir", type=str, default="output",
help="Output directory for checkpoints and tokenizer")
parser.add_argument("--epochs", type=int, default=10,
help="Number of training epochs")
parser.add_argument("--batch_size", type=int, default=32,
help="Training batch size")
parser.add_argument("--lr", type=float, default=3e-4,
help="Learning rate")
parser.add_argument("--warmup_steps", type=int, default=100,
help="Number of warmup steps")
parser.add_argument("--device", type=str, default=None,
help="Device (cuda/cpu/mps)")
return parser.parse_args()
def main():
args = parse_args()
print(f"\n{'='*60}")
print(f" ViuAI-1M Training Pipeline")
print(f"{'='*60}\n")
# ── Step 1: Load Config ──────────────────────────────────────
print("[1/6] Loading configuration...")
config_path = os.path.join(project_root, args.config)
with open(config_path, "r") as f:
config_data = json.load(f)
config = ViuAIConfig(**config_data["architecture"])
print(f" Config: {config_path}")
print(f" Estimated parameters: {config.num_parameters:,}")
# ── Step 2: Train Tokenizer ──────────────────────────────────
print(f"\n[2/6] Training BPE tokenizer...")
tokenizer = ViuAITokenizer(vocab_size=config.vocab_size)
corpus_path = os.path.join(project_root, args.tokenizer_corpus)
if os.path.exists(corpus_path):
with open(corpus_path, "r", encoding="utf-8") as f:
corpus_texts = [line.strip() for line in f if line.strip()]
tokenizer.train(corpus_texts, verbose=True)
else:
print(f" WARNING: Corpus not found at {corpus_path}")
print(f" Run 'python scripts/prepare_data.py' first!")
return
# Update config vocab size to match actual tokenizer
config.vocab_size = tokenizer.vocab_size
# Save tokenizer
tokenizer_dir = os.path.join(project_root, args.output_dir, "tokenizer")
tokenizer.save(tokenizer_dir)
# ── Step 3: Create Datasets ──────────────────────────────────
print(f"\n[3/6] Creating datasets...")
train_path = os.path.join(project_root, args.train_data)
eval_path = os.path.join(project_root, args.eval_data)
train_dataset = ViuAIDataset(train_path, tokenizer, max_seq_len=config.max_seq_len)
eval_dataset = ViuAIDataset(eval_path, tokenizer, max_seq_len=config.max_seq_len)
collator = ChatDataCollator(pad_token_id=config.pad_token_id)
train_loader = DataLoader(
train_dataset,
batch_size=args.batch_size,
shuffle=True,
collate_fn=collator,
num_workers=0,
)
eval_loader = DataLoader(
eval_dataset,
batch_size=args.batch_size,
shuffle=False,
collate_fn=collator,
num_workers=0,
)
print(f" Train batches: {len(train_loader)}")
print(f" Eval batches: {len(eval_loader)}")
# ── Step 4: Initialize Model ─────────────────────────────────
print(f"\n[4/6] Initializing ViuAI-1M model...")
model = ViuAIForCausalLM(config)
print(model.get_parameter_summary())
# ── Step 5: Train ────────────────────────────────────────────
print(f"\n[5/6] Starting training...")
checkpoint_dir = os.path.join(project_root, args.output_dir, "checkpoints")
trainer = ViuAITrainer(
model=model,
train_dataloader=train_loader,
eval_dataloader=eval_loader,
learning_rate=args.lr,
weight_decay=0.01,
warmup_steps=args.warmup_steps,
num_epochs=args.epochs,
save_dir=checkpoint_dir,
log_interval=10,
eval_interval=50,
device=args.device,
)
metrics = trainer.train()
# ── Step 6: Save Final Model ─────────────────────────────────
print(f"\n[6/6] Saving final model...")
final_dir = os.path.join(project_root, args.output_dir, "final")
os.makedirs(final_dir, exist_ok=True)
torch.save(model.state_dict(), os.path.join(final_dir, "model.pt"))
config.save(os.path.join(final_dir, "config.json"))
print(f" Model saved to: {final_dir}")
print(f"\n To chat with your model:")
print(f" python scripts/chat.py --checkpoint {final_dir} --tokenizer {tokenizer_dir}")
if __name__ == "__main__":
main()
|