benchmark/data.py

"""Loading book corpus and evaluation queries."""

import json
from dataclasses import dataclass, field


@dataclass
class EvalQuery:
    query: str
    normalized_query: str
    relevance: dict  # book_id -> graded relevance (1-3)
    categories: list = field(default_factory=list)

    def text_for_mode(self, mode):
        if mode == "normalized":
            return self.normalized_query
        return self.query


def load_books(path):
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def load_queries(path):
    with open(path, "r", encoding="utf-8") as f:
        raw = json.load(f)
    return [
        EvalQuery(
            query=item["query"],
            normalized_query=item.get("normalized_query", item["query"]),
            relevance=item["relevance"],
            categories=item.get("categories", []),
        )
        for item in raw
    ]


def book_to_text(book):
    return (
        f"Title: {book['title']}\n"
        f"Author: {book['author']}\n"
        f"Genre: {', '.join(book['genres'])}\n"
        f"Tags: {', '.join(book['tags'])}\n"
        f"Description: {book['description']}\n"
    )