Make AbstentionBench pip-installable

.gitignore

@@ -204,4 +204,5 @@ __marimo__/
 
 
 playground*.ipynb
-data/*
+data/*
+uv.lock

README.md

@@ -37,7 +37,7 @@ You can explore existing abstention results without special installations. Simpl
 ```python
 import pandas as pd
 
-df = pd.read_csv("analysis/abstention_performance.csv")
+df = pd.read_csv("abstentionbench/analysis/abstention_performance.csv")
 df
 ```
 
@@ -151,7 +151,7 @@ To set up to run with FreshQA:
 1. Review the list of available FreshQA Google Sheets on the project's [README](https://github.com/freshllms/freshqa?tab=readme-ov-file)
 2. Pick a 'baseline' date, which should be before the model's pretraining cut-off, and an 'updated' date, which should be after the cut-off
 3. For both baseline and updated, open the Google Sheet and export their contents as CSV (File > Download > Comma Separated Values)
-4. Move these two CSVs into `data/freshqa/`, and update the paths in `configs/dataset/freshqa.yaml` accordingly
+4. Move these two CSVs into `abstentionbench/data/freshqa/`, and update the paths in `abstentionbench/configs/dataset/freshqa.yaml` accordingly
 
 Note: To exactly replicate our work, use [FreshQA_v10282024](https://docs.google.com/spreadsheets/d/1j6qr14l8oK_7gJ_XdnTBi8Pj1NVt5yeQEBKxFPkKn_g/edit?gid=334049794#gid=334049794) and [FreshQA_v12182024](https://docs.google.com/spreadsheets/d/1llFQDYuwX95L7yYDQ4aLCwJmkEh9VOSHNu6g7HjT8e0/edit?gid=334049794#gid=334049794) as baseline and updated respectively.
 
@@ -169,7 +169,7 @@ For a simple test of the end-to-end pipeline on your local machine, run:
 python main.py -m mode=local model=dummy abstention_detector=contains_abstention_keyword run_single_job_for_inference_and_judge=True
 ```
 
-Where the `dataset` arg corresponds to config file names in `configs/dataset/`, and `model` to `config/model/`. See below for supported models and datasets.
+Where the `dataset` arg corresponds to config file names in `abstentionbench/configs/dataset/`, and `model` to `abstentionbench/configs/model/`. See below for supported models and datasets.
 
 By default, `main.py` runs the whole pipeline from inference, through abstention detection, to evaluation. The pipeline launches seaparate jobs for inference and the rest of the evaluations to support a separate LLM judge. To launch everything in a single job for debugging use the `run_single_job_for_inference_and_judge=True` flag.
 
@@ -178,7 +178,7 @@ Model responses and evaluation results will be saved in `save_dir` specified in
 
 To run a fast subset on with at most 100 examples per dataset: 
 ```
-python main.py -m mode=cluster dataset='glob(*,exclude=dummy)' model=llama_3_1_8B_instruct sweep_folder=fast-subset dataset_indices_path=data/fast-subset-indices.json
+python main.py -m mode=cluster dataset='glob(*,exclude=dummy)' model=llama_3_1_8B_instruct sweep_folder=fast-subset dataset_indices_path=abstentionbench/data/fast-subset-indices.json
 ```
 
 Note `mode=cluster` expects you're running the code on a cluster that supports SLURM.
@@ -227,7 +227,7 @@ Note `mode=cluster` expects you're running the code on a cluster that supports S
 
 Note that all datasets with over a threshold 3.5k samples were truncated using uniform subsampling.
 These datasets are marked with "[subsampled]" in the list above. 
-Sample indices can be found in `data/subsampling-indices.json`. 
+Sample indices can be found in `abstentionbench/data/subsampling-indices.json`. 
 You can disable this by setting the `dataset_indices_path` config attribute to `null`. 
 
 

abstentionbench/__init__.py

@@ -0,0 +1,7 @@
+"""
+AbstentionBench: A Holistic Benchmark for LLM Abstention
+
+This package provides tools and datasets for evaluating LLM abstention capabilities.
+"""
+
+__version__ = "1.0.0"

abstentionbench/analysis/__init__.py

@@ -0,0 +1,10 @@
+"""
+Analysis package for AbstentionBench.
+
+Provides utilities for loading results and generating tables.
+"""
+
+from .load_results import Results
+from .tables import AbstentionF1ScoreTable
+
+__all__ = ["Results", "AbstentionF1ScoreTable"]

analysis/load_results.py → abstentionbench/analysis/load_results.py

@@ -11,16 +11,16 @@
 
 import pandas as pd
 
-from recipe.abstention import Responses
-from recipe.job_manager import JobManager
+from ..recipe.abstention import Responses
+from ..recipe.job_manager import JobManager
 
 logger = logging.getLogger(__name__)
 
 
 class Results:
     """Loads results.
     By default uses the paths specified in
-        analysis/results_path.py
+        abstentionbench/analysis/results_path.py
 
     Attributes:
         df: pandas dataframe that turns saved results into a table

abstentionbench/configs/__init__.py

@@ -0,0 +1 @@
+# Configuration files for AbstentionBench

configs/dataset/alcuna.yaml → abstentionbench/configs/dataset/alcuna.yaml

@@ -3,6 +3,6 @@
 dataset_name: ALCUNADataset
 
 datamodule:
-  _target_: recipe.abstention_datasets.alcuna.${dataset_name}
+  _target_: abstentionbench.recipe.abstention_datasets.alcuna.${dataset_name}
   data_dir: ${data_dir}/alcuna
   max_num_samples: null

configs/dataset/bbq.yaml → abstentionbench/configs/dataset/bbq.yaml

@@ -3,7 +3,7 @@
 dataset_name: BBQDataset
 
 datamodule:
-  _target_: recipe.abstention_datasets.bbq.${dataset_name}
+  _target_: abstentionbench.recipe.abstention_datasets.bbq.${dataset_name}
   data_dir: ${data_dir}/bbq
   data_url_base: "https://raw.githubusercontent.com/nyu-mll/BBQ/refs/heads/main/data/"
   max_num_samples: null

configs/dataset/big_bench_disambiguate.yaml → abstentionbench/configs/dataset/big_bench_disambiguate.yaml

@@ -3,7 +3,7 @@
 dataset_name: BigBenchDisambiguateDataset
 
 datamodule:
-  _target_: recipe.abstention_datasets.big_bench.${dataset_name}
+  _target_: abstentionbench.recipe.abstention_datasets.big_bench.${dataset_name}
   data_dir: ${data_dir}/big_bench_disambiguate
   data_url: "https://raw.githubusercontent.com/suzgunmirac/BIG-Bench-Hard/refs/heads/main/bbh/disambiguation_qa.json"
   max_num_samples: null

configs/dataset/big_bench_known_unknowns.yaml → abstentionbench/configs/dataset/big_bench_known_unknowns.yaml

@@ -3,5 +3,5 @@
 dataset_name: BigBenchKnownUnknownsDataset
 
 datamodule:
-  _target_: recipe.abstention_datasets.big_bench.${dataset_name}
+  _target_: abstentionbench.recipe.abstention_datasets.big_bench.${dataset_name}
   max_num_samples: null

configs/dataset/coconot.yaml → abstentionbench/configs/dataset/coconot.yaml

@@ -3,5 +3,5 @@
 dataset_name: CoCoNotDataset
 
 datamodule:
-  _target_: recipe.abstention_datasets.coconot.${dataset_name}
+  _target_: abstentionbench.recipe.abstention_datasets.coconot.${dataset_name}
   max_num_samples: null

abstentionbench/configs/dataset/dummy.yaml

@@ -0,0 +1,7 @@
+# @package _global_
+
+dataset_name: DummyDataset
+
+datamodule:
+  _target_: abstentionbench.recipe.abstention_datasets.abstract_abstention_dataset.DummyDataset
+  max_num_samples: 100

configs/dataset/falseqa.yaml → abstentionbench/configs/dataset/falseqa.yaml

@@ -3,7 +3,7 @@
 dataset_name: FalseQADataset
 
 datamodule:
-  _target_: recipe.abstention_datasets.false_qa.${dataset_name}
+  _target_: abstentionbench.recipe.abstention_datasets.false_qa.${dataset_name}
   data_dir: ${data_dir}/falseqa
   data_url: "https://raw.githubusercontent.com/thunlp/FalseQA/refs/heads/main/dataset/test.csv"
   max_num_samples: null

configs/dataset/freshqa.yaml → abstentionbench/configs/dataset/freshqa.yaml

@@ -3,7 +3,7 @@
 dataset_name: FreshQADataset
 
 datamodule:
-  _target_: recipe.abstention_datasets.freshqa.${dataset_name}
+  _target_: abstentionbench.recipe.abstention_datasets.freshqa.${dataset_name}
   path_baseline: ${data_dir}/freshqa/FreshQA_v10282024.csv
   path_updated: ${data_dir}/freshqa/FreshQA_v12182024.csv
   max_num_samples: null

configs/dataset/gpqa.yaml → abstentionbench/configs/dataset/gpqa.yaml

@@ -3,6 +3,6 @@
 dataset_name: GPQA
 
 datamodule:
-  _target_: recipe.abstention_datasets.gpqa.GPQA
+  _target_: abstentionbench.recipe.abstention_datasets.gpqa.GPQA
   subset: "gpqa_diamond"
   max_num_samples: null

configs/dataset/gsm8k.yaml → abstentionbench/configs/dataset/gsm8k.yaml

@@ -3,6 +3,6 @@
 dataset_name: GSM8K
 
 datamodule:
-  _target_: recipe.abstention_datasets.gsm8k.GSM8K
+  _target_: abstentionbench.recipe.abstention_datasets.gsm8k.GSM8K
   split: "test"
   max_num_samples: null

abstentionbench/configs/dataset/kuq.yaml

@@ -0,0 +1,8 @@
+# @package _global_
+
+dataset_name: KUQDataset
+
+datamodule:
+  _target_: abstentionbench.recipe.abstention_datasets.kuq.KUQDataset
+  max_num_samples: null
+  category_map_path: ${data_dir}/kuq/new-category-mapping.csv

configs/dataset/mediq.yaml → abstentionbench/configs/dataset/mediq.yaml

@@ -3,7 +3,7 @@
 dataset_name: MediQDataset
 
 datamodule:
-  _target_: recipe.abstention_datasets.mediq.${dataset_name}
+  _target_: abstentionbench.recipe.abstention_datasets.mediq.${dataset_name}
   data_dir: ${data_dir}/mediq
   icraftmd_url: "https://raw.githubusercontent.com/stellalisy/mediQ/refs/heads/main/data/all_craft_md.jsonl"
   imedqa_url: "https://raw.githubusercontent.com/stellalisy/mediQ/refs/heads/main/data/all_dev_good.jsonl"

configs/dataset/mmlu_history.yaml → abstentionbench/configs/dataset/mmlu_history.yaml

@@ -3,6 +3,6 @@
 dataset_name: MMLUHistory
 
 datamodule:
-  _target_: recipe.abstention_datasets.mmlu.MMLUHistory
+  _target_: abstentionbench.recipe.abstention_datasets.mmlu.MMLUHistory
   split: "test"
   max_num_samples: null

configs/dataset/mmlu_math.yaml → abstentionbench/configs/dataset/mmlu_math.yaml

@@ -3,6 +3,6 @@
 dataset_name: MMLUMath
 
 datamodule:
-  _target_: recipe.abstention_datasets.mmlu.MMLUMath
+  _target_: abstentionbench.recipe.abstention_datasets.mmlu.MMLUMath
   split: "test"
   max_num_samples: null

configs/dataset/moralchoice.yaml → abstentionbench/configs/dataset/moralchoice.yaml

@@ -3,7 +3,7 @@
 dataset_name: MoralChoiceDataset
 
 datamodule:
-  _target_: recipe.abstention_datasets.moralchoice.${dataset_name}
+  _target_: abstentionbench.recipe.abstention_datasets.moralchoice.${dataset_name}
   template_type: "ab"
   system_instruction: True
   max_num_samples: null

configs/dataset/musique.yaml → abstentionbench/configs/dataset/musique.yaml

@@ -3,6 +3,6 @@
 dataset_name: MusiqueDataset
 
 datamodule:
-  _target_: recipe.abstention_datasets.musique.${dataset_name}
+  _target_: abstentionbench.recipe.abstention_datasets.musique.${dataset_name}
   data_dir: ${data_dir}/musique
   max_num_samples: null

configs/dataset/qaqa.yaml → abstentionbench/configs/dataset/qaqa.yaml

@@ -3,6 +3,6 @@
 dataset_name: QAQADataset
 
 datamodule:
-  _target_: recipe.abstention_datasets.qaqa.${dataset_name}
+  _target_: abstentionbench.recipe.abstention_datasets.qaqa.${dataset_name}
   data_dir: ${data_dir}/qaqa
   max_num_samples: null

configs/dataset/qasper.yaml → abstentionbench/configs/dataset/qasper.yaml

@@ -7,6 +7,6 @@ dataset_name: QASPERDataset
 # empty responses. 
 
 datamodule:
-  _target_: recipe.abstention_datasets.qasper.${dataset_name}
+  _target_: abstentionbench.recipe.abstention_datasets.qasper.${dataset_name}
   data_dir: ${data_dir}/qasper
   max_num_samples: null

configs/dataset/self_aware.yaml → abstentionbench/configs/dataset/self_aware.yaml

@@ -3,7 +3,7 @@
 dataset_name: SelfAwareDataset
 
 datamodule:
-  _target_: recipe.abstention_datasets.self_aware.SelfAwareDataset
+  _target_: abstentionbench.recipe.abstention_datasets.self_aware.SelfAwareDataset
   data_dir: ${data_dir}/selfaware
   data_url: "https://raw.githubusercontent.com/yinzhangyue/SelfAware/refs/heads/main/data/SelfAware.json"
   max_num_samples: null

configs/dataset/situated_qa.yaml → abstentionbench/configs/dataset/situated_qa.yaml

@@ -3,5 +3,5 @@
 dataset_name: SituatedQAGeoDataset
 
 datamodule:
-  _target_: recipe.abstention_datasets.situated_qa.${dataset_name}
+  _target_: abstentionbench.recipe.abstention_datasets.situated_qa.${dataset_name}
   max_num_samples: null

configs/dataset/squad2.yaml → abstentionbench/configs/dataset/squad2.yaml

@@ -3,5 +3,5 @@
 dataset_name: Squad2Dataset
 
 datamodule:
-  _target_: recipe.abstention_datasets.squad.${dataset_name}
+  _target_: abstentionbench.recipe.abstention_datasets.squad.${dataset_name}
   max_num_samples: null

configs/dataset/umwp.yaml → abstentionbench/configs/dataset/umwp.yaml

@@ -3,6 +3,6 @@
 dataset_name: UMWP
 
 datamodule:
-  _target_: recipe.abstention_datasets.umwp.${dataset_name}
+  _target_: abstentionbench.recipe.abstention_datasets.umwp.${dataset_name}
   data_dir: ${data_dir}/umwp
   max_num_samples: null

configs/dataset/worldsense.yaml → abstentionbench/configs/dataset/worldsense.yaml

@@ -3,7 +3,7 @@
 dataset_name: WorldSenseDataset
 
 datamodule:
-  _target_: recipe.abstention_datasets.world_sense.${dataset_name}
+  _target_: abstentionbench.recipe.abstention_datasets.world_sense.${dataset_name}
   data_dir: ${data_dir}/world_sense
   raw_data_url: "https://github.com/facebookresearch/worldsense/raw/refs/heads/main/data/worldsense/training_set/trials_10k.jsonl.bz2"
   max_num_samples: null

configs/default_pipeline.yaml → abstentionbench/configs/default_pipeline.yaml

@@ -43,8 +43,8 @@ only_run_inference: False
 # name to list of indices.
 dataset_indices_subset: null
 # fast subset
-# dataset_indices_path: data/fast-subset-indices.json
-dataset_indices_path: data/subsampling-indices.json
+# dataset_indices_path: abstentionbench/data/fast-subset-indices.json
+dataset_indices_path: abstentionbench/data/subsampling-indices.json
 
 seed: 123456789
 

abstentionbench/data/__init__.py

@@ -0,0 +1,35 @@
+"""
+Data files and utilities for AbstentionBench datasets.
+"""
+import os
+from pathlib import Path
+try:
+    from importlib import resources
+except ImportError:
+    # Python < 3.9 fallback
+    import importlib_resources as resources
+
+
+def get_data_dir(subdir=None):
+    """Get the path to the data directory, creating it if it doesn't exist.
+
+    Args:
+        subdir: Optional subdirectory within the data directory
+
+    Returns:
+        Path to the data directory or subdirectory
+    """
+    try:
+        # When installed as a package, use the installed data directory
+        with resources.path(__name__, '.') as data_path:
+            base_dir = Path(data_path)
+    except (ImportError, AttributeError, FileNotFoundError):
+        # Fallback for development or when resources don't work
+        base_dir = Path(__file__).parent
+
+    if subdir:
+        full_path = base_dir / subdir
+        full_path.mkdir(parents=True, exist_ok=True)
+        return str(full_path)
+    else:
+        return str(base_dir)

recipe/abstention.py → abstentionbench/recipe/abstention.py

@@ -14,7 +14,7 @@
 
 from pydantic import BaseModel
 
-from recipe.inference import LoadSaveDataMixin, RawResponse, RawResponses
+from .inference import LoadSaveDataMixin, RawResponse, RawResponses
 
 
 class Response(RawResponse):