scripts/hf_job_eval.py

#!/usr/bin/env python3
"""HF Job evaluation script for FHIR reproduction."""
import json
import csv
import random
import time
import os
import argparse
import subprocess
import sys

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default="Qwen/Qwen3-8B")
    parser.add_argument("--data-path", required=True)
    parser.add_argument("--num-questions", type=int, default=20)
    parser.add_argument("--split", default="valid")
    parser.add_argument("--output", default="/data/results.json")
    args = parser.parse_args()

    # Install vLLM if not present
    try:
        from vllm import LLM, SamplingParams
    except ImportError:
        print("Installing vLLM...")
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "vllm==1.2.0"])
        from vllm import LLM, SamplingParams

    # 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)
    random.seed(42)
    random.shuffle(rows)
    rows = rows[:args.num_questions]
    print(f"Loaded {len(rows)} {args.split} questions")

    # Init vLLM
    print(f"Loading {args.model}...")
    llm = LLM(
        model=args.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=0.1,
        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[:120],
            "true_answer": str(true_answer)[:120] if true_answer else "",
            "generated": generated[:500],
            "latency_seconds": t1 - t0,
        })

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

    total_time = time.time() - start_time

    output = {
        "model": args.model,
        "num_questions": args.num_questions,
        "split": args.split,
        "total_time_seconds": total_time,
        "avg_latency_seconds": total_time / len(results) if results else 0,
        "results": results,
    }

    os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)
    with open(args.output, "w") as f:
        json.dump(output, f, indent=2)

    print(f"\nDone! {total_time:.1f}s for {len(results)} questions")
    print(f"Results saved to {args.output}")

if __name__ == "__main__":
    main()