scripts/vllm_eval.py

#!/usr/bin/env python3
"""
Full reproduction: Evaluate Qwen3-8B on FHIR-AgentBench using vLLM.
Run on HF GPU Job or Modal.
"""
import json
import csv
import random
import time
import os
from vllm import LLM, SamplingParams

MODEL = os.environ.get("MODEL", "Qwen/Qwen3-8B")
DATA_PATH = os.environ.get("DATA_PATH", "/data/questions_answers_sql_fhir.csv")
OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "/data")
NUM_QUESTIONS = int(os.environ.get("NUM_QUESTIONS", "50"))
SPLIT = os.environ.get("SPLIT", "valid")
MAX_TURNS = int(os.environ.get("MAX_TURNS", "6"))
TEMPERATURE = float(os.environ.get("TEMPERATURE", "0.1"))

print(f"Model: {MODEL}")
print(f"Data: {DATA_PATH}")
print(f"Questions: {NUM_QUESTIONS} (split={SPLIT})")

# Load dataset
rows = []
with open(DATA_PATH) as f:
    reader = csv.DictReader(f)
    for r in reader:
        if r["split"] == SPLIT:
            rows.append(r)
random.seed(42)
random.shuffle(rows)
rows = rows[:NUM_QUESTIONS]
print(f"Loaded {len(rows)} {SPLIT} questions")

# Initialize vLLM
print("Loading model...")
llm = LLM(
    model=MODEL,
    tensor_parallel_size=1,
    max_model_len=8192,
    dtype="bfloat16",
    enable_auto_tool_choice=True,
    tool_call_parser="qwen3",
)
tokenizer = llm.get_tokenizer()

sampling_params = SamplingParams(
    temperature=TEMPERATURE,
    max_tokens=2048,
    stop=["<|im_end|>", "<|tool_call|>"],
)

SYSTEM_PROMPT = """You are a FHIR data analyst. Answer patient data questions by querying a FHIR server.

Rules:
- Every claim must trace to a print() output or computation.
- If unsure about a resource's schema, print a sample first.
- Keep your reasoning brief.
- When done, call finish."""

results = []
start_time = time.time()

for idx, row in enumerate(rows):
    question = row["question"]
    patient_fhir_id = row["patient_fhir_id"]
    true_answer = row["true_answer"]

    prompt = tokenizer.apply_chat_template(
        [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": f"Patient FHIR ID: {patient_fhir_id}\n\nQuestion: {question}"},
        ],
        tokenize=False,
        add_generation_prompt=True,
    )

    t0 = time.time()
    outputs = llm.generate([prompt], sampling_params)
    generated = outputs[0].outputs[0].text.strip()
    t1 = time.time()

    results.append({
        "question_id": row["question_id"],
        "question": question,
        "patient_fhir_id": patient_fhir_id,
        "true_answer": true_answer,
        "generated": generated,
        "latency_seconds": t1 - t0,
    })

    if (idx + 1) % 10 == 0:
        print(f"[{idx+1}/{len(rows)}] done. Elapsed: {time.time()-start_time:.1f}s")

total_time = time.time() - start_time

# Save results
out_path = os.path.join(OUTPUT_DIR, "vllm_results.json")
with open(out_path, "w") as f:
    json.dump({
        "model": MODEL,
        "num_questions": NUM_QUESTIONS,
        "split": SPLIT,
        "total_time_seconds": total_time,
        "avg_latency": total_time / len(results) if results else 0,
        "results": results,
    }, f, indent=2)

print(f"\nDone! {total_time:.1f}s total, {total_time/len(results):.1f}s avg per question")
print(f"Results saved to {out_path}")