Merge pull request #1 from cxcscmu/haok/refactoring

Haok/refactoring

Author: haok1402

.gitignore

@@ -1 +1,2 @@
+wandb
 **/__pycache__

.vscode/settings.json

@@ -9,7 +9,6 @@
   },
   "[python]": {
     "editor.tabSize": 4,
-    "editor.formatOnSave": true,
     "editor.defaultFormatter": "ms-python.black-formatter",
   },
   "[jsonc]": {

environment.yml

@@ -6,25 +6,29 @@ channels:
   - conda-forge
 
 dependencies:
-  - python==3.12
-  - pip==24.2
+  - python=3.12.0
+  - pip=24.2
   - pytorch=2.4.0
-  - pytorch-cuda==12.4
-  - numpy==1.26.4
-  - rich==13.7.1
-  - pyarrow==16.1.0
-  - mypy==1.10.0
-  - aiofiles==22.1.0
-  - aiohttp==3.9.5
-  - transformers==4.44.2
-  - pytest==7.4.4
-  - pytest-asyncio==0.20.3
-  - types-aiofiles==24.1.0.20240626
-  - attrs==23.1.0
-  - wandb==0.16.6
-  - faiss-gpu==1.8.0
-  - blobfile==3.0.0
-  - elasticsearch==8.15.1
-  - pillow==9.5.0
+  - pytorch-cuda=12.4
+  - numpy=1.26.4
+  - rich=13.7.1
+  - pyarrow=16.1.0
+  - mypy=1.10.0
+  - aiofiles=22.1.0
+  - aiohttp=3.9.5
+  - transformers=4.44.2
+  - pytest=7.4.4
+  - pytest-asyncio=0.20.3
+  - types-aiofiles=24.1.0.20240626
+  - attrs=23.1.0
+  - wandb
+  - faiss-gpu=1.8.0
+  - blobfile=3.0.0
+  - elasticsearch=8.15.1
+  - pillow
+  - seaborn
+  - sentencepiece
   - pip:
-      - treevizer==0.2.4
+      - treevizer
+      - beir
+      - flash-attn

script/acquire.sh

@@ -1,9 +1,17 @@
 #!/usr/bin/bash
 
-# Acquire a compute node with 32 CPUs, 96GB RAM, and 1 A6000 GPU.
+# Validate the arguments
+if [ -z "$1" ]; then
+    echo "Usage: $0 <num_gpu>"
+    exit 1
+fi
+
+NUM_GPU=$1
+
+# Acquire a compute node with 32 CPUs, 96GB RAM, and specified number of GPUs.
 # The node will be acquired for 7 days, and the session will be interactive.
 # Please run this script with tmux to avoid losing the session.
 srun \
-    --partition=long --time=07-00:00:00 \
-    --cpus-per-task=32 --mem=96GB --gres=gpu:A6000:1 \
-    --pty bash
+        --partition=long --time=07-00:00:00 \
+        --cpus-per-task=32 --mem=96GB --gres=gpu:A6000:$NUM_GPU \
+        --pty bash

source/dataset/__init__.py

@@ -5,3 +5,4 @@
 workspace.mkdir(mode=0o770, parents=True, exist_ok=True)
 
 from source.dataset.msMarco import MsMarcoDataset
+from source.dataset.beir import BeirDataset