update beaker config

README.md

@@ -14,9 +14,16 @@ git clone https://github.com/davidheineman/acl-search
 pip install -r requirements.txt 
 python src/server.py # (this will download a pre-built index!)
 
-# (getting pip errors?)
+# getting pip errors? (install sentencepiece deps)
 sudo apt-get update
 sudo apt-get install -y pkg-config libsentencepiece-dev
+
+# running on CUDA? (fix broken package path)
+INSTALL_PATH=PATH_TO_YOUR_PYTHON_INSTALL # e.g., /root/ai2/miniconda3/envs/acl_search/lib/python3.10
+cp ./src/extras/segmented_maxsim.cpp $INSTALL_PATH/site-packages/colbert/modeling/segmented_maxsim.cpp
+cp ./src/extras/decompress_residuals.cpp $INSTALL_PATH/site-packages/colbert/search/decompress_residuals.cpp
+cp ./src/extras/filter_pids.cpp $INSTALL_PATH/site-packages/colbert/search/filter_pids.cpp
+cp ./src/extras/segmented_lookup.cpp $INSTALL_PATH/site-packages/colbert/search/segmented_lookup.cpp
 ```
 
 ## More Features
@@ -71,15 +78,14 @@ fly launch
 
 **Update Index on HF**
 ```sh
-# For a full pipeline to update an index, see: src/scrape/beaker/index.sh
+# Download a fresh set of papers, index and push to hf:
+chmod +x src/scrape/beaker/index.sh
+./src/scrape/beaker/index.sh
 
 # Build and deploy container for auto-updating:
 docker build -t acl-search -f src/scrape/beaker/Dockerfile .
 docker run -it -e HF_TOKEN=$HF_TOKEN acl-search # (Optional) test it out!
 
-# Run on slurm
-sbatch src/scrape/slurm.sh
-
 # Run on beaker
 beaker image delete davidh/acl-search
 beaker image create --name acl-search acl-search

src/index.py

@@ -60,7 +60,7 @@ def main():
     with open(DATASET_PATH, 'r', encoding='utf-8') as f:
         dataset = json.loads(f.read())
 
-    # dataset = dataset[:5000] # 5K in 48s/iter on 2 A40s (67K in 2hr)
+    # dataset = dataset[:100] # 5K in 48s/iter on 2 A40s (67K in 2hr)
 
     # Get the abstracts + titles for indexing
     collection = [e.get('title', '') + '\n\n' + e.get('abstract', '') for e in dataset]

src/scrape/beaker/Dockerfile

@@ -1,11 +1,15 @@
-FROM ghcr.io/allenai/pytorch:2.4.0-cuda12.1-python3.11
+FROM ghcr.io/allenai/cuda:12.1-cudnn8-dev-ubuntu20.04-v1.2.118
 
-ENV CUDA_HOME=/opt/conda
+# Set environment variables for CUDA
+ENV CUDA_HOME=/usr/local/cuda
+ENV PATH=/usr/local/cuda/bin:$PATH
+ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
 
 WORKDIR /app
 
-# Install sentencepiece deps
+# Install sentencepiece deps (and git lfs)
 RUN apt-get update && apt-get install -y \
+    git-lfs \
     pkg-config \
     libsentencepiece-dev \
     && apt-get clean
@@ -19,7 +23,7 @@ COPY ./src/extras/decompress_residuals.cpp /opt/conda/lib/python3.11/site-packag
 COPY ./src/extras/filter_pids.cpp /opt/conda/lib/python3.11/site-packages/colbert/search/filter_pids.cpp
 COPY ./src/extras/segmented_lookup.cpp /opt/conda/lib/python3.11/site-packages/colbert/search/segmented_lookup.cpp
 
-# Copy repo
+# Copy repo, openreview creds
 COPY . .
 COPY .openreview .
 
@@ -32,7 +36,7 @@ CMD ["./src/scrape/beaker/index.sh"]
 
 # docker build -t acl-search -f src/scrape/beaker/Dockerfile .
 # docker run -it acl-search
-# docker run -it -e HF_TOKEN=$HF_TOKEN acl-search 
+# docker run -it -e HF_TOKEN=$HF_TOKEN acl-search
 # docker run -it --gpus '"device=0"' -e HF_TOKEN=$HF_TOKEN acl-search
 # # docker run --rm acl-search
 # beaker image delete davidh/acl-search

src/scrape/beaker/beaker-conf.yml

@@ -11,15 +11,15 @@ tasks:
     result:
       path: /output
     resources:
-      gpuCount: 1
+      gpuCount: 4
     context:
       priority: low
       preemptible: true
     constraints:
       cluster:
-        - ai2/general-cirrascale # no weka
-        - ai2/pluto-cirrascale # no weka
-        - ai2/jupiter-cirrascale-2 # h100
+        # - ai2/general-cirrascale # no weka
+        # - ai2/pluto-cirrascale # no weka
+        # - ai2/jupiter-cirrascale-2 # h100 # <- doesn't work with ColBERT's PyTorch 1
+        # - ai2/allennlp-elara-cirrascale # h100 # <- doesn't work with ColBERT's PyTorch 1
         - ai2/saturn-cirrascale # a100
-        - ai2/allennlp-elara-cirrascale # h100
         - ai2/neptune-cirrascale # a100

src/scrape/beaker/index.sh

@@ -1,8 +1,14 @@
 #!/bin/bash
 
+echo "Git LFS version:"
+git lfs version
+
 echo "Logged in as:"
 huggingface-cli whoami
 
+git config --global user.email "dheineman3@gatech.edu"
+git config --global user.name "davidheineman"
+
 # clone hf repository
 git clone https://huggingface.co/davidheineman/colbert-acl
 mv colbert-acl ../hf
@@ -21,9 +27,15 @@ mkdir -p ../hf
 cp -r index ../hf
 cp data/papers.json ../hf/papers.json
 
+# Fix hf token spacing
+export HF_TOKEN=$(echo "$HF_TOKEN" | tr -d '
+' | tr -d '')
+
 # push changes
 cd ../hf
+git remote set-url origin https://davidheineman:$HF_TOKEN@huggingface.co/davidheineman/colbert-acl
 git add .
+git status
 git commit -m "update index"
 git push