utils/sample_creator.py

"""Create Inspect AI samples from paper configurations."""

from pathlib import Path

from inspect_ai.dataset import Sample

from logger import get_logger
from schemas import CompressionSample
from utils.config_loader import load_papers_config
from utils.constants import DEFAULT_COMPRESSION_LEVELS
from utils.pdf_utils import list_papers

log = get_logger("utils.sample_creator")


def create_samples(
    config_path: str | None = None,
    papers_dir: str = "papers",
    output_dir: str = "outputs",
) -> list[Sample]:
    """Create Inspect AI samples from papers configuration.

    Args:
        config_path: Path to YAML config file (optional)
        papers_dir: Directory containing papers (default: papers/)
        output_dir: Directory for results (default: outputs/)

    Returns:
        List of Inspect AI Sample objects for evaluation
    """
    papers = []

    if config_path:
        # Load from YAML config
        if not Path(config_path).exists():
            log.error(f"Config file not found: {config_path}")
            raise FileNotFoundError(f"Config file not found: {config_path}")

        try:
            paper_configs = load_papers_config(config_path)
            log.info(f"Loaded {len(paper_configs)} papers from config")
        except Exception as e:
            log.error(f"Failed to load config: {e}", exc_info=True)
            raise
    else:
        # Auto-discover from directory
        paper_paths = list_papers(papers_dir)
        if not paper_paths:
            log.warning(f"No papers found in {papers_dir}")
            return []

        paper_configs = [
            {
                "path": p,
                "name": Path(p).name,
                "token_limits": DEFAULT_COMPRESSION_LEVELS,
            }
            for p in paper_paths
        ]
        log.info(f"Discovered {len(paper_configs)} papers in {papers_dir}")

    # Create samples - one per paper (solver will loop through all token limits)
    for paper in paper_configs:
        paper_name = paper.get("name", Path(paper["path"]).name)
        token_limits = paper.get("token_limits", DEFAULT_COMPRESSION_LEVELS)
        token_tolerances = paper.get("token_tolerances")

        sample_data = CompressionSample(
            paper_path=paper["path"],
            paper_name=paper_name,
            token_limit=None,  # Not used; solver will loop through all limits
            token_limits=sorted(token_limits, reverse=True),
            token_tolerances=token_tolerances,
            output_dir=output_dir,
        )
        papers.append(Sample(input=paper_name, metadata=sample_data.model_dump()))

    log.info(
        f"Created {len(papers)} samples total (solver will process all token limits)"
    )
    return papers