File size: 9,877 Bytes
d7228c8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 | """
Command Line Interface (CLI) for training, generation, tokenization, and benchmarking.
"""
import argparse
import os
import sys
from typing import List
import torch
from slm.config.config_loader import load_config, save_config
from slm.config.model_config import ModelConfig
from slm.config.train_config import TrainConfig
from slm.dataset.cleaner import DataCleaner
from slm.dataset.readers import DatasetReader
from slm.dataset.dataset import CausalLMDataset
from slm.dataset.loader import create_dataloader
from slm.tokenizer.bpe import BPETokenizer
from slm.model.transformer_lm import SLMForCausalLM
from slm.sampling.generator import TextGenerator
from slm.training.trainer import Trainer
from slm.evaluation.benchmark import benchmark_inference, benchmark_training_step
from slm.checkpoint.manager import CheckpointManager
from slm.utils.logger import get_logger
from slm.utils.utils import set_seed, get_device
logger = get_logger("slm.cli")
def cmd_train(args: argparse.Namespace) -> None:
"""Handles 'train' CLI command."""
logger.info(f"Loading configuration from {args.config}...")
model_config, train_config = load_config(args.config)
if args.dataset:
train_config.train_dataset_path = args.dataset
set_seed(train_config.seed)
# Ingest text documents
if train_config.train_dataset_path and os.path.exists(train_config.train_dataset_path):
docs = DatasetReader.load_file(train_config.train_dataset_path)
else:
logger.info("No training dataset path supplied. Using synthetic demonstration corpus.")
docs = [
"Deep learning allows computational models composed of multiple processing layers to learn representations of data with multiple levels of abstraction.",
"Transformer architectures rely entirely on self-attention mechanisms to compute representations of their input and output without using sequence-aligned RNNs.",
"Byte Pair Encoding subword tokenization enables vocabulary scaling across arbitrary natural language text corpora.",
] * 100
# Train or load BPE tokenizer
tokenizer = BPETokenizer()
tokenizer.train_on_texts(docs, vocab_size=model_config.vocab_size)
# Build dataset and dataloader
dataset = CausalLMDataset(docs, tokenizer, max_seq_len=model_config.max_seq_len)
loader = create_dataloader(dataset, batch_size=train_config.batch_size, shuffle=True)
# Build model and trainer
model = SLMForCausalLM(model_config)
trainer = Trainer(model, train_config, loader, tokenizer=tokenizer)
# Execute training loop
summary = trainer.train()
logger.info(f"Training finished successfully! Final checkpoint saved to: {summary['final_checkpoint']}")
def cmd_generate(args: argparse.Namespace) -> None:
"""Handles 'generate' CLI command."""
logger.info(f"Initializing text generator for prompt: '{args.prompt}'")
if args.checkpoint and os.path.exists(args.checkpoint):
ckpt_manager = CheckpointManager(output_dir=os.path.dirname(args.checkpoint))
# Inspect checkpoint to get saved model_config if present
try:
ckpt_data = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
except Exception:
ckpt_data = torch.load(args.checkpoint, map_location="cpu")
if isinstance(ckpt_data, dict) and "model_config" in ckpt_data:
config = ModelConfig.from_dict(ckpt_data["model_config"])
else:
config = ModelConfig(vocab_size=1000, d_model=128, n_heads=4, n_layers=2)
model = SLMForCausalLM(config)
ckpt_manager.load_checkpoint(args.checkpoint, model)
tokenizer = BPETokenizer.load(os.path.join(os.path.dirname(args.checkpoint), "tokenizer"))
else:
logger.info("No checkpoint provided. Initializing lightweight untrained SLM model for demo.")
config = ModelConfig(vocab_size=1000, d_model=128, n_heads=4, n_layers=2, d_ff=512)
model = SLMForCausalLM(config)
tokenizer = BPETokenizer()
tokenizer.train_on_texts([args.prompt, "Language model generation testing completely from scratch."], vocab_size=1000)
generator = TextGenerator(model, tokenizer)
output = generator.generate(
prompt=args.prompt,
max_new_tokens=args.max_tokens,
temperature=args.temperature,
top_k=args.top_k,
top_p=args.top_p
)
print("\n" + "=" * 50)
print(" GENERATED TEXT OUTPUT")
print("=" * 50)
print(output)
print("=" * 50 + "\n")
def cmd_tokenize(args: argparse.Namespace) -> None:
"""Handles 'tokenize' CLI command."""
tokenizer = BPETokenizer()
tokenizer.train_on_texts([args.text], vocab_size=500)
token_ids = tokenizer.encode(args.text)
decoded = tokenizer.decode(token_ids)
print("\n" + "=" * 50)
print(f" Input Text: {args.text}")
print(f" Token IDs: {token_ids}")
print(f" Decoded Text: {decoded}")
print("=" * 50 + "\n")
def cmd_benchmark(args: argparse.Namespace) -> None:
"""Handles 'benchmark' CLI command."""
device = get_device(args.device)
logger.info(f"Running SLM Benchmark on device={device}...")
config = ModelConfig(vocab_size=10000, d_model=256, n_heads=8, n_layers=4, d_ff=1024)
model = SLMForCausalLM(config)
res = benchmark_inference(model, batch_size=args.batch_size, seq_len=args.seq_len, device=device)
print("\n" + "=" * 50)
print(" SLM INFERENCE BENCHMARK RESULTS")
print("=" * 50)
print(f" Device: {res['device']}")
print(f" Batch Size: {res['batch_size']}")
print(f" Sequence Length: {res['seq_len']}")
print(f" Latency (ms): {res['avg_latency_ms']} ms")
print(f" Throughput: {res['tokens_per_sec']} tokens/sec")
print("=" * 50 + "\n")
def cmd_chat(args: argparse.Namespace) -> None:
"""Handles 'chat' CLI command."""
from scripts.chat_run import start_chat
start_chat(args.checkpoint)
def cmd_pipeline(args: argparse.Namespace) -> None:
"""Handles 'pipeline' CLI command for end-to-end automated execution."""
from slm.pipeline import PipelineOrchestrator
docs = [
"User: hello\nSLM: Hello! How can I assist you with language modeling today?\n",
"User: What is law?\nSLM: Law is a system of rules created and enforced by institutions.\n",
"User: What is Python?\nSLM: Python is a high-level programming language.\n"
] * 20
orchestrator = PipelineOrchestrator(output_dir=args.output_dir)
orchestrator.run_full_pipeline(
raw_documents=docs,
vocab_size=args.vocab_size,
max_steps=args.max_steps,
epochs=args.epochs,
batch_size=args.batch_size,
device=args.device
)
def main() -> None:
"""Main CLI entry point."""
parser = argparse.ArgumentParser(description="SLM CLI - Small Language Model Tools")
subparsers = parser.add_subparsers(dest="command", required=True)
# Pipeline command
pipe_parser = subparsers.add_parser("pipeline", help="Run end-to-end automated SLM pipeline")
pipe_parser.add_argument("--output_dir", type=str, default="checkpoints_pipeline", help="Output directory")
pipe_parser.add_argument("--vocab_size", type=int, default=1000, help="Vocabulary size")
pipe_parser.add_argument("--max_steps", type=int, default=100, help="Max training steps")
pipe_parser.add_argument("--epochs", type=int, default=20, help="Epochs")
pipe_parser.add_argument("--batch_size", type=int, default=4, help="Batch size")
pipe_parser.add_argument("--device", type=str, default="cpu", help="Device (cpu/cuda)")
# Train command
train_parser = subparsers.add_parser("train", help="Train model")
train_parser.add_argument("--config", type=str, default="configs/nano_config.yaml", help="Path to config file")
train_parser.add_argument("--dataset", type=str, default=None, help="Path to training dataset file")
# Generate command
gen_parser = subparsers.add_parser("generate", help="Generate text")
gen_parser.add_argument("--prompt", type=str, required=True, help="Input prompt string")
gen_parser.add_argument("--checkpoint", type=str, default=None, help="Path to model checkpoint")
gen_parser.add_argument("--max_tokens", type=int, default=64, help="Max tokens to generate")
gen_parser.add_argument("--temperature", type=float, default=0.8, help="Sampling temperature")
gen_parser.add_argument("--top_k", type=int, default=40, help="Top-K limit")
gen_parser.add_argument("--top_p", type=float, default=0.9, help="Top-P nucleus threshold")
# Chat command
chat_parser = subparsers.add_parser("chat", help="Start interactive chat REPL")
chat_parser.add_argument("--checkpoint", type=str, default=None, help="Path to model checkpoint")
# Tokenize command
tok_parser = subparsers.add_parser("tokenize", help="Tokenize text")
tok_parser.add_argument("--text", type=str, required=True, help="Text to tokenize")
# Benchmark command
bench_parser = subparsers.add_parser("benchmark", help="Run benchmark")
bench_parser.add_argument("--batch_size", type=int, default=4, help="Batch size")
bench_parser.add_argument("--seq_len", type=int, default=256, help="Sequence context length")
bench_parser.add_argument("--device", type=str, default="auto", help="Device choice")
args = parser.parse_args()
if args.command == "pipeline":
cmd_pipeline(args)
elif args.command == "train":
cmd_train(args)
elif args.command == "generate":
cmd_generate(args)
elif args.command == "chat":
cmd_chat(args)
elif args.command == "tokenize":
cmd_tokenize(args)
elif args.command == "benchmark":
cmd_benchmark(args)
if __name__ == "__main__":
main()
|