benchmark/genai_bench.py

from pathlib import Path
from typing import Iterator, List, Tuple

import requests


class GenAIBenchPrompts:
    def __init__(self):
        super().__init__()
        self._download_genai_bench_files()
        prompts_path = Path("downloads/genai_bench/prompts.txt")
        with open(prompts_path, "r") as f:
            self.prompts = [line.strip() for line in f if line.strip()]

    def __iter__(self) -> Iterator[Tuple[str, Path]]:
        for i, prompt in enumerate(self.prompts):
            yield prompt, Path(f"{i}.png")

    def _download_genai_bench_files(self) -> None:
        folder_name = Path("downloads/genai_bench")
        folder_name.mkdir(parents=True, exist_ok=True)
        prompts_url = "https://huggingface.co/datasets/zhiqiulin/GenAI-Bench-527/raw/main/prompts.txt"
        prompts_path = folder_name / "prompts.txt"
        if not prompts_path.exists():
            response = requests.get(prompts_url)
            with open(prompts_path, "w") as f:
                f.write(response.text)

    @property
    def name(self) -> str:
        return "genai_bench"

    @property
    def size(self) -> int:
        return len(self.prompts)

    @property
    def metrics(self) -> List[str]:
        return ["vqa"]