PRIME/test_prime.sh at main · HySonLab/PRIME · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
#!/bin/bash

# ==========================================
# Testing Script (PRIME)
# ==========================================

# -----------------------------
# Config
# -----------------------------
DATA_CONFIG="./config/data_config.yaml"
MODEL_CONFIG="./config/model_config.yaml"

TASK="GeneOntology"   # FoldClassification | ECReaction | GeneOntology | BindingSite
BATCH_SIZE=32
CUDA_DEVICE=1

# -----------------------------
# Hierarchy Ablation
# -----------------------------
ACTIVE_LEVELS=("surface" "atom" "residue" "sse" "protein")
# ACTIVE_LEVELS=("surface" "atom" "residue" "sse")
READOUT_LEVEL="residue"

# -----------------------------
# Cross-Attention Option
# Set to "true" to test PRIME_CrossAttention
# Set to "false" to test standard PRIME
# -----------------------------
CROSS_ATTENTION="false"

# -----------------------------
# Optional for GO
# -----------------------------
GO_BRANCH="BP"   # MF | BP | CC

# -----------------------------
# Optional for FoldClassification
# -----------------------------
TEST_SET_SPLIT="fold"   # family | superfamily | fold

echo "===================================="
echo "Testing PRIME"
echo "Task:            $TASK"
echo "Batch Size:      $BATCH_SIZE"
echo "CUDA device:     $CUDA_DEVICE"
echo "Active Levels:   ${ACTIVE_LEVELS[@]}"
echo "Readout Level:   $READOUT_LEVEL"
echo "Cross Attention: $CROSS_ATTENTION"
echo "===================================="

# -----------------------------
# Build base command
# -----------------------------
CMD="CUDA_VISIBLE_DEVICES=$CUDA_DEVICE python test_prime.py \
    --data_config $DATA_CONFIG \
    --model_config $MODEL_CONFIG \
    --task $TASK \
    --batch_size $BATCH_SIZE \
    --active_levels ${ACTIVE_LEVELS[@]} \
    --readout_level $READOUT_LEVEL"

if [ "$CROSS_ATTENTION" == "true" ]; then
    CMD="$CMD --cross_attention"
fi

if [ "$TASK" == "GeneOntology" ]; then
    echo "GO Branch: $GO_BRANCH"
    CMD="$CMD --go_branch $GO_BRANCH"
elif [ "$TASK" == "FoldClassification" ]; then
    echo "Test Split: $TEST_SET_SPLIT"
    CMD="$CMD --test_set_split $TEST_SET_SPLIT"
fi

# -----------------------------
# Run
# -----------------------------
eval $CMD