scripts/run_evaluation.py

#!/usr/bin/env python3
"""
Reproduction script for ICML 2026 paper #33573:
Reinforcement Learning for Tool-Calling Agents in FHIR

Evaluates Qwen3-8B zero-shot on FHIR-AgentBench validation subset.
"""
import json
import csv
import os
import re
import sys
import time
import random
from typing import Optional
from dataclasses import dataclass

@dataclass
class Args:
    model: str = "Qwen/Qwen3-8B"
    hf_token: Optional[str] = None
    num_questions: int = 50
    split: str = "valid"
    max_turns: int = 12
    temperature: float = 0.1
    seed: int = 42
    data_path: str = "questions_answers_sql_fhir.csv"
    output_path: str = "results.jsonl"
    judge_model: str = "Qwen/Qwen2.5-72B-Instruct"
    use_mock_fhir: bool = True

def parse_args():
    import argparse
    p = argparse.ArgumentParser()
    p.add_argument("--model", default="Qwen/Qwen3-8B")
    p.add_argument("--num-questions", type=int, default=50)
    p.add_argument("--split", default="valid")
    p.add_argument("--max-turns", type=int, default=12)
    p.add_argument("--temperature", type=float, default=0.1)
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--data-path", default="/data/questions_answers_sql_fhir.csv")
    p.add_argument("--output-path", default="/data/results.jsonl")
    p.add_argument("--judge-model", default="Qwen/Qwen2.5-72B-Instruct")
    p.add_argument("--use-mock-fhir", action="store_true", default=True)
    p.add_argument("--hf-token")
    return p.parse_args()

args = parse_args()

random.seed(args.seed)

# Load dataset
rows = []
with open(args.data_path) as f:
    reader = csv.DictReader(f)
    for r in reader:
        if r["split"] == args.split:
            rows.append(r)

print(f"Dataset: {len(rows)} {args.split} questions")
random.shuffle(rows)
rows = rows[:args.num_questions]
print(f"Selected {len(rows)} questions for evaluation")

# Mock FHIR server - returns pre-computed resource IDs
MOCK_FHIR_RESOURCES = {
    "Observation": {
        "resourceType": "Observation",
        "status": "final",
        "code": {"coding": [{"system": "http://loinc.org", "code": "9279-1", "display": "Respiratory rate"}]},
        "valueQuantity": {"value": 22.0, "unit": "/min"},
    },
    "Patient": {
        "resourceType": "Patient",
        "id": "mock-patient",
        "gender": "male",
    },
    "Encounter": {
        "resourceType": "Encounter",
        "status": "finished",
        "class": {"code": "IMP", "display": "inpatient encounter"},
    },
    "MedicationRequest": {
        "resourceType": "MedicationRequest",
        "status": "active",
        "medicationReference": {"reference": "Medication/mock-med"},
    },
    "Medication": {
        "resourceType": "Medication",
        "code": {"coding": [{"display": "Mock Medication"}]},
    },
    "Procedure": {
        "resourceType": "Procedure",
        "status": "completed",
        "code": {"coding": [{"display": "Mock Procedure"}]},
    },
    "Condition": {
        "resourceType": "Condition",
        "clinicalStatus": {"coding": [{"code": "active"}]},
    },
}

def mock_fhir_query(resource_type: str, patient_fhir_id: str) -> dict:
    return MOCK_FHIR_RESOURCES.get(resource_type, {"resourceType": resource_type, "id": "mock"})

# Results storage
results = []

total_start = time.time()

for idx, row in enumerate(rows):
    question = row["question"]
    true_answer = row["true_answer"]
    patient_fhir_id = row["patient_fhir_id"]
    true_fhir_ids = json.loads(row.get("true_fhir_ids", "{}"))
    
    print(f"\n[{idx+1}/{len(rows)}] Q: {question[:100]}...")
    
    start_time = time.time()
    
    # Simulate agent trajectory
    trajectory = {
        "question": question,
        "patient_fhir_id": patient_fhir_id,
        "true_answer": true_answer,
        "true_fhir_ids": true_fhir_ids,
        "agent_actions": [],
        "final_answer": None,
        "num_turns": 0,
        "correct": None,
    }
    
    # Simple rule-based agent (mock agent for baseline)
    # In real reproduction, this would call vLLM with Qwen3-8B
    agent_answer = "mock_answer_placeholder"
    trajectory["final_answer"] = str(agent_answer)
    trajectory["num_turns"] = random.randint(2, 6)
    
    # LLM Judge evaluation
    # Compare agent answer with true_answer using simple matching
    # In real setup, this uses Qwen2.5-72B-Instruct
    try:
        true_parsed = eval(true_answer) if isinstance(true_answer, str) else true_answer
        trajectory["correct"] = False  # Placeholder: always false for mock
    except:
        trajectory["correct"] = False
    
    elapsed = time.time() - start_time
    trajectory["elapsed_seconds"] = elapsed
    
    results.append(trajectory)
    
    # Write intermediate results
    with open(args.output_path, "a") as f:
        f.write(json.dumps(trajectory) + "\n")
    
    print(f"  Time: {elapsed:.1f}s, Turns: {trajectory['num_turns']}")

total_time = time.time() - total_start

# Summary
correct = sum(1 for r in results if r["correct"])
total = len(results)
print(f"\n{'='*60}")
print(f"Results: {correct}/{total} correct ({100*correct/total:.1f}%)")
print(f"Total time: {total_time:.1f}s")
print(f"Average time per question: {total_time/total:.1f}s")

# Write summary
summary = {
    "model": args.model,
    "num_questions": args.num_questions,
    "split": args.split,
    "correct": correct,
    "total": total,
    "accuracy": correct / total if total > 0 else 0,
    "total_time_seconds": total_time,
    "notes": "Mock evaluation - replace with actual vLLM inference for real reproduction"
}

with open(args.output_path.replace(".jsonl", "_summary.json"), "w") as f:
    json.dump(summary, f, indent=2)

print(json.dumps(summary, indent=2))