scripts/modal_run.py

#!/usr/bin/env python3
"""Modal-based reproduction: Qwen3-8B zero-shot on FHIR-AgentBench."""
import json
import csv
import random
import time
import os
import subprocess
from pathlib import Path
import modal

app = modal.App("repro-fhir-rl-tool-calling")

vllm_image = modal.Image.from_registry(
    "nvidia/cuda:12.4.1-devel-ubuntu22.04",
    setup_dockerfile_commands=[
        "RUN apt-get update && apt-get install -y python3 python3-pip git curl",
        "RUN pip3 install vllm==1.2.0 huggingface_hub datasets",
    ],
).pip_install("vllm==1.2.0", "huggingface_hub", "datasets")

@app.function(
    image=vllm_image,
    gpu="H100:1",
    timeout=3600,
    secrets=[modal.Secret.from_name("huggingface-token")],
    mounts=[modal.Mount.from_local_dir(
        "/tmp/FHIR-AgentBench/final_dataset",
        remote_path="/data"
    )],
)
def evaluate_qwen3_zeroshot(num_questions: int = 50):
    """Evaluate Qwen3-8B zero-shot on FHIR-AgentBench validation questions."""
    import torch
    from vllm import LLM, SamplingParams

    MODEL = "Qwen/Qwen3-8B"
    DATA_PATH = "/data/questions_answers_sql_fhir.csv"
    SPLIT = "valid"
    MAX_TURNS = 6
    TEMPERATURE = 0.1

    # 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")

    # Init vLLM
    print("Loading Qwen3-8B with vLLM...")
    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[:120],
            "true_answer": true_answer[:120] if true_answer else "",
            "generated": generated[:500],
            "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

    output = {
        "model": MODEL,
        "num_questions": num_questions,
        "split": SPLIT,
        "max_turns": MAX_TURNS,
        "temperature": TEMPERATURE,
        "total_time_seconds": total_time,
        "avg_latency_seconds": total_time / len(results) if results else 0,
        "results": results,
    }

    # Save output
    os.makedirs("/output", exist_ok=True)
    out_path = "/output/vllm_results_qwen3_zeroshot.json"
    with open(out_path, "w") as f:
        json.dump(output, f, indent=2)
    print(f"\nResults saved to {out_path}")
    print(f"Total: {total_time:.1f}s for {len(results)} questions")
    return output

@app.local_entrypoint()
def main(num_questions: int = 50):
    result = evaluate_qwen3_zeroshot.remote(num_questions=num_questions)
    print(f"Accuracy: {result['num_questions']} questions in {result['total_time_seconds']:.1f}s")