Implement head_qa benchmark for inference evaluation pipeline

.gitignore

@@ -0,0 +1,4 @@
+.env
+test.ipynb
+test.sh
+benchmarks

evaluation/benchmarks.py

@@ -28,6 +28,7 @@ def benchmark_factory(name):
     """
     # Note: benchmark is instantiated *after* selection.
     factories = {
+        "head_qa": HeadQA,
         "medmcqa": MedMCQA,
         "pubmedqa": ClosedPubMedQA,
         "open_pubmedqa": PubMedQA,
@@ -315,6 +316,33 @@ def load_generations(self, checkpoint_name):
         print("Loading {} generations from the following path: {}".format(self.name, path))
         self.generations = pd.read_json(path)
 
+class HeadQA(Benchmark):
+    '''
+    HEAD-QA is a multi-choice HEAlthcare Dataset. The questions come from exams 
+    to access a specialized position in the Spanish healthcare system, and are 
+    challenging even for highly specialized humans.
+
+    Huggingface card: https://huggingface.co/datasets/boysack/head_qa
+    '''
+    def __init__(self, name='head_qa') -> None:
+        super().__init__(name)
+        self.hub_name = 'boysack/head_qa'
+        self.dir_name = 'boysack___head_qa'
+        self.path = os.path.join(ROOT_DIR, 'benchmarks', 'datasets', self.dir_name)
+        self.splits = ['train', 'validation', 'test']
+        self.num_options = 5 # 5 train, 4 validation and test
+        #self.subsets = ['<If subset exist, specify here>'] # es and en
+
+    @staticmethod
+    def custom_preprocessing(row):
+        options = sorted(row["answers"], key=lambda answer: answer["aid"]) # needed?
+        options = [option["atext"] for option in options]
+        answer = int(row['ra'])
+        row['prompt'] = format_mcq(row['qtext'], options)
+        row['gold'] = chr(ord('A')+answer-1) if answer in range(1, len(options)+1) else None # answers are 1-based
+
+
+        return row
 
 class MedMCQA(Benchmark):
     '''

evaluation/evaluate.py

@@ -16,6 +16,7 @@
 # nltk.download('punkt')
 
 benchmark_output_type = {
+    "head_qa": "mcq",
     'pubmedqa': 'boolean',
     'newpubmedqa': 'boolean',
     'medmcqa': 'mcq',
@@ -368,7 +369,7 @@ def match_truthfulqa(generations):
             generation["subset"] = "Unknown"
 
 def main(args):
-    args.out_dir = f'{args.out_dir}/{args.benchmark}'
+    args.out_dir = f'{args.out_dir}'
 
     if args.shots > 0:
         path = f'{args.out_dir}/{args.benchmark}-{args.checkpoint}-{args.shots}-shot.jsonl'
@@ -440,7 +441,7 @@ def main(args):
         metrics["model"] = model,
         del metrics["ignored"]
 
-        wandb.init(project=args.wandb_project, entity=args.wandb_entity, name=run_name)
+        wandb.init(project=args.wandb_project, name=run_name)
         artifact = wandb.Artifact(run_name, type="dataset", metadata=metrics)
         artifact.add_file(path)
         wandb.log_artifact(artifact)

evaluation/inference.py

@@ -16,6 +16,7 @@
 logger.setLevel(logging.INFO)
 
 INSTRUCTIONS = {
+    'head_qa': {'task': 'mcq', 'partition': 'validation', 'instructions': 'head_qa'},
     'truthfulqa': {'task': 'mcq', 'partition': 'validation', 'instructions': 'truthfulqa', 'cot_col': 'exp'},
     'medmcqa': {'task': 'mcq', 'partition': 'validation', 'instructions': 'medmcqa', 'cot_col': 'exp'},
     'pubmedqa': {'task': 'mcq', 'partition': 'test', 'instructions': 'pubmedqa', 'cot_col': 'long_answer'},
@@ -28,6 +29,7 @@
 }
 
 INSTRUCTIONS_SIMPLE = {
+    'head_qa': {'task': 'mcq', 'partition': 'validation', 'instructions': 'mcp'},
     'truthfulqa': {'task': 'mcq', 'partition': 'validation', 'instructions': 'mcp', 'cot_col': 'exp'},
     'medmcqa': {'task': 'mcq', 'partition': 'validation', 'instructions': 'mcp', 'cot_col': 'exp'},
     'pubmedqa': {'task': 'mcq', 'partition': 'test', 'instructions': 'open_question', 'cot_col': 'long_answer'},
@@ -83,32 +85,20 @@ def vllm_infer(client, tokenizer, prompt, stop_seq, max_new_tokens=1024, cot=Fal
     :param temperature: float, the temperature to use for sampling
     """
 
-    response = client.generate(prompt, sampling_params=vllm.SamplingParams(
-        # See https://github.com/vllm-project/vllm/blob/main/vllm/sampling_params.py
-        best_of=1,
+    sampling_params = vllm.SamplingParams(
         presence_penalty=0.0,
         frequency_penalty=1.0,
-        top_k=-1,
-        top_p=1.0,
         temperature=temperature,
         stop=stop_seq,
-        use_beam_search=False,
         max_tokens=max_new_tokens,
-        logprobs=5
-    ))
+        # 'best_of' and 'use_beam_search' removed to prevent crashes
+        # removed 'top_k' and 'top_p' (redundant for temp=0)
+        # and dropped 'logprobs=5' to eliminate unnecessary GPU overhead.
+    )
 
-    def top_answer(logprob):
-        top_token = max(logprob, key=logprob.get)
-        output_text = tokenizer.decode(top_token, skip_special_tokens=True)
-        return output_text
+    response = client.generate(prompt, sampling_params=sampling_params)
 
-    if len(response) > 0:
-        return [r.outputs[0].text for r in response]
-
-    if not cot:
-        return top_answer(response[0].outputs[0].logprobs[0])
-    else:
-        return response[0].outputs[0].text
+    return [r.outputs[0].text for r in response]
 
 
 def format_prompt(prompt, args):
@@ -257,8 +247,14 @@ def main(args):
         logging.info(f"/pure-mlo-scratch/trial-runs/{args.checkpoint_name}")
         kwargs["download_dir"] = f"/pure-mlo-scratch/trial-runs/{args.checkpoint_name}"
 
-    if "7b" in args.checkpoint:
-        kwargs["tensor_parallel_size"] = 4
+    #if "7b" in args.checkpoint:
+    #    kwargs["tensor_parallel_size"] = 4
+
+    kwargs["enforce_eager"] = True
+    kwargs["gpu_memory_utilization"] = 0.8
+
+    if "awq" in args.checkpoint:
+        kwargs["quantization"] = "awq"
 
     client = vllm.LLM(**kwargs)
 

evaluation/inference_pipeline.sh

@@ -8,42 +8,37 @@ checkpoints=(["mpt"]="mosaicml/mpt-7b" \
              ["falcon"]="tiiuae/falcon-7b" \
              ["mistral"]="mistralai/Mistral-7B-Instruct-v0.1" \
              ["zephyr"]="HuggingFaceH4/zephyr-7b-beta" \
+             ["meditron-7b"]="epfl-llm/meditron-7b" \
+             ["meditron-7b-awq"]="TheBloke/meditron-7B-AWQ" \
              ["baseline-7b"]="/pure-mlo-scratch/llama2/converted_HF_7B_8shard/" \
              ["pmc-7b"]="/pure-mlo-scratch/trial-runs/pmc-7b/hf_checkpoints/raw/pmc-llama-7b" \
-             ["meditron-7b"]="${CHECKPOINT_DIR}meditron-7b/hf_checkpoints/raw/release/" \
              ["clinical-camel"]="wanglab/ClinicalCamel-70B" \
              ["med42"]="m42-health/med42-70b" \
-
              ["baseline-70b"]="${CHECKPOINT_DIR}baseline-70b/hf_checkpoints/raw/release/" \
              ["meditron-70b"]="${CHECKPOINT_DIR}meditron-70b/hf_checkpoints/raw/iter_23000/" \
-
              ["baseline-medmcqa"]="${CHECKPOINT_DIR}baseline-7b/hf_checkpoints/instruct/medmcqa/" \
              ["baseline-pubmedqa"]="${CHECKPOINT_DIR}baseline-7b/hf_checkpoints/instruct/pubmedqa/" \
              ["baseline-medqa"]="${CHECKPOINT_DIR}baseline-7b/hf_checkpoints/instruct/medqa/" \
              ["baseline-cotmedmcqa"]="${CHECKPOINT_DIR}baseline-7b/hf_checkpoints/instruct/cotmedmcqa/" \
              ["baseline-cotpubmedqa"]="${CHECKPOINT_DIR}baseline-7b/hf_checkpoints/instruct/cotpubmedqa/" \
              ["baseline-medical"]="${CHECKPOINT_DIR}baseline-7b/hf_checkpoints/instruct/medical/" \
-
              ["pmc-medmcqa"]="${CHECKPOINT_DIR}pmc-7b/hf_checkpoints/instruct/medmcqa/" \
              ["pmc-medqa"]="${CHECKPOINT_DIR}pmc-7b/hf_checkpoints/instruct/medqa-32/" \
              ["pmc-pubmedqa"]="${CHECKPOINT_DIR}pmc-7b/hf_checkpoints/instruct/pubmedqa/" \
              ["pmc-cotpubmedqa"]="${CHECKPOINT_DIR}pmc-7b/hf_checkpoints/instruct/cotpubmedqa/" \
              ["pmc-cotmedmcqa"]="${CHECKPOINT_DIR}pmc-7b/hf_checkpoints/instruct/cotmedmcqa/"\
              ["pmc-medical"]="${CHECKPOINT_DIR}pmc-7b/hf_checkpoints/instruct/medical/"\
-
              ["meditron-7b-medmcqa"]="${CHECKPOINT_DIR}meditron-7b/hf_checkpoints/instruct/medmcqa/" \
              ["meditron-7b-pubmedqa"]="${CHECKPOINT_DIR}meditron-7b/hf_checkpoints/instruct/pubmedqa/" \
              ["meditron-7b-medqa"]="${CHECKPOINT_DIR}meditron-7b/hf_checkpoints/instruct/medqa/" \
              ["meditron-7b-cotpubmedqa"]="${CHECKPOINT_DIR}meditron-7b/hf_checkpoints/instruct/cotpubmedqa/" \
              ["meditron-7b-cotmedmcqa"]="${CHECKPOINT_DIR}meditron-7b/hf_checkpoints/instruct/cotmedmcqa/" \
-
              ["baseline-70b-medqa"]="${CHECKPOINT_DIR}baseline-70b/hf_checkpoints/instruct/medqa/" \
              ["baseline-70b-medmcqa"]="${CHECKPOINT_DIR}baseline-70b/hf_checkpoints/instruct/medmcqa/" \
              ["baseline-70b-pubmedqa"]="${CHECKPOINT_DIR}baseline-70b/hf_checkpoints/instruct/pubmedqa/" \
              ["baseline-70b-cotmedqa"]="${CHECKPOINT_DIR}baseline-70b/hf_checkpoints/instruct/cotmedqa/" \
              ["baseline-70b-cotmedmcqa"]="${CHECKPOINT_DIR}baseline-70b/hf_checkpoints/instruct/cotmedmcqa/" \
              ["baseline-70b-cotpubmedqa"]="${CHECKPOINT_DIR}baseline-70b/hf_checkpoints/instruct/cotpubmedqa/" \
-
              ["meditron-70b-medmcqa"]="${CHECKPOINT_DIR}meditron-70b/hf_checkpoints/instruct/medmcqa/" \
              ["meditron-70b-pubmedqa"]="${CHECKPOINT_DIR}meditron-70b/hf_checkpoints/instruct/pubmedqa" \
              ["meditron-70b-medqa"]="${CHECKPOINT_DIR}meditron-70b/hf_checkpoints/instruct/medqa/" \
@@ -52,15 +47,15 @@ checkpoints=(["mpt"]="mosaicml/mpt-7b" \
              ["meditron-70b-cotmedqa-qbank"]="${CHECKPOINT_DIR}meditron-70b/hf_checkpoints/instruct/cotmedqa/" \
              ["meditron-70b-instruct"]="${CHECKPOINT_DIR}meditron-70b/hf_checkpoints/instruct/medical")
 
-CHECKPOINT_NAME=meditron-70b
-BENCHMARK=medmcqa
+CHECKPOINT_NAME=meditron-7b-awq
+BENCHMARK=head_qa
 SHOTS=0
 COT=0
 SC_COT=0
 MULTI_SEED=0
 BACKEND=vllm
 WANDB=1
-BATCH_SIZE=16
+BATCH_SIZE=4
 
 HELP_STR="[--checkpoint=$CHECKPOINT_NAME] [--benchmark=$BENCHMARK] [--help]"
 
@@ -72,8 +67,7 @@ if [[ $# = 1 ]] && [[ $1 = "-h" ]] || [[ $1 = "--help" ]]; then
 	help
 	exit 0
 elif [[ $# = 0 ]]; then
-	help
-	exit 1
+	echo "Running with default parameters..."
 fi
 
 while getopts c:b:s:r:e:m:t:d: flag
@@ -92,6 +86,10 @@ done
 
 CHECKPOINT=${checkpoints[$CHECKPOINT_NAME]}
 
+if [[ -z "$CHECKPOINT" ]]; then
+    CHECKPOINT=$CHECKPOINT_NAME
+fi
+
 echo
 echo "Running inference pipeline"
 echo "Checkpoint name: $CHECKPOINT_NAME"
@@ -105,12 +103,12 @@ echo "SC COT: $SC_COT"
 echo "BATCH_SIZE: $BATCH_SIZE"
 echo
 
-COMMON_ARGS="--checkpoint  $CHECKPOINT \
+COMMON_ARGS="--checkpoint $CHECKPOINT \
     --checkpoint_name ${CHECKPOINT_NAME} \
     --benchmark $BENCHMARK \
     --shots $SHOTS \
     --batch_size $BATCH_SIZE"
-ACC_ARGS="--checkpoint $CHECKPOINT_NAME \
+ACC_ARGS="--checkpoint $CHECKPOINT \
     --benchmark $BENCHMARK \
     --shots $SHOTS"
 
@@ -136,6 +134,14 @@ if [[ $WANDB = 1 ]]; then
     ACC_ARGS="$ACC_ARGS --wandb"
 fi
 
-echo inference.py $COMMON_ARGS
+echo python inference.py $COMMON_ARGS
 python inference.py $COMMON_ARGS
-python evaluate.py $ACC_ARGS
+
+# Safety net: only evaluate if inference succeeded
+if [ $? -eq 0 ]; then
+    echo python evaluate.py $ACC_ARGS
+    python evaluate.py $ACC_ARGS
+else
+    echo "Inference failed. Skipping evaluation."
+    exit 1
+fi

evaluation/instructions.json

@@ -1,4 +1,10 @@
 {
+    "head_qa": {
+        "system": "You are a medical expert taking a Spanish healthcare specialization exam (MiR, EiR, BiR, etc.). You need to demonstrate your understanding of clinical science and medical knowledge. For the following multiple-choice question, select the one correct answer. Base your answer on current medical guidelines.",
+        "user": "The answer is:",
+        "type": "task-oriented",
+        "source": ""
+    },
     "medqa": {
         "system": "You are a medical doctor taking the US Medical Licensing Examination. You need to demonstrate your understanding of basic and clinical science, medical knowledge, and mechanisms underlying health, disease, patient care, and modes of therapy. Show your ability to apply the knowledge essential for medical practice. For the following multiple-choice question, select one correct answer from A to E. Base your answer on the current and standard practices referenced in medical guidelines.",
         "user": "The answer is:",