Update app.py

app.py

@@ -1,9 +1,9 @@
 from __future__ import annotations
 
 import os
+import re
 import sys
 from pathlib import Path
-from typing import Any
 
 import gradio as gr
 import torch
@@ -13,7 +13,9 @@ from huggingface_hub import snapshot_download
 MODEL_ID = "MarkChenX/lfm2-quantum-128m-sft-v2"
 HF_TOKEN = os.getenv("HF_TOKEN")
 
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+DEVICE = torch.device(
+    "cuda" if torch.cuda.is_available() else "cpu"
+)
 
 
 print(f"Downloading custom checkpoint: {MODEL_ID}")
@@ -22,259 +24,85 @@ 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",
+            "*.bin",
+        ],
     )
 )
 
-# The repository contains the custom nanochat package.
-sys.path.insert(0, str(model_dir))
-
 print(f"Model repository: {model_dir}")
 
+sys.path.insert(0, str(model_dir))
 
 try:
     from nanochat.checkpoint_manager import build_model
 except ImportError as error:
     raise RuntimeError(
-        "Could not import nanochat.checkpoint_manager from the model "
-        f"repository at {model_dir}. Ensure the repository contains "
-        "nanochat/__init__.py and nanochat/checkpoint_manager.py."
+        "Could not import nanochat.checkpoint_manager. "
+        "Confirm that the model repository contains the nanochat directory."
     ) from error
 
 
-def find_checkpoint_file() -> Path:
-    candidates = [
-        model_dir / "model.pt",
-        model_dir / "model.safetensors",
-        model_dir / "checkpoint.pt",
-        model_dir / "checkpoint_849.pt",
-    ]
-
-    for path in candidates:
-        if path.exists():
-            return path
+def find_latest_checkpoint(
+    repository_dir: Path,
+) -> tuple[Path, int]:
+    checkpoint_candidates: list[tuple[int, Path]] = []
 
-    discovered = list(model_dir.glob("*.pt"))
-    discovered += list(model_dir.glob("*.pth"))
-    discovered += list(model_dir.glob("*.safetensors"))
-
-    if not discovered:
-        raise FileNotFoundError(
-            f"No checkpoint file found in {model_dir}"
+    for checkpoint in repository_dir.glob("model_*.pt"):
+        match = re.fullmatch(
+            r"model_(\d+)\.pt",
+            checkpoint.name,
         )
 
-    return discovered[0]
-
-
-checkpoint_path = find_checkpoint_file()
-
-print(f"Checkpoint file: {checkpoint_path}")
-print("Building custom nanochat model...")
-
-
-# Adjust these arguments only if your actual build_model signature differs.
-model, tokenizer = build_model(
-    checkpoint_path=str(checkpoint_path),
-    device=DEVICE,
-)
-
-model.eval()
-
-print(f"Model loaded on {DEVICE}")
-
-
-def normalize_history(
-    history: list[Any] | None,
-) -> list[dict[str, str]]:
-    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,
-                    }
-                )
-
-        elif isinstance(item, (list, tuple)) and len(item) == 2:
-            user_text, assistant_text = item
-
-            if isinstance(user_text, str) and user_text.strip():
-                messages.append(
-                    {
-                        "role": "user",
-                        "content": user_text,
-                    }
-                )
-
-            if (
-                isinstance(assistant_text, str)
-                and assistant_text.strip()
-            ):
-                messages.append(
-                    {
-                        "role": "assistant",
-                        "content": assistant_text,
-                    }
-                )
-
-    return messages
-
-
-def build_prompt(
-    message: str,
-    history: list[Any],
-    system_message: str,
-) -> str:
-    messages: list[dict[str, str]] = []
+        if match:
+            step = int(match.group(1))
+            checkpoint_candidates.append((step, checkpoint))
 
-    if system_message.strip():
-        messages.append(
-            {
-                "role": "system",
-                "content": system_message.strip(),
-            }
+    if not checkpoint_candidates:
+        raise FileNotFoundError(
+            f"No model_XXXXXX.pt checkpoint found in {repository_dir}"
         )
 
-    messages.extend(normalize_history(history))
-
-    messages.append(
-        {
-            "role": "user",
-            "content": message.strip(),
-        }
+    checkpoint_candidates.sort(
+        key=lambda item: item[0],
+        reverse=True,
     )
 
-    parts: list[str] = []
-
-    for item in messages:
-        role = item["role"]
-        content = item["content"]
-        parts.append(f"<|im_start|>{role}\n{content}<|im_end|>\n")
-
-    parts.append("<|im_start|>assistant\n")
-
-    return "".join(parts)
-
-
-def encode_prompt(prompt: str) -> torch.Tensor:
-    encoded = tokenizer.encode(prompt)
-
-    if isinstance(encoded, torch.Tensor):
-        input_ids = encoded
-    else:
-        input_ids = torch.tensor(encoded, dtype=torch.long)
+    latest_step, latest_path = checkpoint_candidates[0]
+    return latest_path, latest_step
 
-    if input_ids.ndim == 1:
-        input_ids = input_ids.unsqueeze(0)
 
-    return input_ids.to(DEVICE)
+checkpoint_path, checkpoint_step = find_latest_checkpoint(model_dir)
 
+expected_meta_path = model_dir / f"meta_{checkpoint_step:06d}.json"
 
-def decode_tokens(token_ids: torch.Tensor) -> str:
-    ids = token_ids.detach().cpu().tolist()
-
-    if ids and isinstance(ids[0], list):
-        ids = ids[0]
-
-    return tokenizer.decode(ids)
-
-
-def respond(
-    message: str,
-    history: list[Any],
-    system_message: str,
-    max_tokens: int,
-    temperature: float,
-    top_p: float,
-) -> str:
-    if not message or not message.strip():
-        return "Please enter a message."
-
-    try:
-        prompt = build_prompt(
-            message=message,
-            history=history,
-            system_message=system_message,
-        )
-
-        input_ids = encode_prompt(prompt)
-        input_length = input_ids.shape[-1]
-
-        generation_kwargs = {
-            "input_ids": input_ids,
-            "max_new_tokens": int(max_tokens),
-            "temperature": float(temperature),
-            "top_p": float(top_p),
-        }
-
-        with torch.inference_mode():
-            output = model.generate(**generation_kwargs)
-
-        generated_ids = output[:, input_length:]
-        response = decode_tokens(generated_ids).strip()
-
-        return response or "The model generated an empty response."
-
-    except Exception as error:
-        return (
-            "Generation failed: "
-            f"{type(error).__name__}: {error}"
-        )
+if not expected_meta_path.exists():
+    raise FileNotFoundError(
+        "The checkpoint metadata file is missing. Expected: "
+        f"{expected_meta_path}"
+    )
 
 
-with gr.Blocks(title="LFM2 Quantum Chat") as demo:
-    gr.Markdown(
-        """
-        # LFM2 Quantum 128M
+print(f"Checkpoint file: {checkpoint_path}")
+print(f"Checkpoint step: {checkpoint_step}")
+print(f"Metadata file: {expected_meta_path}")
+print("Building custom nanochat model...")
 
-        Custom Nanochat LFM2-style hybrid quantum model.
-        """
-    )
 
-    gr.ChatInterface(
-        fn=respond,
-        additional_inputs=[
-            gr.Textbox(
-                value="You are a helpful and friendly assistant.",
-                label="System message",
-                lines=3,
-            ),
-            gr.Slider(
-                minimum=1,
-                maximum=1024,
-                value=256,
-                step=1,
-                label="Max new tokens",
-            ),
-            gr.Slider(
-                minimum=0.0,
-                maximum=2.0,
-                value=0.3,
-                step=0.1,
-                label="Temperature",
-            ),
-            gr.Slider(
-                minimum=0.1,
-                maximum=1.0,
-                value=0.95,
-                step=0.05,
-                label="Top-p",
-            ),
-        ],
-    )
+model, tokenizer, metadata = build_model(
+    str(model_dir),
+    checkpoint_step,
+    DEVICE,
+    "eval",
+)
 
+model.eval()
 
-if __name__ == "__main__":
-    demo.queue().launch(
-        share=False,
-        ssr_mode=False,
-    )
+print(f"Model loaded successfully on {DEVICE}")
+print(
+    "Model config:",
+    metadata.get("model_config", "not present"),
+)