#!/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()