Update app.py

app.py

@@ -1,4 +1,4 @@
-"""Gradio interface for nanochat model (using Guilherme34 repo)."""
+"""Gradio interface for nanochat model."""
 
 from __future__ import annotations
 
@@ -12,37 +12,18 @@ from huggingface_hub import snapshot_download
 
 from model import NanochatModel
 
-# --- Configuration: point to YOUR repo and subfolders ---
-# Defaults can still be overridden by env vars if you want.
-MODEL_REPO = os.environ.get("MODEL_REPO", "Guilherme34/nanochat-retrained-pytorch")
+MODEL_REPO = os.environ.get("MODEL_REPO", "Guilherme34/nanochat-retrained-pytorch-duplicated")
 MODEL_DIR = os.environ.get("MODEL_DIR", "./model_cache")
-
-# The model files live in "sft/" and the tokenizer in "tokenizer/" inside the repo.
-# We'll mirror that structure locally under MODEL_DIR.
-MODEL_SUBDIR = "sft"
-TOKENIZER_SUBDIR = "tokenizer"
-
 _model: NanochatModel | None = None
 
 
 def download_model() -> None:
-    """Download the model + tokenizer folders from Hugging Face if needed."""
+    """Download the model from Hugging Face if needed."""
     model_path = Path(MODEL_DIR)
-    sft_path = model_path / MODEL_SUBDIR
-    tok_path = model_path / TOKENIZER_SUBDIR
-
-    # If either folder is missing/empty, fetch both to keep them in sync.
-    need_download = (
-        not sft_path.exists() or not any(sft_path.iterdir()) or
-        not tok_path.exists() or not any(tok_path.iterdir())
-    )
-
-    if need_download:
-        # Only pull what we need to keep downloads light.
+    if not model_path.exists() or not any(model_path.iterdir()):
         snapshot_download(
             repo_id=MODEL_REPO,
             local_dir=MODEL_DIR,
-            allow_patterns=[f"{MODEL_SUBDIR}/**", f"{TOKENIZER_SUBDIR}/**"],
         )
 
 
@@ -51,13 +32,6 @@ def load_model() -> None:
     global _model
     if _model is None:
         download_model()
-        # If your NanochatModel supports passing a tokenizer_dir, uncomment and use it:
-        # _model = NanochatModel(
-        #     model_dir=str(Path(MODEL_DIR) / MODEL_SUBDIR),
-        #     tokenizer_dir=str(Path(MODEL_DIR) / TOKENIZER_SUBDIR),
-        #     device="cpu",
-        # )
-        # Otherwise, keep the original and let your class discover subfolders:
         _model = NanochatModel(model_dir=MODEL_DIR, device="cpu")
 
 
@@ -69,7 +43,7 @@ def respond(
     history: list[dict[str, str]],
     temperature: float,
     top_k: int,
-    system_prompt: str,
+    system_prompt: str,  # NEW
 ) -> Generator[str, Any, None]:
     """Generate a response using the nanochat model.
 
@@ -85,11 +59,10 @@ def respond(
     """
     conversation: list[dict[str, str]] = []
 
-    # Only include a system message if one was actually provided
-    if system_prompt and system_prompt.strip():
-        conversation.append({"role": "system", "content": system_prompt.strip()})
+    # If a system message is provided, put it at the start of the conversation.
+    conversation.append({"role": "system", "content": system_prompt.strip()})
 
-    # Replay prior turns (already in gradio "messages" format)
+    # Replay prior turns
     for msg in history:
         conversation.append(msg)
 
@@ -112,19 +85,24 @@ chatbot = gr.ChatInterface(
     type="messages",
     additional_inputs=[
         gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(minimum=1, maximum=200, value=50, step=1, label="Top-k sampling"),
-        gr.Textbox(
+        gr.Slider(
+            minimum=1,
+            maximum=200,
+            value=50,
+            step=1,
+            label="Top-k sampling",
+        ),
+        gr.Textbox(  # NEW
             label="System message (optional)",
             placeholder="e.g., You are a concise assistant that answers in markdown.",
             lines=3,
-            value="",  # ensure we always pass a string
         ),
     ],
 )
 
-with gr.Blocks(title="nanochat (Guilherme34)") as demo:
-    gr.Markdown("# nanochat (retrained, PyTorch)")
-    gr.Markdown("Using model weights from `sft/` and tokenizer from `tokenizer/` in your repo.")
+with gr.Blocks(title="nanochat") as demo:
+    gr.Markdown("# nanochat")
+    gr.Markdown("Chat with an AI trained in 4 hours for $100")
     gr.Markdown(
         "**Note:** This model is a research experiment. "
         "Obviously do not rely on the outputs!",
@@ -133,6 +111,4 @@ with gr.Blocks(title="nanochat (Guilherme34)") as demo:
 
 
 if __name__ == "__main__":
-    # To expose publicly or change port, use:
-    # demo.launch(server_name='0.0.0.0', server_port=7860, share=True)
-    demo.launch()
+    demo.launch()