Nova / scripts /terminal_gui.py

Upload folder using huggingface_hub

46f6126 verified 2 months ago

12.1 kB

"""
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 ["<extra_id_1>", "\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()

	"""
	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 ["<extra_id_1>", "\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()