Fix scripts

Author: nielstron

eval/download_models.py

@@ -16,6 +16,21 @@
     DEEPSEEK_CODER_33B_MODEL,
 )
 
+MRI_MODELS = [
+    CODEGEMMA_7B_MODEL,
+    STARCODER_2_7B_MODEL,
+    DEEPSEEK_CODER_1B_MODEL,
+    DEEPSEEK_CODER_7B_MODEL,
+    DEEPSEEK_CODER_33B_MODEL,
+]
+
+DLLM_MODELS = [
+    LLADA_MODEL_NAME,
+    DIFFU_CODER_MODEL_NAME,
+    DREAM_CODER_MODEL_NAME,
+    DREAM_MODEL_NAME,
+]
+
 
 def main():
     # Default models
@@ -39,10 +54,14 @@ def main():
         "--models",
         type=str,
         default=",".join(default_models),
-        help=f"Comma-separated list of model names to load.\nDefault: {','.join(default_models)}",
+        help=f"Comma-separated list of model names to load.\nDefault: {','.join(default_models)}.",
     )
     args = parser.parse_args()
     models = args.models.split(",")
+    if "mri" in models:
+        models = MRI_MODELS
+    elif "dllm" in models:
+        models = DLLM_MODELS
 
     # Load models
     for model in models:

install_datacrunch.sh

@@ -0,0 +1,24 @@
+set -ex
+
+sudo apt update -y
+sudo apt install python3-pip python3-venv git wget -y #  nvitop
+python3 -m venv ~/venv
+source ~/venv/bin/activate
+# set up huggingface
+pip install -U "huggingface_hub[cli]"
+huggingface-cli login --add-to-git-credential
+
+# set up typesafe_llm
+ssh-keygen -t ed25519 -f ~/.ssh/id_ed25519 -N "" || echo "SSH key already exists"
+cat ~/.ssh/id_ed25519.pub
+echo "Add the above SSH key to your GitLab account at"
+echo "https://gitlab.inf.ethz.ch/-/user_settings/ssh_keys"
+read -n 1 -p "Press Enter after adding the SSH key to continue..."
+git clone git@gitlab.inf.ethz.ch:nmuendler/diffusion_constraint.git
+cd diffusion_constraint
+bash install_rust.sh
+. "$HOME/.cargo/env"
+pip install -r requirements.txt
+echo "Restart the shell again to complete the installation"
+
+echo "export HF_HUB=/ephemeral/hub" >> "$HOME/.bashrc"