app.py

#!/usr/bin/env python3
"""
Hugging Face Spaces entry point for NanoChat.
This file is automatically detected and run by HF Spaces.
"""

import os
import sys

# Set environment variables for HF Spaces
os.environ.setdefault("NANOCHAT_BASE_DIR", "/data")

# Download model from HF if not present
def download_model():
    """Download model weights from Hugging Face."""
    checkpoint_dir = "/data/chatsft_checkpoints"

    if os.path.exists(checkpoint_dir) and os.listdir(checkpoint_dir):
        print(f"Model checkpoints found, skipping download")
        return

    print("Downloading model from BrianGuo/nanochat-d20-chat...")
    from huggingface_hub import snapshot_download

    snapshot_download(
        repo_id="BrianGuo/nanochat-d20-chat",
        local_dir="/data/chatsft_checkpoints"
    )
    print("Model downloaded successfully!")

if __name__ == "__main__":
    # Download model before starting
    download_model()

    # Override sys.argv to pass default arguments for HF Spaces
    sys.argv = [
        "app.py",
        "--port", "7860",  # HF Spaces default port
        "--host", "0.0.0.0",
        "--source", "sft",
        "--model-tag", os.environ.get("MODEL_TAG", "d20"),
        "--step", os.environ.get("MODEL_STEP", "650"),
    ]

    # Import and run the web server
    from scripts.chat_web import app
    import uvicorn

    print("Starting NanoChat on Hugging Face Spaces...")
    print(f"Model: {os.environ.get('MODEL_TAG', 'd20')} - Step: {os.environ.get('MODEL_STEP', '650')}")

    uvicorn.run(app, host="0.0.0.0", port=7860)