make fastapi and uvicorn optional dependencies via --extra web

README.md

@@ -43,6 +43,12 @@ For development (adds pytest, matplotlib, ipykernel, transformers, etc.):
 uv sync --extra gpu --group dev
 ```
 
+If you plan on running scripts.chat_web:
+
+```bash
+uv sync --extra web
+```
+
 ### Reproduce and talk to GPT-2
 
 The most fun you can have is to train your own GPT-2 and talk to it. The entire pipeline to do so is contained in the single file [runs/speedrun.sh](runs/speedrun.sh), which is designed to be run on an 8XH100 GPU node. Boot up a new 8XH100 GPU box from your favorite provider (e.g. I use and like [Lambda](https://lambda.ai/service/gpu-cloud)), and kick off the training script:
@@ -51,9 +57,10 @@ The most fun you can have is to train your own GPT-2 and talk to it. The entire
 bash runs/speedrun.sh
 ```
 
-You may wish to do so in a screen session as this will take ~3 hours to run. Once it's done, you can talk to it via the ChatGPT-like web UI. Make sure again that your local uv virtual environment is active (run `source .venv/bin/activate`), and serve it:
+You may wish to do so in a screen session as this will take ~3 hours to run. Once it's done, you can talk to it via the ChatGPT-like web UI. Make sure again that your local uv virtual environment is active (run `source .venv/bin/activate`) and has the `web` extra installed, and then serve it:
 
 ```bash
+uv sync --extra web
 python -m scripts.chat_web
 ```
 

pyproject.toml

@@ -6,14 +6,12 @@ readme = "README.md"
 requires-python = ">=3.10"
 dependencies = [
     "datasets>=4.0.0",
-    "fastapi>=0.117.1",
     "kernels>=0.11.7",
     "psutil>=7.1.0",
     "rustbpe>=0.1.0",
     "tiktoken>=0.11.0",
     "tokenizers>=0.22.0",
     "torch==2.9.1",
-    "uvicorn>=0.36.0",
     "wandb>=0.21.3",
 ]
 
@@ -59,6 +57,10 @@ cpu = [
 gpu = [
     "torch==2.9.1",
 ]
+web = [
+    "fastapi>=0.117.1",
+    "uvicorn>=0.36.0",
+]
 
 [tool.uv]
 default-groups = []

runs/runcpu.sh

@@ -62,4 +62,5 @@ python -m scripts.chat_sft \
 # python -m scripts.chat_cli -p "What is the capital of France?"
 
 # Chat with the model over a pretty WebUI ChatGPT style
+# uv sync --extra web
 # python -m scripts.chat_web

runs/speedrun.sh

@@ -89,6 +89,7 @@ torchrun --standalone --nproc_per_node=8 -m scripts.chat_eval -- -i sft
 # python -m scripts.chat_cli -p "Why is the sky blue?"
 
 # even better, chat with your model over a pretty WebUI ChatGPT style
+# uv sync --extra web
 # python -m scripts.chat_web
 
 # -----------------------------------------------------------------------------

scripts/chat_web.py

@@ -38,9 +38,12 @@ import asyncio
 import logging
 import random
 from contextlib import asynccontextmanager
-from fastapi import FastAPI, HTTPException
-from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import StreamingResponse, HTMLResponse, FileResponse
+try:
+    from fastapi import FastAPI, HTTPException
+    from fastapi.middleware.cors import CORSMiddleware
+    from fastapi.responses import StreamingResponse, HTMLResponse, FileResponse
+except ImportError as exc:
+    raise SystemExit("Missing web dependencies, install with: uv sync --extra web") from exc
 from pydantic import BaseModel
 from typing import List, Optional, AsyncGenerator
 from dataclasses import dataclass