src/main.py

#!/usr/bin/env python3
"""
Frox AI Morph 1.1 — Main Entry Point

Modes:
    python main.py info                    Print architecture summary for every model-family tier
    python main.py build --family nano      Build a fresh model + tokenizer, verify it runs
    python main.py demo  --family nano      Build + run a tiny smoke-test generation
    python main.py chat  --model PATH       Launch interactive chat (delegates to scripts/chat.py)
    python main.py train ...                Launch training (delegates to scripts/train.py)

Model family: nano / mini / classic / pro / code — each defined in its
own standalone file under config/family/. This entry point loads
exactly one tier per invocation via config.family.load, so running
`--family nano` never imports pro.py or code.py.
"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))

from utils.common import (
    print_banner, set_seed, get_device, describe_device, detect_environment,
    load_family_config, FAMILY_TIERS,
)


def cmd_info(args):
    print_banner()
    print(f"Environment: {detect_environment()}")
    device = get_device()
    print(f"Device: {describe_device(device)}\n")

    for tier in FAMILY_TIERS:
        cfg, module = load_family_config(tier)
        t = cfg.text
        # Rough parameter count estimate
        embed = t.total_vocab_size * t.hidden_size
        per_layer = (
            4 * t.hidden_size * t.hidden_size  # attn q/k/v/o (approx, ignoring GQA ratio)
            + 3 * t.hidden_size * t.intermediate_size  # SwiGLU gate/up/down
        )
        total_approx = embed + per_layer * t.num_hidden_layers

        name = getattr(module, "MODEL_NAME", tier.title())
        desc = getattr(module, "DESCRIPTION", "")
        hw = getattr(module, "RECOMMENDED_HARDWARE", "")

        print(f"── {name} ({tier}) ──")
        print(f"   {desc}")
        print(f"   hidden={t.hidden_size}  layers={t.num_hidden_layers}  "
              f"heads={t.num_attention_heads}/{t.num_key_value_heads} (GQA {t.num_attention_heads//t.num_key_value_heads}:1)")
        print(f"   context={t.max_position_embeddings:,} (YaRN {t.rope_scaling_factor}x)  "
              f"vocab={t.total_vocab_size:,}  qk_norm={t.qk_norm}"
              + ("  fim=True" if getattr(t, "code_fim_enabled", False) else ""))
        print(f"   ~{total_approx/1e9:.2f}B parameters (rough estimate, embed+layers only)")
        print(f"   hardware: {hw}\n")


def cmd_build(args):
    from model.architecture.morph_model import MorphForCausalLM
    from tokenizer.morph_tokenizer import build_morph_tokenizer

    print_banner()
    set_seed(args.seed)

    config, module = load_family_config(args.family)
    name = getattr(module, "MODEL_NAME", args.family.title())

    print(f"Building {name}...")
    tokenizer = build_morph_tokenizer()
    model = MorphForCausalLM(config.text)

    params = model.param_count()
    print(f"\n✅ Built successfully: {params['total_billions']}B parameters")

    # Sanity forward pass — sample ids across the FULL vocab (base + special
    # tokens) so the special-token embedding rows are actually exercised;
    # otherwise an embedding table sized to vocab_size instead of
    # total_vocab_size would still pass this smoke test.
    import torch
    dummy = torch.randint(0, config.text.total_vocab_size, (1, 16))
    with torch.no_grad():
        out = model(input_ids=dummy)
    assert out.logits.shape == (1, 16, config.text.total_vocab_size), "Shape mismatch!"
    assert not torch.isnan(out.logits).any(), "NaN in output logits!"
    print(f"✅ Forward pass OK — logits shape {tuple(out.logits.shape)}, no NaNs")

    if args.save:
        model.save(args.save)
        tokenizer.save_pretrained(f"{args.save}/tokenizer")
        print(f"✅ Saved to {args.save}")


def cmd_demo(args):
    from multimodal.fusion.morph_multimodal import MorphMultimodalModel
    from tokenizer.morph_tokenizer import build_morph_tokenizer
    from inference.engine.morph_engine import MorphInferenceEngine

    print_banner()
    set_seed(args.seed)

    config, module = load_family_config(args.family)
    name = getattr(module, "MODEL_NAME", args.family.title())

    print(f"Building an UNTRAINED {name} for a smoke test...")
    print("(Output will be random noise — this only verifies the pipeline runs end-to-end.)\n")

    tokenizer = build_morph_tokenizer()
    model = MorphMultimodalModel(config)

    device = get_device(args.device)
    engine = MorphInferenceEngine(model=model, tokenizer=tokenizer, config=config, device=device)

    response = engine.generate(
        [{"role": "user", "content": "Hello! Tell me about yourself."}],
        max_new_tokens=32,
    )
    print(f"\nRaw output (untrained, expect gibberish): {response!r}")
    print("\n✅ Demo complete — pipeline is wired correctly end-to-end.")


def cmd_chat(args):
    from scripts.chat import main as chat_main
    sys.argv = ["chat.py", "--model", args.model]
    if args.quantization:
        sys.argv += ["--quantization", args.quantization]
    chat_main()


def cmd_train(args):
    from scripts.train import main as train_main
    sys.argv = ["train.py"] + args.train_args
    train_main()


def main():
    parser = argparse.ArgumentParser(description="Frox AI Morph 1.1")
    sub = parser.add_subparsers(dest="command", required=True)

    p_info = sub.add_parser("info", help="Print architecture summary for every tier")

    p_build = sub.add_parser("build", help="Build + sanity-check a fresh model")
    p_build.add_argument("--family", choices=list(FAMILY_TIERS), default="nano")
    p_build.add_argument("--seed", type=int, default=1337)
    p_build.add_argument("--save", type=str, default=None)

    p_demo = sub.add_parser("demo", help="Build + run a tiny smoke-test generation")
    p_demo.add_argument("--family", choices=list(FAMILY_TIERS), default="nano")
    p_demo.add_argument("--seed", type=int, default=1337)
    p_demo.add_argument("--device", type=str, default=None)

    p_chat = sub.add_parser("chat", help="Interactive chat with a trained model")
    p_chat.add_argument("--model", type=str, required=True)
    p_chat.add_argument("--quantization", choices=["4bit", "8bit"], default=None)

    p_train = sub.add_parser("train", help="Train (forwards args to scripts/train.py)")
    p_train.add_argument("train_args", nargs=argparse.REMAINDER)

    args = parser.parse_args()

    {
        "info":  cmd_info,
        "build": cmd_build,
        "demo":  cmd_demo,
        "chat":  cmd_chat,
        "train": cmd_train,
    }[args.command](args)


if __name__ == "__main__":
    main()