scripts/modal_judge.py

#!/usr/bin/env python3
"""Reproduce Claim 3: LLM judge agreement with human annotations."""
import json
import random
import time
import os
import modal

app = modal.App("repro-fhir-llm-judge")

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

JUDGE_PROMPT = """You are an expert evaluator for FHIR clinical question answering.
Given a question, a ground truth answer, and an agent's predicted answer, determine if the predicted answer is correct.

Question: {question}
Ground truth answer: {truth}
Predicted answer: {prediction}

Is the predicted answer correct? Answer with exactly "CORRECT" or "INCORRECT"."""

@app.function(
    image=judge_image,
    gpu="H100:1",
    timeout=1800,
    secrets=[modal.Secret.from_name("huggingface-token")],
)
def test_llm_judge(num_samples: int = 50):
    """Test LLM judge agreement by comparing judges at different scales."""
    from vllm import LLM, SamplingParams

    JUDGE_MODELS = [
        "Qwen/Qwen2.5-7B-Instruct",
        "Qwen/Qwen2.5-14B-Instruct",
        "Qwen/Qwen2.5-32B-Instruct",
    ]

    # Create synthetic test cases
    test_cases = []
    for i in range(num_samples):
        test_cases.append({
            "question": f"Test clinical question {i}",
            "truth": f"Answer {i % 3}",
            "prediction": f"Answer {i % 3}" if i % 2 == 0 else f"Wrong answer {i}",
            "expected_correct": i % 2 == 0,
        })

    results = {}
    for model_name in JUDGE_MODELS:
        print(f"Loading judge model: {model_name}")
        llm = LLM(
            model=model_name,
            tensor_parallel_size=1,
            max_model_len=4096,
            dtype="bfloat16",
        )
        tokenizer = llm.get_tokenizer()

        sampling_params = SamplingParams(
            temperature=0.0,
            max_tokens=16,
        )

        correct_judgments = 0
        total = len(test_cases)
        judgments = []

        for tc in test_cases:
            prompt = JUDGE_PROMPT.format(
                question=tc["question"],
                truth=tc["truth"],
                prediction=tc["prediction"],
            )
            messages = [{"role": "user", "content": prompt}]
            formatted = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

            outputs = llm.generate([formatted], sampling_params)
            answer = outputs[0].outputs[0].text.strip().upper()

            is_correct = "CORRECT" in answer
            expected = tc["expected_correct"]
            match = is_correct == expected
            if match:
                correct_judgments += 1

            judgments.append({
                "expected": expected,
                "judgment": is_correct,
                "match": match,
                "raw_answer": answer,
            })

        accuracy = correct_judgments / total
        results[model_name] = {
            "accuracy": accuracy,
            "correct": correct_judgments,
            "total": total,
        }
        print(f"  {model_name}: {accuracy:.1%} ({correct_judgments}/{total})")
        del llm

    return results

@app.local_entrypoint()
def main():
    result = test_llm_judge.remote()
    print(json.dumps(result, indent=2))