Updates to flash attn install

Dockerfile

@@ -91,6 +91,11 @@ RUN curl https://sh.rustup.rs -sSf | bash -s -- -y
 ENV PATH="$HOME/.cargo/bin:${PATH}"
 RUN chmod -R 777 $HOME/.cargo/
 
+# Install ninja, which might help speed up flash attn build time
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        ninja-build \
+        && rm -rf /var/lib/apt/lists/*
+
 # Python packages
 RUN --mount=target=requirements.txt,source=requirements.txt \
     pip install --no-cache-dir --upgrade -r requirements.txt

start_server.sh

@@ -3,7 +3,9 @@ JUPYTER_TOKEN="${JUPYTER_TOKEN:=huggingface}"
 
 echo "Starting Jupyter Lab with token $JUPYTER_TOKEN"
 make install
+cd server
 make install-flash-attention
+cd ..
 
 jupyter-lab \
     --ip 0.0.0.0 \