Update app.py

app.py

@@ -1,69 +1,116 @@
+"""
+Gradio chat interface for nanochat models on Hugging Face Spaces.
+Downloads the model from HF Hub and serves a streaming chat UI.
+"""
+
+import os
+import json
+import random
+import torch
 import gradio as gr
-from huggingface_hub import InferenceClient
+from huggingface_hub import snapshot_download
+
+# ---------------------------------------------------------------------------
+# Configuration
+# ---------------------------------------------------------------------------
+MODEL_REPO = os.environ.get("MODEL_REPO", "Slothwolf/MiniChat-0.9B")
+MODEL_STEP = int(os.environ.get("MODEL_STEP", "486"))
+DEFAULT_TEMPERATURE = float(os.environ.get("DEFAULT_TEMPERATURE", "0.8"))
+DEFAULT_MAX_TOKENS = int(os.environ.get("DEFAULT_MAX_TOKENS", "1024"))
+DEFAULT_TOP_K = int(os.environ.get("DEFAULT_TOP_K", "50"))
+
+# ---------------------------------------------------------------------------
+# Download and load model
+# ---------------------------------------------------------------------------
+print(f"Downloading model from {MODEL_REPO}...")
+model_dir = snapshot_download(MODEL_REPO)
+
+# Point nanochat at the downloaded directory so get_tokenizer() finds tokenizer/
+os.environ["NANOCHAT_BASE_DIR"] = model_dir
 
+from nanochat.checkpoint_manager import build_model
+from nanochat.engine import Engine
 
-def respond(
-    message,
-    history: list[dict[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-    hf_token: gr.OAuthToken,
-):
-    """
-    For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-    """
-    client = InferenceClient(token=hf_token.token, model="openai/gpt-oss-20b")
+print(f"Loading model (step {MODEL_STEP}) on CPU...")
+device = torch.device("cpu")
+model, tokenizer, meta = build_model(model_dir, MODEL_STEP, device, "eval")
+engine = Engine(model, tokenizer)
+print("Model loaded and ready!")
 
-    messages = [{"role": "system", "content": system_message}]
+# Grab special token ids once
+bos_id = tokenizer.get_bos_token_id()
+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|>")
 
-    messages.extend(history)
+# ---------------------------------------------------------------------------
+# Chat function (streaming)
+# ---------------------------------------------------------------------------
+def chat(message, history, temperature, max_tokens, top_k):
+    """Generate a streaming response for the chat interface."""
+    # Build the full token sequence from conversation history
+    # history is a list of [user_msg, assistant_msg] tuples
+    tokens = [bos_id]
+    for user_msg, assistant_msg in history:
+        tokens.append(user_start)
+        tokens.extend(tokenizer.encode(user_msg))
+        tokens.append(user_end)
+        if assistant_msg:
+            tokens.append(assistant_start)
+            tokens.extend(tokenizer.encode(assistant_msg))
+            tokens.append(assistant_end)
 
-    messages.append({"role": "user", "content": message})
+    # Add the current user message and prime the assistant
+    tokens.append(user_start)
+    tokens.extend(tokenizer.encode(message))
+    tokens.append(user_end)
+    tokens.append(assistant_start)
 
+    # Generate tokens with streaming
     response = ""
+    accumulated_tokens = []
+    last_clean_text = ""
 
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
+    for token_column, token_masks in engine.generate(
+        tokens,
+        num_samples=1,
+        max_tokens=int(max_tokens),
+        temperature=float(temperature),
+        top_k=int(top_k) if int(top_k) > 0 else None,
+        seed=random.randint(0, 2**31 - 1),
     ):
-        choices = message.choices
-        token = ""
-        if len(choices) and choices[0].delta.content:
-            token = choices[0].delta.content
+        token = token_column[0]
+        if token == assistant_end or token == bos_id:
+            break
 
-        response += token
-        yield response
+        # Accumulate tokens and decode, handling multi-byte UTF-8
+        accumulated_tokens.append(token)
+        current_text = tokenizer.decode(accumulated_tokens)
+        if not current_text.endswith("\ufffd"):
+            new_text = current_text[len(last_clean_text):]
+            if new_text:
+                response += new_text
+                last_clean_text = current_text
+                yield response
 
-
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-chatbot = gr.ChatInterface(
-    respond,
+# ---------------------------------------------------------------------------
+# Gradio UI
+# ---------------------------------------------------------------------------
+demo = gr.ChatInterface(
+    chat,
     additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
+        gr.Slider(0.0, 2.0, value=DEFAULT_TEMPERATURE, step=0.1, label="Temperature"),
+        gr.Slider(1, 2048, value=DEFAULT_MAX_TOKENS, step=1, label="Max Tokens"),
+        gr.Slider(0, 200, value=DEFAULT_TOP_K, step=1, label="Top-K (0 = disabled)"),
     ],
+    title="MiniChat",
+    description=(
+        "Chat with a **0.9B parameter** GPT-like model trained with "
+        "[nanochat](https://github.com/karpathy/nanochat) by Andrej Karpathy. "
+        "Running on CPU — generation will be slow."
+    ),
 )
 
-with gr.Blocks() as demo:
-    with gr.Sidebar():
-        gr.LoginButton()
-    chatbot.render()
-
-
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()