slm/cli/main.py

"""
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()