""" Nova All-In-One Ultimate — Interactive Terminal GUI (TUI) Chat Application Features: - Sleek Rich-powered Terminal UI with live Markdown rendering - Model Architecture Selector (NVIDIA Nemotron-8B, Nemotron-70B, Nova 2.0 MoE, Nova 1.0 HRM) - Autonomous Python Tool Execution Engine visualizer - Dynamic Latent Reasoning Depth control (/depth) """ import os import sys import time import warnings warnings.filterwarnings("ignore") os.environ["PYTHONWARNINGS"] = "ignore" import torch # Add project root to sys.path sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) from rich.console import Console from rich.panel import Panel from rich.markdown import Markdown from rich.prompt import Prompt from rich.theme import Theme from rich.table import Table from rich.text import Text from config.model_config import Nova1Config from src.tokenizer.bpe_tokenizer import Nova1Tokenizer from src.model.nova1_hrm import Nova1HRM from src.model.nova2_moe_vision import Nova2MoEVision from src.memory.chat_memory import NovaChatMemory from src.inference.generate import Nova1Generator from src.inference.tool_engine import parse_and_execute_tools # Custom Rich Console Theme custom_theme = Theme({ "user": "bold cyan", "assistant": "bold magenta", "system": "bold yellow", "tool": "bold green", "dim": "dim white", "border": "purple", }) console = Console(theme=custom_theme) def print_banner(): banner_text = Text() banner_text.append("🌌 NOVA ALL-IN-ONE ULTIMATE", style="bold magenta") banner_text.append(" — Interactive Terminal Workspace\n", style="bold cyan") banner_text.append("⚡ AMD ROCm 96GB VRAM Accelerated • Persistent Memory • Sub-Second Latency", style="dim white") panel = Panel( banner_text, title="[bold green]Frontier AI Terminal GUI[/bold green]", subtitle="[dim]Type '/help' for commands | '/clear' to reset memory | '/exit' to quit[/dim]", border_style="cyan", padding=(1, 2) ) console.print(panel) def load_selected_model(choice: int = 1): root_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) device = "cuda" if torch.cuda.is_available() else "cpu" token = os.environ.get("HF_TOKEN") with console.status("[bold purple]Loading model & warming GPU compilation kernels...", spinner="dots"): if choice == 1: # Nova All-In-One (NVIDIA Nemotron-8B) from transformers import AutoTokenizer, AutoModelForCausalLM model_name = "nvidia/Mistral-NeMo-Minitron-8B-Instruct" tokenizer = AutoTokenizer.from_pretrained(model_name, token=token) model = AutoModelForCausalLM.from_pretrained(model_name, dtype=torch.bfloat16, device_map=device, token=token) generator = None name = "Nova All-In-One (NVIDIA Nemotron-8B Engine)" # One-time GPU compilation warmup pass for sub-second first prompt with torch.inference_mode(): dummy = tokenizer(["Hi"], return_tensors="pt").to(device) model.generate(**dummy, max_new_tokens=1, use_cache=True) elif choice == 2: # Nova 3.5 Nemotron-70B from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig model_name = "nvidia/Llama-3.1-Nemotron-70B-Instruct-HF" tokenizer = AutoTokenizer.from_pretrained(model_name, token=token) bnb_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_quant_type="nf4") model = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=bnb_config, device_map="auto", token=token) generator = None name = "Nova 3.5 Nemotron-70B (NVIDIA 70B Engine)" elif choice == 3: # Nova 2.0 MoE-Vision tok_path = os.path.join(root_dir, "checkpoints", "nova1_gemini_tokenizer.json") tokenizer = Nova1Tokenizer.load(tok_path) ckpt_path = os.path.join(root_dir, "checkpoints", "nova2_moe_multidomain.pt") if not os.path.exists(ckpt_path): ckpt_path = os.path.join(root_dir, "checkpoints", "nova2_moe_final.pt") checkpoint = torch.load(ckpt_path, map_location="cpu", weights_only=False) config = checkpoint.get("config", Nova1Config(vocab_size=tokenizer.vocab_size)) config.device = device model = Nova2MoEVision(config, num_experts=8, top_k=2).to(device) model.load_state_dict(checkpoint["model_state"], strict=False) generator = Nova1Generator(model, tokenizer, config) name = "Nova 2.0 MoE-Vision (~522M Params)" else: # Nova 1.0 HRM tok_path = os.path.join(root_dir, "checkpoints", "nova1_gemini_tokenizer.json") tokenizer = Nova1Tokenizer.load(tok_path) ckpt_path = os.path.join(root_dir, "checkpoints", "nova1_hle_exam.pt") checkpoint = torch.load(ckpt_path, map_location="cpu", weights_only=False) config = checkpoint.get("config", Nova1Config(vocab_size=tokenizer.vocab_size)) config.device = device model = Nova1HRM(config).to(device) model.load_state_dict(checkpoint["model_state"], strict=False) generator = Nova1Generator(model, tokenizer, config) name = "Nova 1.0 HRM (~125M Params)" console.print(f"[bold green]✔ Loaded {name} & Warmed GPU![/bold green]\n") return tokenizer, model, generator, name def run_terminal_gui(): console.clear() print_banner() # Parse initial choice from CLI arg or default to 1 (Nemotron-8B) choice = 1 if len(sys.argv) > 1 and sys.argv[1].isdigit(): choice = int(sys.argv[1]) elif len(sys.argv) > 2 and sys.argv[1] == "--model" and sys.argv[2].isdigit(): choice = int(sys.argv[2]) console.print(f"[bold yellow]Active Engine:[/bold yellow] [bold magenta]Nova All-In-One Ultimate (NVIDIA Nemotron-8B)[/bold magenta] (Choice #{choice})") tokenizer, model, generator, model_name = load_selected_model(choice) memory = NovaChatMemory() n_cycles, t_steps = 4, 4 console.print(Panel( Markdown(f"**Chat Session Started with {model_name}**\n\nPersistent Chat Memory: **Active ({len(memory.get_messages())} past turns loaded)**. Nova remembers conversation history across restarts!"), border_style="magenta" )) while True: try: user_input = Prompt.ask("\n[bold cyan]User[/bold cyan]").strip() if not user_input: continue if user_input.lower() in ["/exit", "/quit"]: console.print("[bold red]Exiting Nova Terminal Workspace. Chat history saved to disk. Goodbye![/bold red]") break if user_input.lower() == "/help": console.print("[bold yellow]Available Commands:[/bold yellow]") console.print(" /model - Switch active model engine") console.print(" /clear - Clear persistent conversation memory") console.print(" /exit - Quit terminal chat") continue if user_input.lower() == "/clear": memory.clear() console.clear() print_banner() console.print("[bold green]Persistent conversation memory cleared.[/bold green]") continue if user_input.lower() == "/model": choice_str = Prompt.ask("Select Model [1. Nemotron-8B | 2. Nemotron-70B | 3. MoE-522M | 4. HRM-125M]", default="1") choice = int(choice_str) if choice_str.isdigit() else 1 tokenizer, model, generator, model_name = load_selected_model(choice) continue start_t = time.time() memory.add_user_message(user_input) with console.status("[bold magenta]Nova is reasoning on AMD ROCm GPU...", spinner="earth"): if generator is None: # Foundation Model Generation with low-latency inference_mode & KV cache messages = [ {"role": "system", "content": "You are Nova, an advanced AI assistant powered by NVIDIA Nemotron. When asked who or what you are, state clearly: 'I am Nova, an advanced AI assistant powered by NVIDIA Nemotron.' Be direct, concise, and precise. Answer simple questions with short, exact answers without unnecessary conversational filler or fluff. Provide code or detailed step-by-step explanations ONLY when explicitly requested or necessary. If python code execution is needed, output ```python code here ```."}, *memory.get_messages() ] text_inp = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) is_code_request = any(k in user_input.lower() for k in ["code", "script", "python", "function", "write", "calculate", "program"]) max_tokens = 250 if is_code_request else 32 temp = 0.6 if is_code_request else 0.1 with torch.inference_mode(): inputs = tokenizer([text_inp], return_tensors="pt").to("cuda") out_ids = model.generate( **inputs, max_new_tokens=max_tokens, temperature=temp, do_sample=is_code_request, use_cache=True, pad_token_id=tokenizer.eos_token_id ) out_ids = [out[len(inp):] for inp, out in zip(inputs.input_ids, out_ids)] raw_response = tokenizer.batch_decode(out_ids, skip_special_tokens=True)[0] for stop_seq in ["", "\nUser:", "\nAssistant:"]: if stop_seq in raw_response: raw_response = raw_response.split(stop_seq)[0] raw_response = raw_response.strip() else: # Custom Nova Generator formatted_prompt = f"User: {user_input}\nAssistant:" with torch.inference_mode(): raw_response = generator.generate( prompt=formatted_prompt, max_new_tokens=100, temperature=0.7, N_cycles=n_cycles, T_steps=t_steps ) raw_response = raw_response.replace(formatted_prompt, "").strip() latency = time.time() - start_t # Execute Autonomous Tool Engine if code block present clean_text, tool_results = parse_and_execute_tools(raw_response) memory.add_assistant_message(clean_text) # Render Assistant Output Panel response_panel = Panel( Markdown(clean_text), title=f"[bold magenta]{model_name}[/bold magenta]", subtitle=f"[dim]⚡ Latency: {latency:.2f}s | Device: AMD ROCm GPU[/dim]", border_style="purple", padding=(1, 2) ) console.print(response_panel) if tool_results: for res in tool_results: tool_panel = Panel( Text(res, style="bold green"), title="[bold green]⚡ [Autonomous Tool Output][/bold green]", border_style="green" ) console.print(tool_panel) except KeyboardInterrupt: console.print("\n[bold red]Interrupted. Type /exit to quit.[/bold red]") except Exception as e: console.print(f"[bold red]Error during generation: {e}[/bold red]") if __name__ == "__main__": run_terminal_gui()