feat: update experiments

.gitignore

@@ -194,3 +194,4 @@ cplus_transforms/traverse.py
 *.pdf
 
 .references/*
+/results/

AGENTS.md

@@ -136,6 +136,9 @@ cd downstream/Clone-detection-POJ-104
 ./run.sh <pretrained_model_path> <output_dir>
 ```
 
+Pre-generated per-model evaluation scripts live in `experiments_downstream/`, including
+ModernBERT variants. Regenerate them with `python3 experiments_downstream/gen_all.py`.
+
 Metrics: MAP@R for clone detection, accuracy for defect detection and code classification.
 
 ## Code Conventions

README.md

@@ -262,6 +262,23 @@ cd downstream/Clone-detection-POJ-104
 ./run.sh <pretrained_model_path> <output_dir>
 ```
 
+Pre-generated per-model evaluation scripts live in `experiments_downstream/`, including
+ModernBERT variants. Regenerate them with `python3 experiments_downstream/gen_all.py`.
+
+To launch the 8 downstream tasks (clone detection + code classification across POJ and
+CodeNet) in parallel on 8 GPUs, use the helper script in `experiments_downstream/`.
+Each subtask writes to its own output directory (e.g., `results/<model>/<task>/<subset>/`):
+
+```sh
+python experiments_downstream/run_all_downstream.py --loss supcon --model inv-codebert
+```
+
+Override GPU ids with `--gpus` (comma-separated, must be 8):
+
+```sh
+python experiments_downstream/run_all_downstream.py --loss supcon --model inv-codebert --gpus 0,1,2,3,4,5,6,7
+```
+
 To evaluate robustness, use the augmented test scripts:
 
 ```sh
@@ -276,6 +293,14 @@ Generate t-SNE visualizations of code embeddings across models:
 uv run plot/visualize.py --input_test_file dataset/aug_test.jsonl --output_file clusters.png
 ```
 
+### Results Parsing
+
+Parse downstream evaluation outputs into a regular vs augmented table:
+
+```sh
+python experiments_downstream/parse_results.py --results-root results --digits 2
+```
+
 ## Invariant Code Transformations
 
 InvPT uses six semantic-preserving transformation operators:

downstream/Clone-detection-BigCloneBench/code/run.py

@@ -882,14 +882,16 @@ def main():
     if args.do_eval and args.local_rank in [-1, 0]:
         checkpoint_prefix = "checkpoint-best-f1/model.bin"
         output_dir = os.path.join(args.output_dir, "{}".format(checkpoint_prefix))
-        model.load_state_dict(torch.load(output_dir))
+        model_to_load = model.module if hasattr(model, "module") else model
+        model_to_load.load_state_dict(torch.load(output_dir))
         model.to(args.device)
         result = evaluate(args, model, tokenizer, pool=pool)
 
     if args.do_test and args.local_rank in [-1, 0]:
         checkpoint_prefix = "checkpoint-best-f1/model.bin"
         output_dir = os.path.join(args.output_dir, "{}".format(checkpoint_prefix))
-        model.load_state_dict(torch.load(output_dir))
+        model_to_load = model.module if hasattr(model, "module") else model
+        model_to_load.load_state_dict(torch.load(output_dir))
         model.to(args.device)
         test(args, model, tokenizer, pool=pool, best_threshold=0.5)
 

downstream/Clone-detection-CodeNet/code/run.py

@@ -873,7 +873,8 @@ def main():
     if args.do_eval and args.local_rank in [-1, 0]:
         checkpoint_prefix = "checkpoint-best-map/model.bin"
         output_dir = os.path.join(args.output_dir, "{}".format(checkpoint_prefix))
-        model.load_state_dict(torch.load(output_dir), strict=False)
+        model_to_load = model.module if hasattr(model, "module") else model
+        model_to_load.load_state_dict(torch.load(output_dir), strict=False)
         model.to(args.device)
         result = evaluate(args, model, tokenizer)
         logger.info("***** Eval results *****")
@@ -883,7 +884,8 @@ def main():
     if args.do_test and args.local_rank in [-1, 0]:
         checkpoint_prefix = "checkpoint-best-map/model.bin"
         output_dir = os.path.join(args.output_dir, "{}".format(checkpoint_prefix))
-        model.load_state_dict(torch.load(output_dir), strict=False)
+        model_to_load = model.module if hasattr(model, "module") else model
+        model_to_load.load_state_dict(torch.load(output_dir), strict=False)
         model.to(args.device)
         test(args, model, tokenizer)
 

downstream/Clone-detection-CodeNet/run.sh

@@ -5,7 +5,12 @@ subset=$3
 model_type=${4:-roberta}
 tokenizer_name=${5:-roberta-base}
 
-output_dir=$save_path/$subset
+base_name=$(basename "$save_path")
+if [ "$base_name" = "$subset" ]; then
+    output_dir=$save_path
+else
+    output_dir=$save_path/$subset
+fi
 
 mkdir -p $output_dir
 touch $output_dir/train.log
@@ -40,4 +45,4 @@ python evaluator/evaluator.py \
 
 
 echo "Running evaluation for augmented test set..."
-./run_aug_test.sh $model_path $save_path $subset $model_type $tokenizer_name
+./run_aug_test.sh "$model_path" "$save_path" "$subset" "$model_type" "$tokenizer_name"

downstream/Clone-detection-CodeNet/run_aug_test.sh

@@ -5,7 +5,12 @@ subset=$3
 model_type=${4:-roberta}
 tokenizer_name=${5:-roberta-base}
 
-output_dir=$save_path/$subset
+base_name=$(basename "$save_path")
+if [ "$base_name" = "$subset" ]; then
+    output_dir=$save_path
+else
+    output_dir=$save_path/$subset
+fi
 
 mkdir -p $output_dir
 touch $output_dir/aug_train.log

downstream/Clone-detection-POJ104/code/run.py

@@ -873,7 +873,8 @@ def main():
     if args.do_eval and args.local_rank in [-1, 0]:
         checkpoint_prefix = "checkpoint-best-map/model.bin"
         output_dir = os.path.join(args.output_dir, "{}".format(checkpoint_prefix))
-        model.load_state_dict(torch.load(output_dir), strict=False)
+        model_to_load = model.module if hasattr(model, "module") else model
+        model_to_load.load_state_dict(torch.load(output_dir), strict=False)
         model.to(args.device)
         result = evaluate(args, model, tokenizer)
         logger.info("***** Eval results *****")
@@ -883,7 +884,8 @@ def main():
     if args.do_test and args.local_rank in [-1, 0]:
         checkpoint_prefix = "checkpoint-best-map/model.bin"
         output_dir = os.path.join(args.output_dir, "{}".format(checkpoint_prefix))
-        model.load_state_dict(torch.load(output_dir), strict=False)
+        model_to_load = model.module if hasattr(model, "module") else model
+        model_to_load.load_state_dict(torch.load(output_dir), strict=False)
         model.to(args.device)
         test(args, model, tokenizer)
 

downstream/Code-classification-CodeNet/code/run.py

@@ -456,7 +456,8 @@ def main():
     if args.do_eval:
         checkpoint_prefix = "checkpoint-best-acc/model.bin"
         output_dir = os.path.join(args.output_dir, "{}".format(checkpoint_prefix))
-        model.load_state_dict(torch.load(output_dir))
+        model_to_load = model.module if hasattr(model, "module") else model
+        model_to_load.load_state_dict(torch.load(output_dir))
         model.to(args.device)
         result = evaluate(args, model, tokenizer)
         logger.info("***** Eval results *****")
@@ -466,7 +467,8 @@ def main():
     if args.do_test:
         checkpoint_prefix = "checkpoint-best-acc/model.bin"
         output_dir = os.path.join(args.output_dir, "{}".format(checkpoint_prefix))
-        model.load_state_dict(torch.load(output_dir))
+        model_to_load = model.module if hasattr(model, "module") else model
+        model_to_load.load_state_dict(torch.load(output_dir))
         model.to(args.device)
         result = test(args, model, tokenizer)
         logger.info("***** Test results *****")

downstream/Code-classification-CodeNet/run.sh

@@ -5,7 +5,12 @@ subset=$3
 model_type=${4:-roberta}
 tokenizer_name=${5:-microsoft/codebert-base}
 
-output_dir=$save_path/$subset
+base_name=$(basename "$save_path")
+if [ "$base_name" = "$subset" ]; then
+    output_dir=$save_path
+else
+    output_dir=$save_path/$subset
+fi
 
 mkdir -p $output_dir
 touch $output_dir/test_train.log
@@ -31,4 +36,4 @@ python ./code/run.py \
 
 
 echo "Running evaluation for augmented test set..."
-./run_aug_test.sh $model_path $save_path $subset $model_type $tokenizer_name
+./run_aug_test.sh "$model_path" "$save_path" "$subset" "$model_type" "$tokenizer_name"

downstream/Code-classification-CodeNet/run_aug_test.sh

@@ -5,7 +5,12 @@ subset=$3
 model_type=${4:-roberta}
 tokenizer_name=${5:-microsoft/codebert-base}
 
-output_dir=$save_path/$subset
+base_name=$(basename "$save_path")
+if [ "$base_name" = "$subset" ]; then
+    output_dir=$save_path
+else
+    output_dir=$save_path/$subset
+fi
 
 mkdir -p $output_dir
 touch $output_dir/test_train.log

downstream/Code-classification-POJ104/code/run.py

@@ -53,6 +53,22 @@
     "modernbert": (AutoConfig, AutoModelForSequenceClassification, AutoTokenizer),
 }
 
+
+def resolve_model_type(model_type, model_name_or_path):
+    if model_type in MODEL_CLASSES:
+        return model_type
+    candidate = model_name_or_path or model_type or ""
+    name = candidate.lower()
+    resolved = "modernbert" if "modernbert" in name else "roberta"
+    logger.warning(
+        "Unknown model_type '%s'; inferring '%s' from '%s'.",
+        model_type,
+        resolved,
+        candidate,
+    )
+    return resolved
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -82,7 +98,8 @@ def convert_examples_to_features(js, tokenizer, args):
 
 
 class TextDataset(Dataset):
-    def __init__(self, tokenizer, args, file_path=None):
+    def __init__(self, tokenizer, args, file_path: str):
+        file_path = str(file_path)
         self.examples = []
         with open(file_path) as f:
             for line in f:
@@ -429,6 +446,7 @@ def main():
     # Set seed
     set_seed(args.seed)
 
+    args.model_type = resolve_model_type(args.model_type, args.model_name_or_path)
     config_class, model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
     config = config_class.from_pretrained(args.model_name_or_path)
     config.num_labels = 104
@@ -456,7 +474,8 @@ def main():
     if args.do_eval:
         checkpoint_prefix = "checkpoint-best-acc/model.bin"
         output_dir = os.path.join(args.output_dir, "{}".format(checkpoint_prefix))
-        model.load_state_dict(torch.load(output_dir))
+        model_to_load = model.module if hasattr(model, "module") else model
+        model_to_load.load_state_dict(torch.load(output_dir))
         model.to(args.device)
         result = evaluate(args, model, tokenizer)
         logger.info("***** Eval results *****")
@@ -466,7 +485,8 @@ def main():
     if args.do_test:
         checkpoint_prefix = "checkpoint-best-acc/model.bin"
         output_dir = os.path.join(args.output_dir, "{}".format(checkpoint_prefix))
-        model.load_state_dict(torch.load(output_dir))
+        model_to_load = model.module if hasattr(model, "module") else model
+        model_to_load.load_state_dict(torch.load(output_dir))
         model.to(args.device)
         result = test(args, model, tokenizer)
         logger.info("***** Test results *****")

downstream/Code-classification-POJ104/run.sh

@@ -5,22 +5,32 @@ subset=$3
 model_type=${4:-roberta}
 tokenizer_name=${5:-microsoft/codebert-base}
 
-output_dir=$save_path/$subset
+base_name=$(basename "$save_path")
+if [ "$base_name" = "$subset" ]; then
+    output_dir=$save_path
+else
+    output_dir=$save_path/$subset
+fi
 
 mkdir -p $output_dir
 touch $output_dir/test_train.log
 
 echo "Running fine-tuning for POJ104"
+train_suffix=""
+if [ -n "$subset" ]; then
+    train_suffix="/$subset"
+fi
+
 python ./code/run.py \
     --output_dir=$output_dir \
     --model_type=$model_type \
     --tokenizer_name=$tokenizer_name \
     --model_name_or_path=$model_path \
     --do_train \
     --do_test \
-    --train_data_file=./dataset/$subset/train.jsonl \
-    --eval_data_file=./dataset/$subset/valid.jsonl \
-    --test_data_file=./dataset/$subset/test.jsonl \
+    --train_data_file=./dataset${train_suffix}/train.jsonl \
+    --eval_data_file=./dataset${train_suffix}/valid.jsonl \
+    --test_data_file=./dataset${train_suffix}/test.jsonl \
     --num_train_epochs 10 \
     --block_size 512 \
     --train_batch_size 32 \
@@ -30,4 +40,4 @@ python ./code/run.py \
     --seed 123456  2>&1 | tee $output_dir/test_train.log
 
 echo "Running evaluation for augmented test set..."
-./run_aug_test.sh $model_path $save_path $subset $model_type $tokenizer_name
+./run_aug_test.sh "$model_path" "$save_path" "$subset" "$model_type" "$tokenizer_name"

downstream/Code-classification-POJ104/run_aug_test.sh

@@ -5,20 +5,30 @@ subset=$3
 model_type=${4:-roberta}
 tokenizer_name=${5:-microsoft/codebert-base}
 
-output_dir=$save_path/$subset
+base_name=$(basename "$save_path")
+if [ "$base_name" = "$subset" ]; then
+    output_dir=$save_path
+else
+    output_dir=$save_path/$subset
+fi
 
 mkdir -p $output_dir
 touch $output_dir/test_train.log
 
+train_suffix=""
+if [ -n "$subset" ]; then
+    train_suffix="/$subset"
+fi
+
 python ./code/run.py \
     --output_dir=$output_dir \
     --model_type=$model_type \
     --tokenizer_name=$tokenizer_name \
     --model_name_or_path=$model_path \
     --do_test \
-    --train_data_file=./dataset/$subset/train.jsonl \
-    --eval_data_file=./dataset/$subset/valid.jsonl \
-    --test_data_file=./dataset/$subset/aug_test.jsonl \
+    --train_data_file=./dataset${train_suffix}/train.jsonl \
+    --eval_data_file=./dataset${train_suffix}/valid.jsonl \
+    --test_data_file=./dataset${train_suffix}/aug_test.jsonl \
     --num_train_epochs 5 \
     --block_size 256 \
     --train_batch_size 8 \

downstream/Defect-detection/code/run.py

@@ -819,7 +819,8 @@ def main():
     if args.do_eval and args.local_rank in [-1, 0]:
         checkpoint_prefix = "checkpoint-best-acc/model.bin"
         output_dir = os.path.join(args.output_dir, "{}".format(checkpoint_prefix))
-        model.load_state_dict(torch.load(output_dir))
+        model_to_load = model.module if hasattr(model, "module") else model
+        model_to_load.load_state_dict(torch.load(output_dir))
         model.to(args.device)
         result = evaluate(args, model, tokenizer)
         logger.info("***** Eval results *****")
@@ -829,7 +830,8 @@ def main():
     if args.do_test and args.local_rank in [-1, 0]:
         checkpoint_prefix = "checkpoint-best-acc/model.bin"
         output_dir = os.path.join(args.output_dir, "{}".format(checkpoint_prefix))
-        model.load_state_dict(torch.load(output_dir))
+        model_to_load = model.module if hasattr(model, "module") else model
+        model_to_load.load_state_dict(torch.load(output_dir))
         model.to(args.device)
         test(args, model, tokenizer)
 

experiments_downstream/Clone-detection-BigCloneBench/clone-bcb_codebert.sh

@@ -1,8 +1,14 @@
 #!/bin/bash
 # Downstream evaluation: Clone-detection-BigCloneBench with codebert
 set -euo pipefail
+
+# Parse CUDA device argument (default: 0)
+CUDA_DEVICE="${1:-0}"
+
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 ROOT_DIR="$(cd "$SCRIPT_DIR/../.." && pwd)"
 
+export CUDA_VISIBLE_DEVICES="$CUDA_DEVICE"
+
 cd "$ROOT_DIR/downstream/Clone-detection-BigCloneBench"
 ./run.sh microsoft/codebert-base "$ROOT_DIR/results/codebert/Clone-detection-BigCloneBench"

experiments_downstream/Clone-detection-BigCloneBench/clone-bcb_contrabert_c.sh

@@ -1,8 +1,14 @@
 #!/bin/bash
 # Downstream evaluation: Clone-detection-BigCloneBench with contrabert_c
 set -euo pipefail
+
+# Parse CUDA device argument (default: 0)
+CUDA_DEVICE="${1:-0}"
+
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 ROOT_DIR="$(cd "$SCRIPT_DIR/../.." && pwd)"
 
+export CUDA_VISIBLE_DEVICES="$CUDA_DEVICE"
+
 cd "$ROOT_DIR/downstream/Clone-detection-BigCloneBench"
 ./run.sh "$ROOT_DIR/saved_models/ContraBERT_C" "$ROOT_DIR/results/contrabert_c/Clone-detection-BigCloneBench"