feat: add data loaders

README.md

@@ -9,7 +9,7 @@
 
 # FanOutQA
 
-Read the paper! | [Download the dataset!](/data)
+Read the paper! | [Download the dataset!](/fanoutqa/data)
 
 FanOutQA is a high quality, multi-hop, multi-document benchmark for large language models using English Wikipedia as its
 knowledge base. Compared to other question-answering benchmarks, FanOutQA requires reasoning over a greater number of
@@ -88,14 +88,19 @@ used in the human-written decompositions for our Evidence Provided task.
 class TestQuestion:
     id: str
     question: str
-    necessary_evidence: list[FinalEvidence]
+    necessary_evidence: list[Evidence]
     categories: list[str]
 ```
 
 ## Wikipedia Retrieval
 
 TODO
 
+To save on time waiting for requests and computation power (both locally and on Wikipedia's end), this package
+aggressively caches retrieved Wikipedia pages.
+
+TODO: instructions for setting cache and downloading cache from server
+
 ## Evaluation
 
 To evaluate a model's generation, first ensure that you have installed all the evaluation dependencies (see above).

fanoutqa/__init__.py

@@ -0,0 +1,2 @@
+from .models import DevQuestion, TestQuestion
+from .utils import load_dev, load_test

fanoutqa/models.py

@@ -0,0 +1,99 @@
+from dataclasses import dataclass
+from typing import Optional, Union
+
+Primitive = Union[bool, int, float, str]
+
+
+@dataclass
+class Evidence:
+    """A reference to a Wikipedia article at a given point in time."""
+
+    pageid: int
+    """Wikipedia page ID"""
+
+    revid: int
+    """Wikipedia revision ID of page as of dataset epoch"""
+
+    title: str
+    """Title of page"""
+
+    url: str
+    """Link to page"""
+
+    @classmethod
+    def from_dict(cls, d):
+        return cls(**d)
+
+
+@dataclass
+class DevSubquestion:
+    """A human-written decomposition of a top-level question."""
+
+    id: str
+    question: str
+    decomposition: list["DevSubquestion"]
+    answer: Union[dict[str, Primitive], list[Primitive], Primitive]
+    """the answer to this subquestion"""
+
+    depends_on: list[str]
+    """the IDs of subquestions that this subquestion requires answering first"""
+
+    evidence: Optional[Evidence]
+    """if this is None, the question will have a decomposition"""
+
+    @classmethod
+    def from_dict(cls, d):
+        decomposition = [DevSubquestion.from_dict(dc) for dc in d["decomposition"]]
+        evidence = None if d["evidence"] is None else Evidence.from_dict(d["evidence"])
+        return cls(
+            id=d["id"],
+            question=d["question"],
+            decomposition=decomposition,
+            answer=d["answer"],
+            depends_on=d["depends_on"],
+            evidence=evidence,
+        )
+
+
+@dataclass
+class DevQuestion:
+    """A top-level question in the FOQA dataset and its decomposition."""
+
+    id: str
+    question: str
+    """the top-level question to answer"""
+    decomposition: list[DevSubquestion]
+    """human-written decomposition of the question"""
+    answer: Union[dict[str, Primitive], list[Primitive], Primitive]
+    categories: list[str]
+
+    @classmethod
+    def from_dict(cls, d):
+        decomposition = [DevSubquestion.from_dict(dc) for dc in d["decomposition"]]
+        return cls(
+            id=d["id"],
+            question=d["question"],
+            decomposition=decomposition,
+            answer=d["answer"],
+            categories=d["categories"],
+        )
+
+
+@dataclass
+class TestQuestion:
+    """A top-level question in the FOQA dataset, without its decomposition or answer."""
+
+    id: str
+    question: str
+    necessary_evidence: list[Evidence]
+    categories: list[str]
+
+    @classmethod
+    def from_dict(cls, d):
+        evidence = [Evidence.from_dict(e) for e in d["evidence"]]
+        return cls(
+            id=d["id"],
+            question=d["question"],
+            necessary_evidence=evidence,
+            categories=d["categories"],
+        )

fanoutqa/utils.py

@@ -0,0 +1,35 @@
+import json
+import os
+from pathlib import Path
+from typing import TypeAlias, Union
+
+from .models import DevQuestion, TestQuestion
+
+AnyPath: TypeAlias = Union[str, bytes, os.PathLike]
+PKG_ROOT = Path(__file__).parent
+
+
+def load_dev(fp: AnyPath = None) -> list[DevQuestion]:
+    """Load all questions from the development set.
+
+    :param fp: The path to load the questions from (defaults to bundled FOQA).
+    """
+    if fp is None:
+        fp = PKG_ROOT / "data/fanout-final-dev.json"
+
+    with open(fp) as f:
+        data = json.load(f)
+    return [DevQuestion.from_dict(d) for d in data]
+
+
+def load_test(fp: AnyPath = None) -> list[TestQuestion]:
+    """Load all questions from the test set.
+
+    :param fp: The path to load the questions from (defaults to bundled FOQA).
+    """
+    if fp is None:
+        fp = PKG_ROOT / "data/fanout-final-test.json"
+
+    with open(fp) as f:
+        data = json.load(f)
+    return [TestQuestion.from_dict(d) for d in data]