Frontier Judge QA Filter

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "data-designer>=0.5.6",
# ]
# ///
"""Long-Document Understanding Frontier Model QA Judge Recipe
Use a frontier VLM as an LLM-as-a-judge to evaluate the quality
of (question, answer) pairs generated by the upstream visual QA recipes. The
judge scores each example across five rubrics:
  1. **Answer Correctness** – factual accuracy against the visible document
  2. **Question Quality** – reasoning depth, ambiguity, specificity
  3. **Visual Grounding** – reliance on visual elements vs. plain text
  4. **Format Compliance** – answer format matches the question type
  5. **Training Signal Strength** – overall value as VLM training data
A weighted composite score (0–1) is computed from the five rubric scores.
Prerequisites:
    - A seed parquet file containing output from an upstream QA recipe
      (e.g. 05-visual-qa-sdg.py, 06-single-page-qa-sdg.py, or
      08-whole-document-qa-sdg.py) with at least:
        * `png_images_base64` – JSON array of base64-encoded PNG(s) of
          document pages.
        * `question_type`  – classification of the question.
        * `question`       – the generated question.
        * `answer`         – the generated answer.
    - Access to a frontier model endpoint that exposes an OpenAI-compatible
      API. Provide the model ID, endpoint URL, and the name of the
      environment variable holding the API key via the CLI flags
      ``--model-id``, ``--endpoint``, and ``--api-key-env``.
Run:
    # Basic usage (judges 5 records by default)
    uv run 09-frontier-judge-sdg.py --seed-path my_qa_output.parquet \
        --model-id <model-id> --endpoint <endpoint-url> --api-key-env <ENV_VAR>
    # Custom record count
    uv run 09-frontier-judge-sdg.py --seed-path my_qa_output.parquet \
        --model-id <model-id> --endpoint <endpoint-url> --api-key-env <ENV_VAR> \
        --num-records 100
    # For help message and available options
    uv run 09-frontier-judge-sdg.py --help
"""
from pathlib import Path
import data_designer.config as dd
from data_designer.interface import DataDesigner, DatasetCreationResults
PROVIDER_NAME = "frontier"
# =============================================================================
# Score weights for the weighted composite
# =============================================================================
FINAL_SCORE_WEIGHTS = {
    "Answer Correctness": 0.35,
    "Training Signal Strength": 0.30,
    "Question Quality": 0.15,
    "Visual Grounding": 0.10,
    "Format Compliance": 0.10,
}
# =============================================================================
# Custom column: weighted composite score
# =============================================================================
@dd.custom_column_generator(required_columns=["qa_quality_judge"])
def compute_weighted_score(row: dict) -> dict:
    """Weighted composite score normalized to 0-1."""
    judge = row["qa_quality_judge"]
    raw = sum(float(judge[k]["score"]) * w for k, w in FINAL_SCORE_WEIGHTS.items())
    row["weighted_quality_score"] = round(raw / 5.0, 2)
    return row
# =============================================================================
# Judge prompt
# =============================================================================
PROMPT_JUDGE = """\
You are an expert evaluator of visual document question-answering (VQA) training data
for the MMLongBench-Doc benchmark.
Your task is to assess the quality of a (question, answer) pair generated from a PDF
document image. The goal is to determine how strong of a training signal this example
would provide for improving VLM performance.
You will be given:
- One or more images of document pages (with tables, charts, diagrams, etc.)
- A question type classification
- A question about the document
- An answer to the question
<question-type>
{{ question_type }}
</question-type>
<question>
{{ question }}
</question>
<answer>
{{ answer }}
</answer>
Evaluate the example across the following rubrics. For each rubric, provide a brief
reasoning and a score. Be objective and critical -- do not inflate scores.
{
  "Answer Correctness": {
    "reasoning": "Your brief analysis here",
    "score": "X"
  },
  "Question Quality": {
    "reasoning": "Your brief analysis here",
    "score": "X"
  },
  "Visual Grounding": {
    "reasoning": "Your brief analysis here",
    "score": "X"
  },
  "Format Compliance": {
    "reasoning": "Your brief analysis here",
    "score": "X"
  },
  "Training Signal Strength": {
    "reasoning": "Your brief analysis here",
    "score": "X"
  }
}
Provide your evaluation in the exact JSON format above with ALL 5 rubrics.
Keep your reasoning for each rubric short and to the point.
"""
# =============================================================================
# Score rubric definitions
# =============================================================================
answer_correctness_score = dd.Score(
    name="Answer Correctness",
    description=(
        "Is the answer factually correct given the visible document content? "
        "Verify by examining the image yourself. For calculations, redo the math. "
        "For counts, recount. For lists, check completeness."
    ),
    options={
        "5": "Exactly correct: answer matches the visible content precisely, calculations are accurate, lists are complete",
        "4": "Substantially correct: answer is right with minor imprecision (e.g., rounding differences within +/-5%, equivalent formats like '25%' vs '0.25')",
        "3": "Partially correct: core answer is right but has notable issues (missing list items, slightly off calculation, incomplete but not wrong)",
        "2": "Mostly incorrect: answer has the right idea but wrong values, wrong entity, or significant calculation errors",
        "1": "Incorrect: answer contradicts the visible content, uses wrong data, or is completely off",
        "0": "Not answerable or refused: answer is a refusal, 'Not answerable', or nonsensical when a real answer exists",
    },
)
question_quality_score = dd.Score(
    name="Question Quality",
    description=(
        "Is the question well-formed, unambiguous, and appropriately challenging? "
        "Does it require genuine reasoning (comparison, calculation, counting) rather than trivial lookup? "
        "Is it specific to the visual content and not generic?"
    ),
    options={
        "5": "Excellent: requires clear reasoning (comparison, calculation, or cross-element synthesis), unambiguous, has exactly one correct answer, well-matched to the visual element type",
        "4": "Good: requires some reasoning, mostly unambiguous, well-grounded in the visual content with minor issues",
        "3": "Adequate: reasonable question but either too easy (direct lookup), slightly ambiguous, or not well-matched to the visual element type",
        "2": "Poor: trivial lookup, ambiguous wording, or asks about content not well-suited to the visual element type",
        "1": "Very poor: unanswerable from the image, contains the answer, or is about irrelevant content",
        "0": "Invalid: nonsensical, empty, or completely unrelated to the document",
    },
)
visual_grounding_score = dd.Score(
    name="Visual Grounding",
    description=(
        "Does the question target the actual visual elements (tables, charts, diagrams) in the image? "
        "Does answering require examining the visual structure, not just reading plain text? "
        "Is the question grounded in specific, identifiable elements?"
    ),
    options={
        "5": "Excellent: question directly targets specific visual elements (chart data, table cells, diagram nodes), answering requires visual perception and spatial understanding",
        "4": "Good: question is grounded in visual content with clear references to identifiable elements, requires examining the visual structure",
        "3": "Adequate: question relates to visual content but could partially be answered from text alone, or uses vague references ('the table' without specificity)",
        "2": "Poor: question mostly targets plain text content, minimal visual grounding, could be answered without seeing the visual elements",
        "1": "Very poor: question has no meaningful connection to the visual elements, purely text-based",
        "0": "No grounding: question is about content not present in the image at all",
    },
)
format_compliance_score = dd.Score(
    name="Format Compliance",
    description=(
        "Does the answer match the expected format for its question type? "
        "Check: multiple choice uses 'A. option' format, yes/no is exactly 'Yes'/'No', "
        "percentages include '%', integers are digits only, lists are JSON arrays, "
        "and the answer contains no reasoning traces or meta-commentary."
    ),
    options={
        "5": "Perfect compliance: answer format exactly matches the question type requirements, no extraneous content",
        "4": "Good compliance: correct format with trivial deviations (e.g., extra whitespace, minor punctuation)",
        "3": "Adequate: answer is usable but has format issues (e.g., missing units, prose instead of JSON array, includes 'Based on the image...')",
        "2": "Poor: significant format violations (e.g., includes reasoning steps, wrong answer structure, contains <think> tags)",
        "1": "Very poor: answer format is fundamentally wrong for the question type",
        "0": "No compliance: answer is empty, garbled, or completely ignores format requirements",
    },
)
training_signal_score = dd.Score(
    name="Training Signal Strength",
    description=(
        "Overall, how valuable is this (question, answer) pair as training data for improving "
        "VLM performance on document understanding? Consider: does it exercise visual perception, "
        "require non-trivial reasoning, demand multi-page evidence gathering, and provide a clear learning signal?"
    ),
    options={
        "5": "Excellent: requires combining evidence from multiple pages, exercises visual perception + reasoning, non-trivial, clear correct answer. Would meaningfully improve a VLM on document QA benchmarks",
        "4": "Strong: good training example with cross-page reasoning or strong single-page visual grounding and reasoning, minor issues don't significantly reduce value",
        "3": "Moderate: decent training signal but answerable from a single page, or doesn't fully exercise multi-page or visual understanding",
        "2": "Weak: limited training value -- trivial question, wrong answer, single-page lookup, or doesn't require visual reasoning",
        "1": "Very weak: almost no training value -- incorrect, ambiguous, or completely text-based with no multi-page dependency",
        "0": "No value: harmful to training -- wrong answer presented as correct, nonsensical, or would teach bad patterns",
    },
)
# =============================================================================
# Config builder
# =============================================================================
def build_config(
    seed_path: str = "seed.parquet",
    model_alias: str = "frontier-judge-vlm",
    model_id: str = "",
) -> dd.DataDesignerConfigBuilder:
    """Build the Data Designer config for frontier-model QA judging."""
    config_builder = dd.DataDesignerConfigBuilder(
        model_configs=[
            dd.ModelConfig(
                alias=model_alias,
                model=model_id,
                provider=PROVIDER_NAME,
                inference_parameters=dd.ChatCompletionInferenceParams(
                    timeout=300,
                    max_tokens=40000,
                    max_parallel_requests=32,
                ),
            ),
        ]
    )
    config_builder.with_seed_dataset(
        dd.LocalFileSeedSource(path=seed_path),
        sampling_strategy=dd.SamplingStrategy.ORDERED,
    )
    config_builder.add_column(
        dd.LLMJudgeColumnConfig(
            name="qa_quality_judge",
            model_alias=model_alias,
            prompt=PROMPT_JUDGE,
            scores=[
                answer_correctness_score,
                question_quality_score,
                visual_grounding_score,
                format_compliance_score,
                training_signal_score,
            ],
            multi_modal_context=[
                dd.ImageContext(
                    column_name="png_images_base64",
                    data_type=dd.ModalityDataType.BASE64,
                    image_format=dd.ImageFormat.PNG,
                ),
            ],
        )
    )
    config_builder.add_column(
        dd.CustomColumnConfig(
            name="weighted_quality_score",
            generator_function=compute_weighted_score,
        )
    )
    return config_builder
# =============================================================================
# Dataset creation
# =============================================================================
def create_dataset(
    config_builder: dd.DataDesignerConfigBuilder,
    num_records: int,
    endpoint: str = "",
    api_key_env: str = "",
    artifact_path: Path | str | None = None,
) -> DatasetCreationResults:
    """Create the judged dataset."""
    model_providers = [
        dd.ModelProvider(
            name=PROVIDER_NAME,
            endpoint=endpoint,
            provider_type="openai",
            api_key=api_key_env,
        ),
    ]
    data_designer = DataDesigner(
        artifact_path=artifact_path,
        model_providers=model_providers,
    )
    data_designer.set_run_config(
        dd.RunConfig(disable_early_shutdown=True, progress_bar=True),
    )
    results = data_designer.create(config_builder, num_records=num_records, dataset_name="frontier_judge")
    return results
# =============================================================================
# CLI entry point
# =============================================================================
if __name__ == "__main__":
    from argparse import ArgumentParser
    parser = ArgumentParser()
    parser.add_argument("--seed-path", type=str, required=True, help="Path to the seed parquet file")
    parser.add_argument("--model-alias", type=str, default="frontier-judge-vlm")
    parser.add_argument("--model-id", type=str, required=True, help="ID of the model to use for judging")
    parser.add_argument("--endpoint", type=str, required=True, help="OpenAI-compatible API endpoint URL")
    parser.add_argument(
        "--api-key-env", type=str, required=True, help="Environment variable name containing the API key"
    )
    parser.add_argument("--num-records", type=int, default=5)
    parser.add_argument("--artifact-path", type=str, default=None)
    args = parser.parse_args()
    config_builder = build_config(
        seed_path=args.seed_path,
        model_alias=args.model_alias,
        model_id=args.model_id,
    )
    results = create_dataset(
        config_builder,
        num_records=args.num_records,
        endpoint=args.endpoint,
        api_key_env=args.api_key_env,
        artifact_path=args.artifact_path,
    )
    print(f"Dataset saved to: {results.artifact_storage.final_dataset_path}")
    results.load_analysis().to_report()

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "data-designer>=0.5.6",
# ]
# ///
"""Long-Document Understanding Frontier Model QA Judge Recipe
Use a frontier VLM as an LLM-as-a-judge to evaluate the quality
of (question, answer) pairs generated by the upstream visual QA recipes. The
judge scores each example across five rubrics:
  1. **Answer Correctness** – factual accuracy against the visible document
  2. **Question Quality** – reasoning depth, ambiguity, specificity
  3. **Visual Grounding** – reliance on visual elements vs. plain text
  4. **Format Compliance** – answer format matches the question type
  5. **Training Signal Strength** – overall value as VLM training data
A weighted composite score (0–1) is computed from the five rubric scores.
Prerequisites:
    - A seed parquet file containing output from an upstream QA recipe
      (e.g. 05-visual-qa-sdg.py, 06-single-page-qa-sdg.py, or
      08-whole-document-qa-sdg.py) with at least:
        * `png_images_base64` – JSON array of base64-encoded PNG(s) of
          document pages.
        * `question_type`  – classification of the question.
        * `question`       – the generated question.
        * `answer`         – the generated answer.
    - Access to a frontier model endpoint that exposes an OpenAI-compatible
      API. Provide the model ID, endpoint URL, and the name of the
      environment variable holding the API key via the CLI flags
      ``--model-id``, ``--endpoint``, and ``--api-key-env``.
Run:
    # Basic usage (judges 5 records by default)
    uv run 09-frontier-judge-sdg.py --seed-path my_qa_output.parquet \
        --model-id <model-id> --endpoint <endpoint-url> --api-key-env <ENV_VAR>
    # Custom record count
    uv run 09-frontier-judge-sdg.py --seed-path my_qa_output.parquet \
        --model-id <model-id> --endpoint <endpoint-url> --api-key-env <ENV_VAR> \
        --num-records 100
    # For help message and available options
    uv run 09-frontier-judge-sdg.py --help
"""
from pathlib import Path
import data_designer.config as dd
from data_designer.interface import DataDesigner, DatasetCreationResults
PROVIDER_NAME = "frontier"
# =============================================================================
# Score weights for the weighted composite
# =============================================================================
FINAL_SCORE_WEIGHTS = {
    "Answer Correctness": 0.35,
    "Training Signal Strength": 0.30,
    "Question Quality": 0.15,
    "Visual Grounding": 0.10,
    "Format Compliance": 0.10,
}
# =============================================================================
# Custom column: weighted composite score
# =============================================================================
@dd.custom_column_generator(required_columns=["qa_quality_judge"])
def compute_weighted_score(row: dict) -> dict:
    """Weighted composite score normalized to 0-1."""
    judge = row["qa_quality_judge"]
    raw = sum(float(judge[k]["score"]) * w for k, w in FINAL_SCORE_WEIGHTS.items())
    row["weighted_quality_score"] = round(raw / 5.0, 2)
    return row
# =============================================================================
# Judge prompt
# =============================================================================
PROMPT_JUDGE = """\
You are an expert evaluator of visual document question-answering (VQA) training data
for the MMLongBench-Doc benchmark.
Your task is to assess the quality of a (question, answer) pair generated from a PDF
document image. The goal is to determine how strong of a training signal this example
would provide for improving VLM performance.
You will be given:
- One or more images of document pages (with tables, charts, diagrams, etc.)
- A question type classification
- A question about the document
- An answer to the question
<question-type>
{{ question_type }}
</question-type>
<question>
{{ question }}
</question>
<answer>
{{ answer }}
</answer>
Evaluate the example across the following rubrics. For each rubric, provide a brief
reasoning and a score. Be objective and critical -- do not inflate scores.
{
  "Answer Correctness": {
    "reasoning": "Your brief analysis here",
    "score": "X"
  },
  "Question Quality": {
    "reasoning": "Your brief analysis here",
    "score": "X"
  },
  "Visual Grounding": {
    "reasoning": "Your brief analysis here",
    "score": "X"
  },
  "Format Compliance": {
    "reasoning": "Your brief analysis here",
    "score": "X"
  },
  "Training Signal Strength": {
    "reasoning": "Your brief analysis here",
    "score": "X"
  }
}
Provide your evaluation in the exact JSON format above with ALL 5 rubrics.
Keep your reasoning for each rubric short and to the point.
"""
# =============================================================================
# Score rubric definitions
# =============================================================================
answer_correctness_score = dd.Score(
    name="Answer Correctness",
    description=(
        "Is the answer factually correct given the visible document content? "
        "Verify by examining the image yourself. For calculations, redo the math. "
        "For counts, recount. For lists, check completeness."
    ),
    options={
        "5": "Exactly correct: answer matches the visible content precisely, calculations are accurate, lists are complete",
        "4": "Substantially correct: answer is right with minor imprecision (e.g., rounding differences within +/-5%, equivalent formats like '25%' vs '0.25')",
        "3": "Partially correct: core answer is right but has notable issues (missing list items, slightly off calculation, incomplete but not wrong)",
        "2": "Mostly incorrect: answer has the right idea but wrong values, wrong entity, or significant calculation errors",
        "1": "Incorrect: answer contradicts the visible content, uses wrong data, or is completely off",
        "0": "Not answerable or refused: answer is a refusal, 'Not answerable', or nonsensical when a real answer exists",
    },
)
question_quality_score = dd.Score(
    name="Question Quality",
    description=(
        "Is the question well-formed, unambiguous, and appropriately challenging? "
        "Does it require genuine reasoning (comparison, calculation, counting) rather than trivial lookup? "
        "Is it specific to the visual content and not generic?"
    ),
    options={
        "5": "Excellent: requires clear reasoning (comparison, calculation, or cross-element synthesis), unambiguous, has exactly one correct answer, well-matched to the visual element type",
        "4": "Good: requires some reasoning, mostly unambiguous, well-grounded in the visual content with minor issues",
        "3": "Adequate: reasonable question but either too easy (direct lookup), slightly ambiguous, or not well-matched to the visual element type",
        "2": "Poor: trivial lookup, ambiguous wording, or asks about content not well-suited to the visual element type",
        "1": "Very poor: unanswerable from the image, contains the answer, or is about irrelevant content",
        "0": "Invalid: nonsensical, empty, or completely unrelated to the document",
    },
)
visual_grounding_score = dd.Score(
    name="Visual Grounding",
    description=(
        "Does the question target the actual visual elements (tables, charts, diagrams) in the image? "
        "Does answering require examining the visual structure, not just reading plain text? "
        "Is the question grounded in specific, identifiable elements?"
    ),
    options={
        "5": "Excellent: question directly targets specific visual elements (chart data, table cells, diagram nodes), answering requires visual perception and spatial understanding",
        "4": "Good: question is grounded in visual content with clear references to identifiable elements, requires examining the visual structure",
        "3": "Adequate: question relates to visual content but could partially be answered from text alone, or uses vague references ('the table' without specificity)",
        "2": "Poor: question mostly targets plain text content, minimal visual grounding, could be answered without seeing the visual elements",
        "1": "Very poor: question has no meaningful connection to the visual elements, purely text-based",
        "0": "No grounding: question is about content not present in the image at all",
    },
)
format_compliance_score = dd.Score(
    name="Format Compliance",
    description=(
        "Does the answer match the expected format for its question type? "
        "Check: multiple choice uses 'A. option' format, yes/no is exactly 'Yes'/'No', "
        "percentages include '%', integers are digits only, lists are JSON arrays, "
        "and the answer contains no reasoning traces or meta-commentary."
    ),
    options={
        "5": "Perfect compliance: answer format exactly matches the question type requirements, no extraneous content",
        "4": "Good compliance: correct format with trivial deviations (e.g., extra whitespace, minor punctuation)",
        "3": "Adequate: answer is usable but has format issues (e.g., missing units, prose instead of JSON array, includes 'Based on the image...')",
        "2": "Poor: significant format violations (e.g., includes reasoning steps, wrong answer structure, contains <think> tags)",
        "1": "Very poor: answer format is fundamentally wrong for the question type",
        "0": "No compliance: answer is empty, garbled, or completely ignores format requirements",
    },
)
training_signal_score = dd.Score(
    name="Training Signal Strength",
    description=(
        "Overall, how valuable is this (question, answer) pair as training data for improving "
        "VLM performance on document understanding? Consider: does it exercise visual perception, "
        "require non-trivial reasoning, demand multi-page evidence gathering, and provide a clear learning signal?"
    ),
    options={
        "5": "Excellent: requires combining evidence from multiple pages, exercises visual perception + reasoning, non-trivial, clear correct answer. Would meaningfully improve a VLM on document QA benchmarks",
        "4": "Strong: good training example with cross-page reasoning or strong single-page visual grounding and reasoning, minor issues don't significantly reduce value",
        "3": "Moderate: decent training signal but answerable from a single page, or doesn't fully exercise multi-page or visual understanding",
        "2": "Weak: limited training value -- trivial question, wrong answer, single-page lookup, or doesn't require visual reasoning",
        "1": "Very weak: almost no training value -- incorrect, ambiguous, or completely text-based with no multi-page dependency",
        "0": "No value: harmful to training -- wrong answer presented as correct, nonsensical, or would teach bad patterns",
    },
)
# =============================================================================
# Config builder
# =============================================================================
def build_config(
    seed_path: str = "seed.parquet",
    model_alias: str = "frontier-judge-vlm",
    model_id: str = "",
) -> dd.DataDesignerConfigBuilder:
    """Build the Data Designer config for frontier-model QA judging."""
    config_builder = dd.DataDesignerConfigBuilder(
        model_configs=[
            dd.ModelConfig(
                alias=model_alias,
                model=model_id,
                provider=PROVIDER_NAME,
                inference_parameters=dd.ChatCompletionInferenceParams(
                    timeout=300,
                    max_tokens=40000,
                    max_parallel_requests=32,
                ),
            ),
        ]
    )
    config_builder.with_seed_dataset(
        dd.LocalFileSeedSource(path=seed_path),
        sampling_strategy=dd.SamplingStrategy.ORDERED,
    )
    config_builder.add_column(
        dd.LLMJudgeColumnConfig(
            name="qa_quality_judge",
            model_alias=model_alias,
            prompt=PROMPT_JUDGE,
            scores=[
                answer_correctness_score,
                question_quality_score,
                visual_grounding_score,
                format_compliance_score,
                training_signal_score,
            ],
            multi_modal_context=[
                dd.ImageContext(
                    column_name="png_images_base64",
                    data_type=dd.ModalityDataType.BASE64,
                    image_format=dd.ImageFormat.PNG,
                ),
            ],
        )
    )
    config_builder.add_column(
        dd.CustomColumnConfig(
            name="weighted_quality_score",
            generator_function=compute_weighted_score,
        )
    )
    return config_builder
# =============================================================================
# Dataset creation
# =============================================================================
def create_dataset(
    config_builder: dd.DataDesignerConfigBuilder,
    num_records: int,
    endpoint: str = "",
    api_key_env: str = "",
    artifact_path: Path | str | None = None,
) -> DatasetCreationResults:
    """Create the judged dataset."""
    model_providers = [
        dd.ModelProvider(
            name=PROVIDER_NAME,
            endpoint=endpoint,
            provider_type="openai",
            api_key=api_key_env,
        ),
    ]
    data_designer = DataDesigner(
        artifact_path=artifact_path,
        model_providers=model_providers,
    )
    data_designer.set_run_config(
        dd.RunConfig(disable_early_shutdown=True, progress_bar=True),
    )
    results = data_designer.create(config_builder, num_records=num_records, dataset_name="frontier_judge")
    return results
# =============================================================================
# CLI entry point
# =============================================================================
if __name__ == "__main__":
    from argparse import ArgumentParser
    parser = ArgumentParser()
    parser.add_argument("--seed-path", type=str, required=True, help="Path to the seed parquet file")
    parser.add_argument("--model-alias", type=str, default="frontier-judge-vlm")
    parser.add_argument("--model-id", type=str, required=True, help="ID of the model to use for judging")
    parser.add_argument("--endpoint", type=str, required=True, help="OpenAI-compatible API endpoint URL")
    parser.add_argument(
        "--api-key-env", type=str, required=True, help="Environment variable name containing the API key"
    )
    parser.add_argument("--num-records", type=int, default=5)
    parser.add_argument("--artifact-path", type=str, default=None)
    args = parser.parse_args()
    config_builder = build_config(
        seed_path=args.seed_path,
        model_alias=args.model_alias,
        model_id=args.model_id,
    )
    results = create_dataset(
        config_builder,
        num_records=args.num_records,
        endpoint=args.endpoint,
        api_key_env=args.api_key_env,
        artifact_path=args.artifact_path,
    )
    print(f"Dataset saved to: {results.artifact_storage.final_dataset_path}")
    results.load_analysis().to_report()

Frontier Judge QA Filter

Frontier Judge QA Filter

Download Recipe

Download Recipe

1	# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2	# SPDX-License-Identifier: Apache-2.0
3	# /// script
4	# requires-python = ">=3.10"
5	# dependencies = [
6	# "data-designer>=0.5.6",
7	# ]
8	# ///
9	"""Long-Document Understanding Frontier Model QA Judge Recipe
10
11	Use a frontier VLM as an LLM-as-a-judge to evaluate the quality
12	of (question, answer) pairs generated by the upstream visual QA recipes. The
13	judge scores each example across five rubrics:
14
15	1. Answer Correctness – factual accuracy against the visible document
16	2. Question Quality – reasoning depth, ambiguity, specificity
17	3. Visual Grounding – reliance on visual elements vs. plain text
18	4. Format Compliance – answer format matches the question type
19	5. Training Signal Strength – overall value as VLM training data
20
21	A weighted composite score (0–1) is computed from the five rubric scores.
22
23	Prerequisites:
24	- A seed parquet file containing output from an upstream QA recipe
25	(e.g. 05-visual-qa-sdg.py, 06-single-page-qa-sdg.py, or
26	08-whole-document-qa-sdg.py) with at least:
27	* `png_images_base64` – JSON array of base64-encoded PNG(s) of
28	document pages.
29	* `question_type` – classification of the question.
30	* `question` – the generated question.
31	* `answer` – the generated answer.
32	- Access to a frontier model endpoint that exposes an OpenAI-compatible
33	API. Provide the model ID, endpoint URL, and the name of the
34	environment variable holding the API key via the CLI flags
35	``--model-id``, ``--endpoint``, and ``--api-key-env``.
36
37	Run:
38	# Basic usage (judges 5 records by default)
39	uv run 09-frontier-judge-sdg.py --seed-path my_qa_output.parquet \
40	--model-id <model-id> --endpoint <endpoint-url> --api-key-env <ENV_VAR>
41
42	# Custom record count
43	uv run 09-frontier-judge-sdg.py --seed-path my_qa_output.parquet \
44	--model-id <model-id> --endpoint <endpoint-url> --api-key-env <ENV_VAR> \
45	--num-records 100
46
47	# For help message and available options
48	uv run 09-frontier-judge-sdg.py --help
49	"""
50
51	from pathlib import Path
52
53	import data_designer.config as dd
54	from data_designer.interface import DataDesigner, DatasetCreationResults
55
56	PROVIDER_NAME = "frontier"
57
58	# =============================================================================
59	# Score weights for the weighted composite
60	# =============================================================================
61
62	FINAL_SCORE_WEIGHTS = {
63	"Answer Correctness": 0.35,
64	"Training Signal Strength": 0.30,
65	"Question Quality": 0.15,
66	"Visual Grounding": 0.10,
67	"Format Compliance": 0.10,
68	}
69
70	# =============================================================================
71	# Custom column: weighted composite score
72	# =============================================================================
73
74
75	@dd.custom_column_generator(required_columns=["qa_quality_judge"])
76	def compute_weighted_score(row: dict) -> dict:
77	"""Weighted composite score normalized to 0-1."""
78	judge = row["qa_quality_judge"]
79	raw = sum(float(judge[k]["score"]) * w for k, w in FINAL_SCORE_WEIGHTS.items())
80	row["weighted_quality_score"] = round(raw / 5.0, 2)
81	return row
82
83
84	# =============================================================================
85	# Judge prompt
86	# =============================================================================
87
88	PROMPT_JUDGE = """\
89	You are an expert evaluator of visual document question-answering (VQA) training data
90	for the MMLongBench-Doc benchmark.
91
92	Your task is to assess the quality of a (question, answer) pair generated from a PDF
93	document image. The goal is to determine how strong of a training signal this example
94	would provide for improving VLM performance.
95
96	You will be given:
97	- One or more images of document pages (with tables, charts, diagrams, etc.)
98	- A question type classification
99	- A question about the document
100	- An answer to the question
101
102	<question-type>
103	{{ question_type }}
104	</question-type>
105
106	<question>
107	{{ question }}
108	</question>
109
110	<answer>
111	{{ answer }}
112	</answer>
113
114	Evaluate the example across the following rubrics. For each rubric, provide a brief
115	reasoning and a score. Be objective and critical -- do not inflate scores.
116
117	{
118	"Answer Correctness": {
119	"reasoning": "Your brief analysis here",
120	"score": "X"
121	},
122	"Question Quality": {
123	"reasoning": "Your brief analysis here",
124	"score": "X"
125	},
126	"Visual Grounding": {
127	"reasoning": "Your brief analysis here",
128	"score": "X"
129	},
130	"Format Compliance": {
131	"reasoning": "Your brief analysis here",
132	"score": "X"
133	},
134	"Training Signal Strength": {
135	"reasoning": "Your brief analysis here",
136	"score": "X"
137	}
138	}
139
140	Provide your evaluation in the exact JSON format above with ALL 5 rubrics.
141	Keep your reasoning for each rubric short and to the point.
142	"""
143
144	# =============================================================================
145	# Score rubric definitions
146	# =============================================================================
147
148	answer_correctness_score = dd.Score(
149	name="Answer Correctness",
150	description=(
151	"Is the answer factually correct given the visible document content? "
152	"Verify by examining the image yourself. For calculations, redo the math. "
153	"For counts, recount. For lists, check completeness."
154	),
155	options={
156	"5": "Exactly correct: answer matches the visible content precisely, calculations are accurate, lists are complete",
157	"4": "Substantially correct: answer is right with minor imprecision (e.g., rounding differences within +/-5%, equivalent formats like '25%' vs '0.25')",
158	"3": "Partially correct: core answer is right but has notable issues (missing list items, slightly off calculation, incomplete but not wrong)",
159	"2": "Mostly incorrect: answer has the right idea but wrong values, wrong entity, or significant calculation errors",
160	"1": "Incorrect: answer contradicts the visible content, uses wrong data, or is completely off",
161	"0": "Not answerable or refused: answer is a refusal, 'Not answerable', or nonsensical when a real answer exists",
162	},
163	)
164
165	question_quality_score = dd.Score(
166	name="Question Quality",
167	description=(
168	"Is the question well-formed, unambiguous, and appropriately challenging? "
169	"Does it require genuine reasoning (comparison, calculation, counting) rather than trivial lookup? "
170	"Is it specific to the visual content and not generic?"
171	),
172	options={
173	"5": "Excellent: requires clear reasoning (comparison, calculation, or cross-element synthesis), unambiguous, has exactly one correct answer, well-matched to the visual element type",
174	"4": "Good: requires some reasoning, mostly unambiguous, well-grounded in the visual content with minor issues",
175	"3": "Adequate: reasonable question but either too easy (direct lookup), slightly ambiguous, or not well-matched to the visual element type",
176	"2": "Poor: trivial lookup, ambiguous wording, or asks about content not well-suited to the visual element type",
177	"1": "Very poor: unanswerable from the image, contains the answer, or is about irrelevant content",
178	"0": "Invalid: nonsensical, empty, or completely unrelated to the document",
179	},
180	)
181
182	visual_grounding_score = dd.Score(
183	name="Visual Grounding",
184	description=(
185	"Does the question target the actual visual elements (tables, charts, diagrams) in the image? "
186	"Does answering require examining the visual structure, not just reading plain text? "
187	"Is the question grounded in specific, identifiable elements?"
188	),
189	options={
190	"5": "Excellent: question directly targets specific visual elements (chart data, table cells, diagram nodes), answering requires visual perception and spatial understanding",
191	"4": "Good: question is grounded in visual content with clear references to identifiable elements, requires examining the visual structure",
192	"3": "Adequate: question relates to visual content but could partially be answered from text alone, or uses vague references ('the table' without specificity)",
193	"2": "Poor: question mostly targets plain text content, minimal visual grounding, could be answered without seeing the visual elements",
194	"1": "Very poor: question has no meaningful connection to the visual elements, purely text-based",
195	"0": "No grounding: question is about content not present in the image at all",
196	},
197	)
198
199	format_compliance_score = dd.Score(
200	name="Format Compliance",
201	description=(
202	"Does the answer match the expected format for its question type? "
203	"Check: multiple choice uses 'A. option' format, yes/no is exactly 'Yes'/'No', "
204	"percentages include '%', integers are digits only, lists are JSON arrays, "
205	"and the answer contains no reasoning traces or meta-commentary."
206	),
207	options={
208	"5": "Perfect compliance: answer format exactly matches the question type requirements, no extraneous content",
209	"4": "Good compliance: correct format with trivial deviations (e.g., extra whitespace, minor punctuation)",
210	"3": "Adequate: answer is usable but has format issues (e.g., missing units, prose instead of JSON array, includes 'Based on the image...')",
211	"2": "Poor: significant format violations (e.g., includes reasoning steps, wrong answer structure, contains <think> tags)",
212	"1": "Very poor: answer format is fundamentally wrong for the question type",
213	"0": "No compliance: answer is empty, garbled, or completely ignores format requirements",
214	},
215	)
216
217	training_signal_score = dd.Score(
218	name="Training Signal Strength",
219	description=(
220	"Overall, how valuable is this (question, answer) pair as training data for improving "
221	"VLM performance on document understanding? Consider: does it exercise visual perception, "
222	"require non-trivial reasoning, demand multi-page evidence gathering, and provide a clear learning signal?"
223	),
224	options={
225	"5": "Excellent: requires combining evidence from multiple pages, exercises visual perception + reasoning, non-trivial, clear correct answer. Would meaningfully improve a VLM on document QA benchmarks",
226	"4": "Strong: good training example with cross-page reasoning or strong single-page visual grounding and reasoning, minor issues don't significantly reduce value",
227	"3": "Moderate: decent training signal but answerable from a single page, or doesn't fully exercise multi-page or visual understanding",
228	"2": "Weak: limited training value -- trivial question, wrong answer, single-page lookup, or doesn't require visual reasoning",
229	"1": "Very weak: almost no training value -- incorrect, ambiguous, or completely text-based with no multi-page dependency",
230	"0": "No value: harmful to training -- wrong answer presented as correct, nonsensical, or would teach bad patterns",
231	},
232	)
233
234
235	# =============================================================================
236	# Config builder
237	# =============================================================================
238
239
240	def build_config(
241	seed_path: str = "seed.parquet",
242	model_alias: str = "frontier-judge-vlm",
243	model_id: str = "",
244	) -> dd.DataDesignerConfigBuilder:
245	"""Build the Data Designer config for frontier-model QA judging."""
246	config_builder = dd.DataDesignerConfigBuilder(
247	model_configs=[
248	dd.ModelConfig(
249	alias=model_alias,
250	model=model_id,
251	provider=PROVIDER_NAME,
252	inference_parameters=dd.ChatCompletionInferenceParams(
253	timeout=300,
254	max_tokens=40000,
255	max_parallel_requests=32,
256	),
257	),
258	]
259	)
260
261	config_builder.with_seed_dataset(
262	dd.LocalFileSeedSource(path=seed_path),
263	sampling_strategy=dd.SamplingStrategy.ORDERED,
264	)
265
266	config_builder.add_column(
267	dd.LLMJudgeColumnConfig(
268	name="qa_quality_judge",
269	model_alias=model_alias,
270	prompt=PROMPT_JUDGE,
271	scores=[
272	answer_correctness_score,
273	question_quality_score,
274	visual_grounding_score,
275	format_compliance_score,
276	training_signal_score,
277	],
278	multi_modal_context=[
279	dd.ImageContext(
280	column_name="png_images_base64",
281	data_type=dd.ModalityDataType.BASE64,
282	image_format=dd.ImageFormat.PNG,
283	),
284	],
285	)
286	)
287
288	config_builder.add_column(
289	dd.CustomColumnConfig(
290	name="weighted_quality_score",
291	generator_function=compute_weighted_score,
292	)
293	)
294
295	return config_builder
296
297
298	# =============================================================================
299	# Dataset creation
300	# =============================================================================
301
302
303	def create_dataset(
304	config_builder: dd.DataDesignerConfigBuilder,
305	num_records: int,
306	endpoint: str = "",
307	api_key_env: str = "",
308	artifact_path: Path \| str \| None = None,
309	) -> DatasetCreationResults:
310	"""Create the judged dataset."""
311	model_providers = [
312	dd.ModelProvider(
313	name=PROVIDER_NAME,
314	endpoint=endpoint,
315	provider_type="openai",
316	api_key=api_key_env,
317	),
318	]
319	data_designer = DataDesigner(
320	artifact_path=artifact_path,
321	model_providers=model_providers,
322	)
323	data_designer.set_run_config(
324	dd.RunConfig(disable_early_shutdown=True, progress_bar=True),
325	)
326	results = data_designer.create(config_builder, num_records=num_records, dataset_name="frontier_judge")
327	return results
328
329
330	# =============================================================================
331	# CLI entry point
332	# =============================================================================
333
334	if __name__ == "__main__":
335	from argparse import ArgumentParser
336
337	parser = ArgumentParser()
338	parser.add_argument("--seed-path", type=str, required=True, help="Path to the seed parquet file")
339	parser.add_argument("--model-alias", type=str, default="frontier-judge-vlm")
340	parser.add_argument("--model-id", type=str, required=True, help="ID of the model to use for judging")
341	parser.add_argument("--endpoint", type=str, required=True, help="OpenAI-compatible API endpoint URL")
342	parser.add_argument(
343	"--api-key-env", type=str, required=True, help="Environment variable name containing the API key"
344	)
345	parser.add_argument("--num-records", type=int, default=5)
346	parser.add_argument("--artifact-path", type=str, default=None)
347	args = parser.parse_args()
348
349	config_builder = build_config(
350	seed_path=args.seed_path,
351	model_alias=args.model_alias,
352	model_id=args.model_id,
353	)
354	results = create_dataset(
355	config_builder,
356	num_records=args.num_records,
357	endpoint=args.endpoint,
358	api_key_env=args.api_key_env,
359	artifact_path=args.artifact_path,
360	)
361
362	print(f"Dataset saved to: {results.artifact_storage.final_dataset_path}")
363
364	results.load_analysis().to_report()