Update app.py

app.py

@@ -1,9 +1,12 @@
 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
@@ -13,18 +16,23 @@ from huggingface_hub import snapshot_download
 MODEL_ID = "MarkChenX/lfm2-quantum-128m-sft-v2"
 HF_TOKEN = os.getenv("HF_TOKEN")
 
-DEVICE = torch.device(
+MAX_CONTEXT_TOKENS = 1024
+
+device = torch.device(
     "cuda" if torch.cuda.is_available() else "cpu"
 )
 
 
-print(f"Downloading custom checkpoint: {MODEL_ID}")
+# ---------------------------------------------------------------------
+# Download repository
+# ---------------------------------------------------------------------
+
+print(f"Downloading model repository: {MODEL_ID}")
 
 model_dir = Path(
     snapshot_download(
         repo_id=MODEL_ID,
         token=HF_TOKEN,
-        # Do not download optimizer files into the Space.
         ignore_patterns=[
             "optim_*.pt",
             "optimizer*.pt",
@@ -35,74 +43,388 @@ model_dir = Path(
 
 print(f"Model repository: {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
 except ImportError as error:
     raise RuntimeError(
-        "Could not import nanochat.checkpoint_manager. "
-        "Confirm that the model repository contains the nanochat directory."
+        "Could not import the bundled NanoChat code. The model repository "
+        "must contain nanochat/checkpoint_manager.py and nanochat/engine.py."
     ) from error
 
 
+# ---------------------------------------------------------------------
+# Locate latest model checkpoint
+# ---------------------------------------------------------------------
+
 def find_latest_checkpoint(
-    repository_dir: Path,
+    directory: Path,
 ) -> tuple[Path, int]:
-    checkpoint_candidates: list[tuple[int, Path]] = []
+    candidates: list[tuple[int, Path]] = []
 
-    for checkpoint in repository_dir.glob("model_*.pt"):
+    for path in directory.rglob("model_*.pt"):
         match = re.fullmatch(
             r"model_(\d+)\.pt",
-            checkpoint.name,
+            path.name,
         )
 
         if match:
-            step = int(match.group(1))
-            checkpoint_candidates.append((step, checkpoint))
+            candidates.append(
+                (int(match.group(1)), path)
+            )
 
-    if not checkpoint_candidates:
+    if not candidates:
         raise FileNotFoundError(
-            f"No model_XXXXXX.pt checkpoint found in {repository_dir}"
+            f"No model_XXXXXX.pt checkpoint found under {directory}"
         )
 
-    checkpoint_candidates.sort(
+    candidates.sort(
         key=lambda item: item[0],
         reverse=True,
     )
 
-    latest_step, latest_path = checkpoint_candidates[0]
-    return latest_path, latest_step
+    step, checkpoint_path = candidates[0]
+    return checkpoint_path, step
 
 
 checkpoint_path, checkpoint_step = find_latest_checkpoint(model_dir)
+checkpoint_dir = checkpoint_path.parent
 
-expected_meta_path = model_dir / f"meta_{checkpoint_step:06d}.json"
+metadata_path = checkpoint_dir / f"meta_{checkpoint_step:06d}.json"
 
-if not expected_meta_path.exists():
+if not metadata_path.exists():
     raise FileNotFoundError(
-        "The checkpoint metadata file is missing. Expected: "
-        f"{expected_meta_path}"
+        "Checkpoint metadata is missing. Expected: "
+        f"{metadata_path}"
     )
 
 
+print(f"Checkpoint directory: {checkpoint_dir}")
 print(f"Checkpoint file: {checkpoint_path}")
 print(f"Checkpoint step: {checkpoint_step}")
-print(f"Metadata file: {expected_meta_path}")
-print("Building custom nanochat model...")
+print(f"Metadata file: {metadata_path}")
+print(f"Loading on device: {device}")
 
 
+# ---------------------------------------------------------------------
+# Load custom NanoChat model
+# ---------------------------------------------------------------------
+
 model, tokenizer, metadata = build_model(
-    str(model_dir),
+    str(checkpoint_dir),
     checkpoint_step,
-    DEVICE,
+    device,
     "eval",
 )
 
 model.eval()
 
-print(f"Model loaded successfully on {DEVICE}")
-print(
-    "Model config:",
-    metadata.get("model_config", "not present"),
-)
+engine = Engine(model, tokenizer)
+
+assistant_end_token = tokenizer.encode_special(
+    "<|assistant_end|>"
+)
+
+bos_token = tokenizer.get_bos_token_id()
+
+
+if device.type == "cuda":
+    autocast_context = lambda: torch.amp.autocast(
+        device_type="cuda",
+        dtype=torch.bfloat16,
+    )
+else:
+    autocast_context = nullcontext
+
+
+print("Model and NanoChat engine loaded successfully.")
+
+
+# ---------------------------------------------------------------------
+# Gradio history conversion
+# ---------------------------------------------------------------------
+
+def normalize_history(
+    history: list[Any] | None,
+) -> list[dict[str, str]]:
+    messages: list[dict[str, str]] = []
+
+    for item in history or []:
+        # Current Gradio message format.
+        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(),
+                    }
+                )
+
+        # Compatibility with older tuple-style Gradio history.
+        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
+
+
+# ---------------------------------------------------------------------
+# Build NanoChat-native conversation tokens
+# ---------------------------------------------------------------------
+
+def build_conversation_tokens(
+    message: str,
+    history: list[Any] | None,
+    system_message: str,
+) -> list[int]:
+    messages = normalize_history(history)
+
+    current_message = message.strip()
+
+    # NanoChat merges system instructions into the first user message.
+    if system_message.strip():
+        current_message = (
+            system_message.strip()
+            + "\n\n"
+            + current_message
+        )
+
+    messages.append(
+        {
+            "role": "user",
+            "content": current_message,
+        }
+    )
+
+    tokens: list[int] = [bos_token]
+
+    user_start = tokenizer.encode_special("<|user_start|>")
+    user_end = tokenizer.encode_special("<|user_end|>")
+
+    assistant_start = tokenizer.encode_special(
+        "<|assistant_start|>"
+    )
+    assistant_end = tokenizer.encode_special(
+        "<|assistant_end|>"
+    )
+
+    for chat_message in messages:
+        content_tokens = tokenizer.encode(
+            chat_message["content"]
+        )
+
+        if chat_message["role"] == "user":
+            tokens.append(user_start)
+            tokens.extend(content_tokens)
+            tokens.append(user_end)
+
+        elif chat_message["role"] == "assistant":
+            tokens.append(assistant_start)
+            tokens.extend(content_tokens)
+            tokens.append(assistant_end)
+
+    # Prime the model to produce the next assistant message.
+    tokens.append(assistant_start)
+
+    # Leave room for generation within the model's 1024-token context.
+    return tokens
+
+
+# ---------------------------------------------------------------------
+# Gradio generation function
+# ---------------------------------------------------------------------
+
+def respond(
+    message: str,
+    history: list[Any],
+    system_message: str,
+    max_new_tokens: int,
+    temperature: float,
+    top_k: int,
+) -> Generator[str, None, None]:
+    if not message or not message.strip():
+        yield "Please enter a message."
+        return
+
+    max_new_tokens = int(max_new_tokens)
+    temperature = float(temperature)
+    top_k = int(top_k)
+
+    try:
+        prompt_tokens = build_conversation_tokens(
+            message=message,
+            history=history,
+            system_message=system_message,
+        )
+
+        # Preserve enough context space for newly generated tokens.
+        maximum_prompt_length = (
+            MAX_CONTEXT_TOKENS - max_new_tokens
+        )
+
+        if maximum_prompt_length <= 0:
+            yield (
+                "Max new tokens must be smaller than the "
+                f"{MAX_CONTEXT_TOKENS}-token context window."
+            )
+            return
+
+        if len(prompt_tokens) > maximum_prompt_length:
+            # Keep BOS and the newest conversation context.
+            prompt_tokens = [
+                bos_token,
+                *prompt_tokens[-(maximum_prompt_length - 1):],
+            ]
+
+        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:
+                token = token_column[0]
+
+                if token in {
+                    assistant_end_token,
+                    bos_token,
+                }:
+                    break
+
+                generated_tokens.append(token)
+
+                current_text = tokenizer.decode(
+                    generated_tokens
+                )
+
+                # Avoid yielding incomplete UTF-8 replacement characters.
+                if current_text.endswith("�"):
+                    continue
+
+                if current_text != last_text:
+                    last_text = current_text
+                    yield current_text
+
+        if not generated_tokens:
+            yield "The model generated an empty response."
+
+    except Exception as error:
+        yield (
+            "Generation failed: "
+            f"{type(error).__name__}: {error}"
+        )
+
+
+# ---------------------------------------------------------------------
+# Gradio interface
+# ---------------------------------------------------------------------
+
+with gr.Blocks(
+    title="LFM2 Quantum 128M",
+) as demo:
+    gr.Markdown(
+        """
+        # LFM2 Quantum 128M
+
+        Custom LFM2-style hybrid quantum language model running through
+        the bundled NanoChat inference engine.
+        """
+    )
+
+    chatbot = gr.ChatInterface(
+        fn=respond,
+        type="messages",
+        chatbot=gr.Chatbot(
+            type="messages",
+            height=520,
+            placeholder=(
+                "Ask the LFM2 Quantum model a question."
+            ),
+        ),
+        additional_inputs=[
+            gr.Textbox(
+                value=(
+                    "You are a helpful, concise, and friendly assistant."
+                ),
+                label="System instruction",
+                lines=3,
+            ),
+            gr.Slider(
+                minimum=16,
+                maximum=512,
+                value=256,
+                step=16,
+                label="Max new tokens",
+            ),
+            gr.Slider(
+                minimum=0.0,
+                maximum=2.0,
+                value=0.7,
+                step=0.1,
+                label="Temperature",
+            ),
+            gr.Slider(
+                minimum=1,
+                maximum=200,
+                value=50,
+                step=1,
+                label="Top-k",
+            ),
+        ],
+        examples=[
+            ["Hello! Introduce yourself."],
+            ["Explain quantum machine learning simply."],
+            ["What is reinforcement learning?"],
+        ],
+    )
+
+
+if __name__ == "__main__":
+    demo.queue(
+        default_concurrency_limit=1,
+    ).launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False,
+        ssr_mode=False,
+    )