app.py

from __future__ import annotations

import os
import random
import re
import sys
from contextlib import nullcontext
from pathlib import Path
from typing import Any, Generator

import gradio as gr
import torch
from huggingface_hub import snapshot_download


# =============================================================================
# Model configuration
# =============================================================================

MODEL_ID = os.getenv(
    "MODEL_ID",
    "MarkChenX/lfm2-quantum-128m-sft-v2-reasoning",
)
HF_TOKEN = os.getenv("HF_TOKEN")

MAX_CONTEXT_TOKENS = 1024

# Set FORCE_HF_DOWNLOAD=1 temporarily after replacing bad files on the Hub.
# Leave it unset/0 normally so a Space restart does not redownload ~445 MB.
FORCE_HF_DOWNLOAD = os.getenv("FORCE_HF_DOWNLOAD", "0").strip().lower() in {
    "1",
    "true",
    "yes",
    "on",
}

BUILT_IN_CONTEXT = """You are Supernova 1 Beta, an experimental custom
LFM2-style hybrid quantum language model running through NanoChat.

Behavior:
- Answer the user's request directly.
- Be concise, clear, and technically precise.
- For research or coding questions, explain the key reasoning in short steps.
- If you are uncertain, say what is uncertain instead of inventing facts.
- Prefer useful answers over long introductions.
- Use Markdown when it improves readability.
"""

# These examples are injected into the actual model prompt.
# Keep them short because the checkpoint has a 1,024-token context window.
FEW_SHOT_MESSAGES: list[dict[str, str]] = [
    {
        "role": "user",
        "content": "What are you?",
    },
    {
        "role": "assistant",
        "content": (
            "I am Supernova 1 Beta, an experimental transformer model "
            "with quantum neural networks and an LFM2-style hybrid "
            "quantum language model running through NanoChat."
        ),
    },
]

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


# =============================================================================
# Helpers used during startup
# =============================================================================

GIT_LFS_POINTER_PREFIX = b"version https://git-lfs.github.com/spec/v1"


def is_git_lfs_pointer(path: Path) -> bool:
    """Return True when a downloaded file is text for a Git LFS pointer."""
    if not path.is_file():
        return False

    try:
        with path.open("rb") as f:
            header = f.read(256)
    except OSError:
        return False

    return header.startswith(GIT_LFS_POINTER_PREFIX)


def find_git_lfs_pointers(directory: Path) -> list[Path]:
    """Find unresolved Git LFS pointer files anywhere under a directory."""
    if not directory.exists():
        return []

    bad_files: list[Path] = []

    for path in directory.rglob("*"):
        if path.is_file() and is_git_lfs_pointer(path):
            bad_files.append(path)

    return bad_files


def validate_tokenizer_directory(tokenizer_dir: Path) -> None:
    """
    Validate the tokenizer payload before NanoChat calls pickle.load().

    This specifically catches the failure:
        _pickle.UnpicklingError: invalid load key, 'v'

    That error commonly occurs when the file being unpickled starts with:
        version https://git-lfs.github.com/spec/v1
    """
    if not tokenizer_dir.exists():
        raise FileNotFoundError(
            "Bundled tokenizer directory is missing.\n"
            f"Expected: {tokenizer_dir}\n\n"
            "Upload the exact tokenizer used to train this checkpoint "
            "under the checkpoint's tokenizer/ directory."
        )

    tokenizer_files = [
        path
        for path in tokenizer_dir.rglob("*")
        if path.is_file()
    ]

    if not tokenizer_files:
        raise RuntimeError(
            "Bundled tokenizer directory exists but contains no files.\n"
            f"Directory: {tokenizer_dir}"
        )

    print("Tokenizer directory:", tokenizer_dir)
    print("Tokenizer files:")

    for path in sorted(tokenizer_files):
        relative = path.relative_to(tokenizer_dir)
        try:
            size = path.stat().st_size
        except OSError:
            size = -1
        print(f"  - {relative} ({size:,} bytes)")

    lfs_pointers = find_git_lfs_pointers(tokenizer_dir)

    if lfs_pointers:
        formatted = "\n".join(
            f"  - {path.relative_to(tokenizer_dir)}"
            for path in lfs_pointers
        )

        raise RuntimeError(
            "\n"
            "The bundled tokenizer contains unresolved Git LFS pointer "
            "file(s), not the real tokenizer binary:\n"
            f"{formatted}\n\n"
            "This is why pickle.load() reports:\n"
            "    _pickle.UnpicklingError: invalid load key, 'v'\n\n"
            "Fix the MODEL repository, not pickle.load():\n"
            "1. Recover the exact tokenizer files used during training.\n"
            "2. Confirm the local files are real binaries and do not begin "
            "with 'version https://git-lfs.github.com/spec/v1'.\n"
            "3. Re-upload the tokenizer/ directory to Hugging Face.\n"
            "4. Restart the Space. You may set FORCE_HF_DOWNLOAD=1 for the "
            "first restart after replacing the files.\n\n"
            "Do NOT substitute an arbitrary AutoTokenizer unless it has "
            "the exact same token-to-id mapping as the training tokenizer."
        )


def find_latest_checkpoint(directory: Path) -> tuple[Path, int]:
    candidates: list[tuple[int, Path]] = []

    for path in directory.rglob("model_*.pt"):
        match = re.fullmatch(r"model_(\d+)\.pt", path.name)
        if match:
            candidates.append((int(match.group(1)), path))

    if not candidates:
        raise FileNotFoundError(
            f"No model_XXXXXX.pt checkpoint found under {directory}"
        )

    candidates.sort(key=lambda item: item[0], reverse=True)
    step, checkpoint_path = candidates[0]
    return checkpoint_path, step


# =============================================================================
# Download repository
# =============================================================================

print("=" * 80)
print(f"Downloading model repository: {MODEL_ID}")
print(f"Force Hub download: {FORCE_HF_DOWNLOAD}")
print(f"Requested runtime device: {device}")
print("=" * 80)

model_dir = Path(
    snapshot_download(
        repo_id=MODEL_ID,
        token=HF_TOKEN,
        force_download=FORCE_HF_DOWNLOAD,
        ignore_patterns=[
            "optim_*.pt",
            "optimizer*.pt",
            "*.bin",
        ],
    )
)

print(f"Model repository: {model_dir}")

# Catch unresolved LFS pointers anywhere in the repository early.
repo_lfs_pointers = find_git_lfs_pointers(model_dir)
if repo_lfs_pointers:
    print("WARNING: unresolved Git LFS pointer files found in snapshot:")
    for path in repo_lfs_pointers:
        print(f"  - {path.relative_to(model_dir)}")

# Make the bundled NanoChat package importable.
sys.path.insert(0, str(model_dir))


# =============================================================================
# Import custom architecture
# =============================================================================

try:
    from nanochat.checkpoint_manager import build_model
    from nanochat.engine import Engine
    from nanochat.reasoning import (
        DEFAULT_BANDS,
        ON_HIGH,
        generate_with_budget,
        with_condition,
    )
except ImportError as error:
    raise RuntimeError(
        "Could not import the bundled NanoChat code. The model repository "
        "must contain nanochat/checkpoint_manager.py, nanochat/engine.py, "
        "and nanochat/reasoning.py."
    ) from error


# =============================================================================
# Locate latest model checkpoint
# =============================================================================

checkpoint_path, checkpoint_step = find_latest_checkpoint(model_dir)
checkpoint_dir = checkpoint_path.parent

metadata_path = checkpoint_dir / f"meta_{checkpoint_step:06d}.json"

if not metadata_path.exists():
    raise FileNotFoundError(
        "Checkpoint metadata is missing. Expected: "
        f"{metadata_path}"
    )

tokenizer_dir = checkpoint_dir / "tokenizer"

print(f"Checkpoint directory: {checkpoint_dir}")
print(f"Checkpoint file: {checkpoint_path}")
print(f"Checkpoint step: {checkpoint_step}")
print(f"Metadata file: {metadata_path}")
print(f"Tokenizer directory: {tokenizer_dir}")
print(f"Loading on device: {device}")

# Validate BEFORE NanoChat reaches pickle.load().
validate_tokenizer_directory(tokenizer_dir)


# =============================================================================
# Load custom NanoChat model
# =============================================================================

try:
    model, tokenizer, metadata = build_model(
        str(checkpoint_dir),
        checkpoint_step,
        device,
        "eval",
    )
except Exception as error:
    raise RuntimeError(
        "\nFailed to load the NanoChat checkpoint.\n"
        f"Checkpoint: {checkpoint_path}\n"
        f"Tokenizer: {tokenizer_dir}\n"
        f"Underlying error: {type(error).__name__}: {error}"
    ) from error

model.eval()
engine = Engine(model, tokenizer)

MODEL_CONTEXT_TOKENS = int(
    getattr(model.config, "sequence_len", MAX_CONTEXT_TOKENS)
)

HIGH_REASONING_SLACK = 1.25

bos_token = tokenizer.get_bos_token_id()
user_start_token = tokenizer.encode_special("<|user_start|>")
user_end_token = tokenizer.encode_special("<|user_end|>")
assistant_start_token = tokenizer.encode_special("<|assistant_start|>")
assistant_end_token = tokenizer.encode_special("<|assistant_end|>")

if device.type == "cuda":
    def autocast_context():
        return torch.amp.autocast(
            device_type="cuda",
            dtype=torch.bfloat16,
        )
else:
    autocast_context = nullcontext

print("Model and NanoChat engine loaded successfully.")
print(
    f"Model context length: {MODEL_CONTEXT_TOKENS} tokens · "
    f"Device: {device.type.upper()}"
)


# =============================================================================
# Conversation helpers
# =============================================================================

def normalize_history(
    history: list[Any] | None,
) -> list[dict[str, str]]:
    """Normalize Gradio chat history into role/content dictionaries."""
    messages: list[dict[str, str]] = []

    for item in history or []:
        if isinstance(item, dict):
            role = item.get("role")
            content = item.get("content")

            if (
                role in {"user", "assistant"}
                and isinstance(content, str)
                and content.strip()
            ):
                messages.append(
                    {
                        "role": role,
                        "content": content.strip(),
                    }
                )

        elif isinstance(item, (tuple, list)) and len(item) == 2:
            user_content, assistant_content = item

            if isinstance(user_content, str) and user_content.strip():
                messages.append(
                    {
                        "role": "user",
                        "content": user_content.strip(),
                    }
                )

            if (
                isinstance(assistant_content, str)
                and assistant_content.strip()
            ):
                messages.append(
                    {
                        "role": "assistant",
                        "content": assistant_content.strip(),
                    }
                )

    return messages


def encode_chat_message(role: str, content: str) -> list[int]:
    content_tokens = tokenizer.encode(content)

    if role == "user":
        return [
            user_start_token,
            *content_tokens,
            user_end_token,
        ]

    if role == "assistant":
        return [
            assistant_start_token,
            *content_tokens,
            assistant_end_token,
        ]

    raise ValueError(f"Unsupported role: {role}")


def group_history_into_turns(
    messages: list[dict[str, str]],
) -> list[list[dict[str, str]]]:
    """
    Group history into user-led turns so old context can be removed cleanly
    without slicing through a message boundary.
    """
    turns: list[list[dict[str, str]]] = []
    current: list[dict[str, str]] = []

    for chat_message in messages:
        if chat_message["role"] == "user":
            if current:
                turns.append(current)
            current = [chat_message]
        elif current:
            current.append(chat_message)

    if current:
        turns.append(current)

    return turns


def trim_tokens_from_left(
    tokens: list[int],
    max_length: int,
) -> list[int]:
    """Keep the newest token suffix, safely handling max_length <= 0."""
    if max_length <= 0:
        return []

    if len(tokens) <= max_length:
        return tokens

    return tokens[-max_length:]


def build_current_user_tokens(
    message: str,
    system_message: str,
    max_tokens: int,
    high_reasoning: bool,
) -> list[int]:
    """
    Build the newest user turn while preserving the behavior instruction.

    The user request is prioritized over optional instructions when the
    1,024-token context window becomes tight.
    """
    instruction = BUILT_IN_CONTEXT.strip()

    if system_message.strip():
        instruction += (
            "\n\nAdditional instruction from the user:\n"
            + system_message.strip()
        )

    if high_reasoning:
        conditioned = with_condition(
            {
                "messages": [
                    {
                        "role": "system",
                        "content": instruction,
                    },
                    {
                        "role": "user",
                        "content": message.strip(),
                    },
                ]
            },
            ON_HIGH,
        )

        # with_condition() adds the exact reasoning directive expected by
        # this reasoning-trained runtime.
        instruction = conditioned["messages"][0]["content"]

    prefix = (
        "[INSTRUCTIONS]\n"
        + instruction
        + "\n\n[CURRENT USER REQUEST]\n"
    )

    prefix_tokens = tokenizer.encode(prefix)
    message_tokens = tokenizer.encode(message.strip())

    wrapper_cost = 2  # user_start + user_end
    payload_budget = max(0, max_tokens - wrapper_cost)

    # Prioritize preserving the newest user request.
    if len(prefix_tokens) + len(message_tokens) > payload_budget:
        minimum_message_budget = min(
            len(message_tokens),
            max(32, payload_budget // 2),
        )

        prefix_budget = max(
            0,
            payload_budget - minimum_message_budget,
        )

        # Keep the beginning of the behavioral instruction.
        prefix_tokens = prefix_tokens[:prefix_budget]

        remaining_for_message = max(
            0,
            payload_budget - len(prefix_tokens),
        )

        message_tokens = trim_tokens_from_left(
            message_tokens,
            remaining_for_message,
        )

    return [
        user_start_token,
        *prefix_tokens,
        *message_tokens,
        user_end_token,
    ]


def build_conversation_tokens(
    message: str,
    history: list[Any] | None,
    system_message: str,
    max_new_tokens: int,
    use_few_shot: bool,
    high_reasoning: bool,
) -> tuple[list[int], dict[str, int]]:
    """
    Build a NanoChat-native prompt with:
      1. optional few-shot demonstrations,
      2. as much recent conversation history as fits,
      3. built-in behavior + newest user request,
      4. optional ON_HIGH reasoning directive,
      5. assistant-start token for generation.
    """
    if high_reasoning:
        high_think_cap = int(
            DEFAULT_BANDS[ON_HIGH.effort][1]
            * HIGH_REASONING_SLACK
        )
        generation_reserve = high_think_cap + max_new_tokens
    else:
        high_think_cap = 0
        generation_reserve = max_new_tokens

    maximum_prompt_length = (
        MODEL_CONTEXT_TOKENS - generation_reserve
    )

    # Keep a minimally useful prompt even if the configured reasoning
    # reserve is large relative to this small model's context.
    maximum_prompt_length = max(
        64,
        min(
            maximum_prompt_length,
            MODEL_CONTEXT_TOKENS - 16,
        ),
    )

    # Reserve BOS + assistant_start outside this helper.
    current_user_tokens = build_current_user_tokens(
        message=message,
        system_message=system_message,
        max_tokens=max(8, maximum_prompt_length - 2),
        high_reasoning=high_reasoning,
    )

    mandatory = [
        bos_token,
        *current_user_tokens,
        assistant_start_token,
    ]

    if len(mandatory) > maximum_prompt_length:
        # Preserve BOS, the newest end of the current user payload, and
        # assistant_start instead of blindly clipping the end marker.
        body_budget = max(0, maximum_prompt_length - 2)
        current_body = trim_tokens_from_left(
            current_user_tokens,
            body_budget,
        )

        clipped = [
            bos_token,
            *current_body,
            assistant_start_token,
        ]

        return clipped, {
            "prompt_tokens": len(clipped),
            "history_turns": 0,
            "few_shot_messages": 0,
            "reasoning_reserve": high_think_cap,
        }

    history_messages = normalize_history(history)
    history_turns = group_history_into_turns(history_messages)

    few_shot_messages = (
        list(FEW_SHOT_MESSAGES)
        if use_few_shot
        else []
    )

    def encode_messages(
        messages: list[dict[str, str]],
    ) -> list[int]:
        tokens: list[int] = []

        for chat_message in messages:
            tokens.extend(
                encode_chat_message(
                    chat_message["role"],
                    chat_message["content"],
                )
            )

        return tokens

    few_shot_tokens = encode_messages(few_shot_messages)

    while (
        few_shot_messages
        and len(mandatory) + len(few_shot_tokens)
        > maximum_prompt_length
    ):
        # Drop the oldest user/assistant demonstration pair.
        few_shot_messages = few_shot_messages[2:]
        few_shot_tokens = encode_messages(
            few_shot_messages
        )

    selected_turns: list[list[dict[str, str]]] = []
    running_length = (
        len(mandatory) + len(few_shot_tokens)
    )

    for turn in reversed(history_turns):
        turn_tokens = encode_messages(turn)

        if (
            running_length + len(turn_tokens)
            <= maximum_prompt_length
        ):
            selected_turns.append(turn)
            running_length += len(turn_tokens)
        else:
            break

    selected_turns.reverse()

    context_tokens: list[int] = [bos_token]
    context_tokens.extend(few_shot_tokens)

    for turn in selected_turns:
        context_tokens.extend(
            encode_messages(turn)
        )

    context_tokens.extend(current_user_tokens)
    context_tokens.append(assistant_start_token)

    return context_tokens, {
        "prompt_tokens": len(context_tokens),
        "history_turns": len(selected_turns),
        "few_shot_messages": len(few_shot_messages),
        "reasoning_reserve": high_think_cap,
    }


# =============================================================================
# Generation
# =============================================================================

def respond(
    message: str,
    history: list[Any] | None,
    custom_instruction: str,
    max_new_tokens: int,
    temperature: float,
    top_k: int,
    use_few_shot: bool,
    high_reasoning: bool,
) -> Generator[
    tuple[str, list[dict[str, str]], str],
    None,
    None,
]:
    """
    Standard:
        Streams directly from engine.generate().

    High reasoning:
        Conditions the prompt with ON_HIGH and calls
        generate_with_budget(). Only the parsed final answer is shown.
    """
    if not message or not message.strip():
        yield (
            "",
            normalize_history(history),
            "Enter a message to begin.",
        )
        return

    clean_message = message.strip()

    # `history` supplied by Gradio is the history BEFORE the current Textbox
    # message is added, so append it to the display history only.
    display_history = normalize_history(history)
    display_history.append(
        {
            "role": "user",
            "content": clean_message,
        }
    )

    mode_label = (
        "High reasoning"
        if high_reasoning
        else "Standard"
    )

    yield (
        "",
        display_history,
        f"**{mode_label}** · Building context…",
    )

    max_new_tokens = int(max_new_tokens)
    temperature = float(temperature)
    top_k = int(top_k)

    try:
        prompt_tokens, stats = build_conversation_tokens(
            message=clean_message,
            history=history,
            system_message=custom_instruction,
            max_new_tokens=max_new_tokens,
            use_few_shot=bool(use_few_shot),
            high_reasoning=bool(high_reasoning),
        )

        base_telemetry = (
            f"Mode **{mode_label}** · "
            f"Context **{stats['prompt_tokens']} / "
            f"{MODEL_CONTEXT_TOKENS}** · "
            f"History **{stats['history_turns']} turns** · "
            f"Few-shot **{stats['few_shot_messages']} messages**"
        )

        # ---------------------------------------------------------------------
        # High reasoning path
        # ---------------------------------------------------------------------
        if high_reasoning:
            yield (
                "",
                display_history,
                base_telemetry + " · **Reasoning…**",
            )

            with autocast_context():
                rollouts = generate_with_budget(
                    engine,
                    tokenizer,
                    prompt_tokens,
                    ON_HIGH,
                    num_samples=1,
                    temperature=temperature,
                    top_k=top_k,
                    seed=random.randint(0, 2**31 - 1),
                    max_answer_tokens=max_new_tokens,
                    slack=HIGH_REASONING_SLACK,
                    max_total_tokens=MODEL_CONTEXT_TOKENS,
                )

            if not rollouts:
                yield (
                    "",
                    [
                        *display_history,
                        {
                            "role": "assistant",
                            "content": (
                                "The reasoning engine generated no rollout."
                            ),
                        },
                    ],
                    base_telemetry,
                )
                return

            rollout = rollouts[0]

            parsed_answer = getattr(
                getattr(rollout, "parsed", None),
                "answer",
                "",
            )

            final_text = (
                parsed_answer.strip()
                if isinstance(parsed_answer, str)
                and parsed_answer.strip()
                else str(
                    getattr(rollout, "text", "")
                ).strip()
            )

            if not final_text:
                final_text = (
                    "The model completed its reasoning pass but "
                    "did not produce a final answer."
                )

            num_think_tokens = int(
                getattr(
                    rollout,
                    "num_think_tokens",
                    0,
                )
            )

            forced = bool(
                getattr(rollout, "forced", False)
            )

            reasoning_telemetry = (
                base_telemetry
                + f" · Think **{num_think_tokens} tokens**"
                + (
                    " · Budget forced **yes**"
                    if forced
                    else " · Budget forced **no**"
                )
            )

            yield (
                "",
                [
                    *display_history,
                    {
                        "role": "assistant",
                        "content": final_text,
                    },
                ],
                reasoning_telemetry,
            )
            return

        # ---------------------------------------------------------------------
        # Standard path
        # ---------------------------------------------------------------------
        generated_tokens: list[int] = []
        last_text = ""

        with autocast_context():
            stream = engine.generate(
                prompt_tokens,
                num_samples=1,
                max_tokens=max_new_tokens,
                temperature=temperature,
                top_k=top_k,
                seed=random.randint(0, 2**31 - 1),
            )

            for token_column, _token_masks in stream:
                raw_token = token_column[0]

                if isinstance(raw_token, torch.Tensor):
                    token = int(raw_token.item())
                else:
                    token = int(raw_token)

                if token in {
                    assistant_end_token,
                    bos_token,
                }:
                    break

                generated_tokens.append(token)

                current_text = tokenizer.decode(
                    generated_tokens
                )

                # Do not stream an incomplete UTF-8 replacement character.
                if current_text.endswith("�"):
                    continue

                if current_text != last_text:
                    last_text = current_text

                    streamed_history = [
                        *display_history,
                        {
                            "role": "assistant",
                            "content": current_text,
                        },
                    ]

                    yield (
                        "",
                        streamed_history,
                        base_telemetry,
                    )

        if not generated_tokens:
            yield (
                "",
                [
                    *display_history,
                    {
                        "role": "assistant",
                        "content": (
                            "The model generated an empty response."
                        ),
                    },
                ],
                base_telemetry,
            )

    except Exception as error:
        yield (
            "",
            [
                *display_history,
                {
                    "role": "assistant",
                    "content": (
                        "**Generation failed**\n\n"
                        f"`{type(error).__name__}: {error}`"
                    ),
                },
            ],
            f"**{mode_label}** · Generation error.",
        )


def clear_chat() -> tuple[
    list[dict[str, str]],
    str,
    str,
]:
    return [], "", "New conversation ready."


# =============================================================================
# Professional Gradio UI
# =============================================================================

CSS = """
:root {
    --app-radius: 18px;
}

.gradio-container {
    max-width: 1440px !important;
    margin: 0 auto !important;
}

#app-shell {
    padding-top: 8px;
}

#hero {
    border: 1px solid var(--border-color-primary);
    border-radius: var(--app-radius);
    padding: 20px 22px;
    margin-bottom: 14px;
    background: linear-gradient(
        135deg,
        color-mix(
            in srgb,
            var(--background-fill-primary) 94%,
            #6366f1 6%
        ),
        var(--background-fill-primary)
    );
}

#hero h1 {
    margin: 0 0 5px 0;
    font-size: 1.5rem;
    line-height: 1.2;
}

#hero p {
    margin: 0;
    opacity: 0.72;
}

.model-badge {
    display: inline-block;
    margin-top: 10px;
    padding: 5px 9px;
    border-radius: 999px;
    border: 1px solid var(--border-color-primary);
    font-size: 0.78rem;
    opacity: 0.85;
}

#chat-card,
#settings-card {
    border: 1px solid var(--border-color-primary);
    border-radius: var(--app-radius);
    background: var(--background-fill-primary);
    overflow: hidden;
}

#chatbot {
    border: none !important;
}

#composer-row {
    padding: 4px 2px 2px 2px;
}

#composer textarea {
    border-radius: 14px !important;
}

.prompt-chip button {
    border-radius: 999px !important;
    font-size: 0.82rem !important;
}

#telemetry {
    min-height: 26px;
    opacity: 0.67;
    font-size: 0.82rem;
}

.section-kicker {
    font-size: 0.76rem;
    font-weight: 700;
    letter-spacing: 0.08em;
    text-transform: uppercase;
    opacity: 0.58;
    margin-bottom: 4px;
}

@media (max-width: 900px) {
    .gradio-container {
        padding-left: 10px !important;
        padding-right: 10px !important;
    }
}
"""


with gr.Blocks(
    title="Supernova 1 Beta · Reasoning Chat",
    theme=gr.themes.Soft(),
    css=CSS,
) as demo:
    with gr.Column(elem_id="app-shell"):
        gr.HTML(
            f"""
            <div id="hero">
                <h1>Supernova 1 Beta</h1>
                <p>
                    Few-shot reasoning chat powered by the bundled NanoChat
                    checkpoint runtime.
                </p>
                <span class="model-badge">
                    {MODEL_ID} · checkpoint {checkpoint_step:,} ·
                    {device.type.upper()}
                </span>
            </div>
            """
        )

        with gr.Row(equal_height=False):
            # -----------------------------------------------------------------
            # Main chat
            # -----------------------------------------------------------------
            with gr.Column(
                scale=8,
                min_width=520,
                elem_id="chat-card",
            ):
                chatbot = gr.Chatbot(
                    value=[],
                    height=610,
                    elem_id="chatbot",
                    layout="bubble",
                    placeholder=(
                        "Start a conversation with Supernova 1 Beta."
                    ),
                    render_markdown=True,
                    buttons=["copy", "copy_all"],
                    feedback_options=["Like", "Dislike"],
                )

                telemetry = gr.Markdown(
                    "New conversation ready.",
                    elem_id="telemetry",
                )

                with gr.Row(elem_id="composer-row"):
                    message = gr.Textbox(
                        placeholder="Message Supernova 1 Beta…",
                        show_label=False,
                        lines=1,
                        max_lines=6,
                        autofocus=True,
                        container=False,
                        scale=8,
                        elem_id="composer",
                    )

                    send_button = gr.Button(
                        "Send",
                        variant="primary",
                        scale=1,
                        min_width=86,
                    )

                with gr.Row():
                    prompt_qml = gr.Button(
                        "Explain QML simply",
                        size="sm",
                        elem_classes="prompt-chip",
                    )

                    prompt_compare = gr.Button(
                        "Compare PPO and GRPO",
                        size="sm",
                        elem_classes="prompt-chip",
                    )

                    prompt_code = gr.Button(
                        "Debug a Python function",
                        size="sm",
                        elem_classes="prompt-chip",
                    )

                    prompt_reason = gr.Button(
                        "Give me a reasoning challenge",
                        size="sm",
                        elem_classes="prompt-chip",
                    )

            # -----------------------------------------------------------------
            # Settings / context
            # -----------------------------------------------------------------
            with gr.Column(
                scale=3,
                min_width=300,
                elem_id="settings-card",
            ):
                gr.Markdown(
                    "<div class='section-kicker'>Model controls</div>"
                )

                high_reasoning = gr.Checkbox(
                    value=False,
                    label="High reasoning",
                    info=(
                        "Uses ON_HIGH with budget-forced reasoning. "
                        "Standard mode streams faster."
                    ),
                )

                use_few_shot = gr.Checkbox(
                    value=True,
                    label="Use few-shot prompting",
                    info=(
                        "Injects compact examples into the actual prompt."
                    ),
                )

                max_new_tokens = gr.Slider(
                    minimum=32,
                    maximum=256,
                    value=128,
                    step=16,
                    label="Max answer tokens",
                )

                temperature = gr.Slider(
                    minimum=0.0,
                    maximum=2.0,
                    value=0.7,
                    step=0.1,
                    label="Temperature",
                )

                top_k = gr.Slider(
                    minimum=1,
                    maximum=200,
                    value=50,
                    step=1,
                    label="Top-k",
                )

                custom_instruction = gr.Textbox(
                    value="",
                    label="Additional instruction",
                    placeholder=(
                        "Example: Answer like a concise research assistant."
                    ),
                    lines=4,
                    max_lines=8,
                )

                with gr.Accordion(
                    "Reasoning behavior",
                    open=False,
                ):
                    gr.Markdown(
                        """
                        **High reasoning** adds the checkpoint's
                        `reasoning: high` condition and uses
                        budget-forced generation. The UI displays the
                        final answer while reporting thinking-token
                        usage in the telemetry line.
                        """
                    )

                with gr.Accordion(
                    "Built-in context",
                    open=False,
                ):
                    gr.Textbox(
                        value=BUILT_IN_CONTEXT.strip(),
                        interactive=False,
                        show_label=False,
                        lines=11,
                    )

                with gr.Accordion(
                    "Few-shot examples",
                    open=False,
                ):
                    gr.Markdown(
                        "\n\n".join(
                            (
                                f"**{item['role'].title()}**\n\n"
                                f"{item['content']}"
                            )
                            for item in FEW_SHOT_MESSAGES
                        )
                    )

                new_chat = gr.Button(
                    "New chat",
                    variant="secondary",
                )

        gr.Markdown(
            """
            <div style="
                text-align:center;
                opacity:.55;
                font-size:.78rem;
                padding:8px 0 2px 0;
            ">
                Experimental model demo · Responses may be incomplete
                or incorrect.
            </div>
            """
        )

    # -------------------------------------------------------------------------
    # Events
    # -------------------------------------------------------------------------

    generation_inputs = [
        message,
        chatbot,
        custom_instruction,
        max_new_tokens,
        temperature,
        top_k,
        use_few_shot,
        high_reasoning,
    ]

    generation_outputs = [
        message,
        chatbot,
        telemetry,
    ]

    message.submit(
        fn=respond,
        inputs=generation_inputs,
        outputs=generation_outputs,
        show_progress="hidden",
        concurrency_limit=1,
    )

    send_button.click(
        fn=respond,
        inputs=generation_inputs,
        outputs=generation_outputs,
        show_progress="hidden",
        concurrency_limit=1,
    )

    new_chat.click(
        fn=clear_chat,
        inputs=None,
        outputs=[
            chatbot,
            message,
            telemetry,
        ],
        show_progress="hidden",
    )

    prompt_qml.click(
        fn=lambda: (
            "Explain quantum machine learning to a software engineer "
            "in under 120 words."
        ),
        outputs=message,
        show_progress="hidden",
    )

    prompt_compare.click(
        fn=lambda: (
            "Compare PPO and GRPO: objective, data flow, advantages, "
            "and when you would use each."
        ),
        outputs=message,
        show_progress="hidden",
    )

    prompt_code.click(
        fn=lambda: (
            "Give me a short Python function containing one subtle bug, "
            "then ask me to find and fix it."
        ),
        outputs=message,
        show_progress="hidden",
    )

    prompt_reason.click(
        fn=lambda: (
            "Give me a difficult but self-contained reasoning problem. "
            "Do not reveal the answer until I respond."
        ),
        outputs=message,
        show_progress="hidden",
    )


if __name__ == "__main__":
    demo.queue(
        default_concurrency_limit=1,
    ).launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        ssr_mode=False,
    )