| """ |
| 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 |
|
|
| |
| 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") |
|
|
| |
| 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:,}") |
|
|
| |
| 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 |
|
|
| |
| config.vocab_size = tokenizer.vocab_size |
|
|
| |
| tokenizer_dir = os.path.join(project_root, args.output_dir, "tokenizer") |
| tokenizer.save(tokenizer_dir) |
|
|
| |
| 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)}") |
|
|
| |
| print(f"\n[4/6] Initializing ViuAI-1M model...") |
| model = ViuAIForCausalLM(config) |
| print(model.get_parameter_summary()) |
|
|
| |
| 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() |
|
|
| |
| 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() |
|
|
