From 6e5c5119544378b87e2b3a6b853666b0b1df0b8d Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Tue, 3 Dec 2024 21:59:07 +0000
Subject: [PATCH] add install cript

---
 .gitignore                 |  1 +
 scripts/install/install.sh | 64 ++++++++++++++++++++++++++++++++++++++
 src/zeroband/data.py       |  2 +-
 3 files changed, 66 insertions(+), 1 deletion(-)
 create mode 100755 scripts/install/install.sh

diff --git a/.gitignore b/.gitignore
index 5411fae9..463b903a 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,6 +1,7 @@
 .vscode/*
 logs/*
 wandb/*
+datasets/*
 
 # Byte-compiled / optimized / DLL files
 __pycache__/
diff --git a/scripts/install/install.sh b/scripts/install/install.sh
new file mode 100755
index 00000000..5a4361de
--- /dev/null
+++ b/scripts/install/install.sh
@@ -0,0 +1,64 @@
+#!/usr/bin/env bash
+
+set -e
+
+# Colors for output
+GREEN='\033[0;32m'
+NC='\033[0m' # No Color
+
+log_info() {
+    echo -e "${GREEN}[INFO]${NC} $1"
+}
+
+main() {
+    # Check if sudo is installed
+    if ! command -v sudo &> /dev/null; then
+        apt update
+        apt install sudo -y
+    fi
+
+    log_info "Updating apt..."
+    sudo apt update
+
+    log_info "Installing cmake python3-dev..."
+    sudo apt install python3-dev cmake -y
+
+    log_info "Installing iperf..."
+    sudo apt install iperf -y
+
+    log_info "Cloning repository..."
+    git clone https://github.com/PrimeIntellect-ai/prime.git
+    
+    log_info "Entering project directory..."
+    cd prime
+    
+    log_info "Installing uv..."
+    curl -LsSf https://astral.sh/uv/install.sh | sh
+    
+    log_info "Sourcing uv environment..."
+    source $HOME/.local/bin/env
+        
+    log_info "Creating virtual environment..."
+    uv venv
+    
+    log_info "Activating virtual environment..."
+    source .venv/bin/activate
+    
+    log_info "Installing dependencies..."
+    uv sync --extra all
+    
+    log_info "Installing flash-attn..."
+    uv pip install flash-attn --no-build-isolation
+    
+    log_info "Updating git submodules..."
+    git submodule update --init --recursive
+    
+    log_info "Downloading data..."
+    mkdir -p datasets
+    uv run python scripts/subset_data.py --dataset_name PrimeIntellect/fineweb-edu --data_world_size 1 --data_rank 0 --max_shards 128
+    mv fineweb-edu/ datasets/fineweb-edu/
+
+    log_info "Installation completed! You can double check that everything is install correctly by running 'GLOO_SOCKET_IFNAME=lo GLOBAL_ADDR=localhost GLOBAL_RANK=0 GLOBAL_UNIQUE_ID=0 GLOBAL_WORLD_SIZE=1 GLOBAL_PORT=8989  uv run torchrun --nproc_per_node=2 src/zeroband/train.py  @configs/debug/diloco.toml'"
+}
+
+main
\ No newline at end of file
diff --git a/src/zeroband/data.py b/src/zeroband/data.py
index 7cb7f9f6..be1219db 100644
--- a/src/zeroband/data.py
+++ b/src/zeroband/data.py
@@ -25,7 +25,7 @@
 
 
 class DataConfig(BaseConfig):
-    dataset_name_or_paths: str = "/data/datasets/fineweb-edu"
+    dataset_name_or_paths: str = "datasets/fineweb-edu"
     val_dataset_name_or_paths: Optional[str] = None
     seq_length: int = 1024
     fake: bool = False