diff --git a/setup.sh b/setup.sh
index a175510..02855e6 100755
--- a/setup.sh
+++ b/setup.sh
@@ -1,41 +1,79 @@
 #!/bin/sh
 
-set -e
+set -eu
+set -o pipefail
 
-check_command() {
-    if [ $? -ne 0 ]; then
-        echo "Error: $1 failed"
-        exit 1
-    fi
-}
+# constants
+WORKSPACE_DIR="/workspace"
+REPO_URL="https://git.hye.su/mira/unsloth-train-scripts.git"
+INSTALL_SCRIPT_URL="https://raw.githubusercontent.com/unslothai/unsloth/main/unsloth/_auto_install.py"
+DATASET_URL="https://drive.google.com/file/d/1mqhq69dsSOK7ep7trTjRd3FMagTFTrzF/view?usp=sharing"
 
-cd /workspace || {
-    echo "Error: Cannot change to /workspace directory"
+error_exit() {
+    echo "Error: $1" >&2
     exit 1
 }
 
-echo "Cloning unsloth-train-scripts repository..."
-git clone https://git.hye.su/mira/unsloth-train-scripts.git
-check_command "Git clone"
+run_command() {
+    if ! "$@"; then
+        error_exit "Command failed: $*"
+    fi
+}
 
-echo "Downloading and running unsloth auto-install script..."
-INSTALL_CMD=$(wget -qO- https://raw.githubusercontent.com/unslothai/unsloth/main/unsloth/_auto_install.py | python -)
-check_command "Auto-install script"
+check_root() {
+    if [ "$(id -u)" -ne 0 ]; then
+        error_exit "This script must be run as root"
+    fi
+}
 
-echo "Executing installation command..."
-eval "$INSTALL_CMD"
-check_command "Installation command"
+setup_workspace() {
+    cd "$WORKSPACE_DIR" || error_exit "Cannot change to $WORKSPACE_DIR directory"
+}
 
-echo "Installing additional dependencies..."
-pip install gdown wandb
-check_command "pip install"
+install_system_packages() {
+    echo "Installing system packages..."
+    run_command apt-get update
+    run_command apt-get install -y cmake tmux vim
+}
 
-echo "Downloading dataset"
-gdown --fuzzy 'https://drive.google.com/file/d/1mqhq69dsSOK7ep7trTjRd3FMagTFTrzF/view?usp=sharing'
-check_command "gdown"
+clone_repository() {
+    echo "Cloning unsloth-train-scripts repository..."
+    run_command git clone "$REPO_URL"
+}
 
-echo "Enter your Weights & Biases API key when prompted"
-wandb login
-check_command "wandb login"
+install_unsloth() {
+    echo "Downloading and running unsloth auto-install script..."
+    INSTALL_CMD=$(wget -qO- "$INSTALL_SCRIPT_URL" | python -) || error_exit "Failed to get install command"
+    [ -n "$INSTALL_CMD" ] || error_exit "Empty install command received"
+    run_command eval "$INSTALL_CMD"
+}
 
-echo "Setup completed successfully."
+install_python_deps() {
+    echo "Installing additional dependencies..."
+    run_command pip install gdown wandb
+}
+
+download_dataset() {
+    echo "Downloading dataset..."
+    run_command gdown --fuzzy "$DATASET_URL"
+}
+
+setup_wandb() {
+    echo "Enter your Weights & Biases API key when prompted"
+    run_command wandb login
+}
+
+# Main execution
+main() {
+    check_root
+    setup_workspace
+    install_system_packages
+    clone_repository
+    install_unsloth
+    install_python_deps
+    download_dataset
+    setup_wandb
+    echo "Setup completed successfully."
+}
+
+main