[REFACTOR] Continue refactor

Author: ZhuochengZhang98

src/flexrag/datasets/__init__.py

@@ -1,49 +0,0 @@
-# datasets
-from .corpora import IterableCorpus, MappingCorpus
-from .dataset import ChainDataset, ConcatDataset, IterableDataset, MappingDataset
-from .hf_dataset import HFDataset, HFDatasetConfig
-from .qa_dataset import (
-    QA_DATASETS,
-    FlashQADataset,
-    FlashQADatasetConfig,
-    QADataset,
-    QADatasetConfig,
-    QAEvalData,
-)
-from .retrieval_datasets import (
-    RETRIEVAL_DATASETS,
-    IREvalData,
-    MSMARCODataset,
-    MSMARCODatasetConfig,
-    MTEBDataset,
-    MTEBDatasetConfig,
-    MultiLongDocRetrievalDataset,
-    MultiLongDocRetrievalDatasetConfig,
-    RetrievalDatasetBase,
-)
-
-__all__ = [
-    "ChainDataset",
-    "IterableDataset",
-    "MappingDataset",
-    "ConcatDataset",
-    "HFDataset",
-    "HFDatasetConfig",
-    "IterableCorpus",
-    "MappingCorpus",
-    "QA_DATASETS",
-    "FlashQADataset",
-    "FlashQADatasetConfig",
-    "QADataset",
-    "QAEvalData",
-    "QADatasetConfig",
-    "MSMARCODatasetConfig",
-    "MSMARCODataset",
-    "MTEBDataset",
-    "MTEBDatasetConfig",
-    "RETRIEVAL_DATASETS",
-    "IREvalData",
-    "MultiLongDocRetrievalDataset",
-    "MultiLongDocRetrievalDatasetConfig",
-    "RetrievalDatasetBase",
-]

src/flexrag/datasets/benchmarks/__init__.py

@@ -1,8 +1,8 @@
 from .multiple_choice import (
     LongBenchV2Dataset,
     LongBenchV2DatasetConfig,
-    NovelQADataset,
     NovelQAConfig,
+    NovelQADataset,
     QuALITYDataset,
     QuALITYDatasetConfig,
 )
@@ -30,6 +30,15 @@
     SQuADDataset,
     SQuADDatasetConfig,
 )
+from .retrieval import (
+    MSMARCODataset,
+    MSMARCODatasetConfig,
+    MTEBDataset,
+    MTEBDatasetConfig,
+    MultiLongDocRetrievalDataset,
+    MultiLongDocRetrievalDatasetConfig,
+    RetrievalDatasetBase,
+)
 from .suites import KiltDataset, KiltDatasetConfig
 
 __all__ = [
@@ -61,6 +70,13 @@
     "SimpleQADatasetConfig",
     "SQuADDataset",
     "SQuADDatasetConfig",
+    "MSMARCODataset",
+    "MSMARCODatasetConfig",
+    "MTEBDataset",
+    "MTEBDatasetConfig",
+    "MultiLongDocRetrievalDataset",
+    "MultiLongDocRetrievalDatasetConfig",
+    "RetrievalDatasetBase",
     "KiltDataset",
     "KiltDatasetConfig",
 ]

…/datasets/retrieval_datasets/__init__.py …atasets/benchmarks/retrieval/__init__.pysrc/flexrag/datasets/retrieval_datasets/__init__.py renamed to src/flexrag/datasets/benchmarks/retrieval/__init__.py src/flexrag/datasets/retrieval_datasets/__init__.py renamed to src/flexrag/datasets/benchmarks/retrieval/__init__.py

@@ -4,16 +4,14 @@
 )
 from .msmarco_dataset import MSMARCODataset, MSMARCODatasetConfig
 from .mteb_dataset import MTEBDataset, MTEBDatasetConfig
-from .retrieval_dataset import RETRIEVAL_DATASETS, IREvalData, RetrievalDatasetBase
+from .retrieval_dataset_base import RetrievalDatasetBase
 
 __all__ = [
+    "RetrievalDatasetBase",
     "MultiLongDocRetrievalDataset",
     "MultiLongDocRetrievalDatasetConfig",
-    "RETRIEVAL_DATASETS",
-    "IREvalData",
-    "RetrievalDatasetBase",
+    "MSMARCODataset",
+    "MSMARCODatasetConfig",
     "MTEBDataset",
     "MTEBDatasetConfig",
-    "MSMARCODatasetConfig",
-    "MSMARCODataset",
 ]

…asets/retrieval_datasets/mldr_dataset.py …ets/benchmarks/retrieval/mldr_dataset.pysrc/flexrag/datasets/retrieval_datasets/mldr_dataset.py renamed to src/flexrag/datasets/benchmarks/retrieval/mldr_dataset.py src/flexrag/datasets/retrieval_datasets/mldr_dataset.py renamed to src/flexrag/datasets/benchmarks/retrieval/mldr_dataset.py

@@ -1,14 +1,15 @@
 from collections import defaultdict
 from pathlib import Path
-from typing import Mapping
+from typing import Annotated, Mapping, Optional
 
 from huggingface_hub import snapshot_download
 
-from flexrag.common import FLEXRAG_CACHE_DIR, LOGGER_MANAGER, configure
+from flexrag.common import FLEXRAG_CACHE_DIR, LOGGER_MANAGER, Choices, configure
 from flexrag.common.dataclasses import Context
 
-from ..reader import LineDelimitedReader
-from .retrieval_dataset import RETRIEVAL_DATASETS, RetrievalDatasetBase
+from ...core import DATASETS
+from ...reader import LineDelimitedReader
+from .retrieval_dataset_base import RetrievalDatasetBase
 
 logger = LOGGER_MANAGER.get_logger("flexrag.datasets.mldr_dataset")
 
@@ -45,11 +46,28 @@ class MultiLongDocRetrievalDatasetConfig:
     """
 
     split: str
-    lang: str = "en"
-    data_path: str | None = None
-
-
-@RETRIEVAL_DATASETS("mldr", config_class=MultiLongDocRetrievalDatasetConfig)
+    lang: Annotated[
+        str,
+        Choices(
+            "ar",
+            "de",
+            "en",
+            "es",
+            "fr",
+            "hi",
+            "it",
+            "ja",
+            "ko",
+            "pt",
+            "ru",
+            "th",
+            "zh",
+        ),
+    ] = "en"
+    data_path: Optional[str] = None
+
+
+@DATASETS("mldr", config_class=MultiLongDocRetrievalDatasetConfig)
 class MultiLongDocRetrievalDataset(RetrievalDatasetBase):
     def __init__(self, config: MultiLongDocRetrievalDatasetConfig) -> None:
         # prepare dataset path
@@ -112,16 +130,13 @@ def __init__(self, config: MultiLongDocRetrievalDatasetConfig) -> None:
         return
 
     @property
-    def _contexts(self) -> Mapping[str, Context]:
-        """Return a mapping from context_id to Context object."""
-        return self._context_data
+    def _qrels(self) -> Mapping[str, Mapping[str, float]]:
+        return self._qrels_data
 
     @property
     def _queries(self) -> Mapping[str, str]:
-        """Return a mapping from query_id to query text."""
         return self._queries_data
 
     @property
-    def _qrels(self) -> Mapping[str, set[str]]:
-        """Return a mapping from query_id to a set of relevant context_ids."""
-        return self._qrels_data
+    def _contexts(self) -> Mapping[str, Context]:
+        return self._context_data

…ts/retrieval_datasets/msmarco_dataset.py …/benchmarks/retrieval/msmarco_dataset.pysrc/flexrag/datasets/retrieval_datasets/msmarco_dataset.py renamed to src/flexrag/datasets/benchmarks/retrieval/msmarco_dataset.py src/flexrag/datasets/retrieval_datasets/msmarco_dataset.py renamed to src/flexrag/datasets/benchmarks/retrieval/msmarco_dataset.py

@@ -13,8 +13,9 @@
 )
 from flexrag.common.dataclasses import Context
 
-from ..reader import LineDelimitedReader
-from .retrieval_dataset import RETRIEVAL_DATASETS, RetrievalDatasetBase
+from ...core import DATASETS
+from ...reader import LineDelimitedReader
+from .retrieval_dataset_base import RetrievalDatasetBase
 
 logger = LOGGER_MANAGER.get_logger("flexrag.datasets.msmarco_dataset")
 
@@ -62,7 +63,7 @@ class MSMARCODatasetConfig:
     load_corpus: bool = False
 
 
-@RETRIEVAL_DATASETS("msmarco", MSMARCODatasetConfig)
+@DATASETS("msmarco", MSMARCODatasetConfig)
 class MSMARCODataset(RetrievalDatasetBase):
     """Dataset for loading MSMARCO Retrieval Dataset."""
 

…asets/retrieval_datasets/mteb_dataset.py …ets/benchmarks/retrieval/mteb_dataset.pysrc/flexrag/datasets/retrieval_datasets/mteb_dataset.py renamed to src/flexrag/datasets/benchmarks/retrieval/mteb_dataset.py src/flexrag/datasets/retrieval_datasets/mteb_dataset.py renamed to src/flexrag/datasets/benchmarks/retrieval/mteb_dataset.py

@@ -6,7 +6,8 @@
 from flexrag.common import configure
 from flexrag.common.dataclasses import Context
 
-from .retrieval_dataset import RETRIEVAL_DATASETS, RetrievalDatasetBase
+from ...core import DATASETS
+from .retrieval_dataset_base import RetrievalDatasetBase
 
 
 @configure
@@ -41,7 +42,7 @@ class MTEBDatasetConfig:
     load_corpus: bool = False
 
 
-@RETRIEVAL_DATASETS("mteb", MTEBDatasetConfig)
+@DATASETS("mteb", MTEBDatasetConfig)
 class MTEBDataset(RetrievalDatasetBase):
     """Dataset for loading MTEB Retrieval Dataset."""
 

…/retrieval_datasets/retrieval_dataset.py …arks/retrieval/retrieval_dataset_base.pysrc/flexrag/datasets/retrieval_datasets/retrieval_dataset.py renamed to src/flexrag/datasets/benchmarks/retrieval/retrieval_dataset_base.py src/flexrag/datasets/retrieval_datasets/retrieval_dataset.py renamed to src/flexrag/datasets/benchmarks/retrieval/retrieval_dataset_base.py

@@ -1,39 +1,13 @@
 from abc import abstractmethod
 from collections.abc import Iterator, Mapping
-from dataclasses import field
 from functools import cached_property
-from typing import Optional
 
-from flexrag.common import Register, data
 from flexrag.common.dataclasses import Context
 
-from ..dataset import MappingDataset
+from ...core import IRSample, MappingDataset
 
 
-@data
-class IREvalData:
-    """The dataclass for Information Retrieval evaluation data.
-
-    :param question: The question for evaluation. Required.
-    :type question: str
-    :param question_id: The unique identifier for the question. Default: None.
-    :type question_id: Optional[str]
-    :param contexts: The contexts related to the question. Default: None.
-    :type contexts: Optional[list[Context]]
-    :param hard_negatives: The hard negatives related to the question. Default: None.
-    :type hard_negatives: Optional[list[Context]]
-    :param meta_data: The metadata of the evaluation data. Default: {}.
-    :type meta_data: dict
-    """
-
-    question: str
-    question_id: Optional[str] = None
-    contexts: Optional[list[Context]] = None
-    hard_negatives: Optional[list[Context]] = None
-    meta_data: dict = field(default_factory=dict)
-
-
-class RetrievalDatasetBase(MappingDataset[IREvalData]):
+class RetrievalDatasetBase(MappingDataset[IRSample]):
     """Base class for Information Retrieval (IR) datasets.
 
     This class provides a unified interface for accessing IR datasets, which typically consist of:
@@ -162,7 +136,7 @@ def __len__(self) -> int:
         """The number of queries in the qrels."""
         return len(self._qids)
 
-    def get_item(self, index: int) -> IREvalData:
+    def get_item(self, index: int) -> IRSample:
         qid = self._qids[index]
         query = self._queries[qid]
         relevant_ctxs = []
@@ -180,12 +154,9 @@ def get_item(self, index: int) -> IREvalData:
             self._contexts.get(ctx_id, Context(context_id=ctx_id))
             for ctx_id in hard_negatives
         ]
-        return IREvalData(
+        return IRSample(
             question=query,
             question_id=qid,
             contexts=rels,
             hard_negatives=negs,
         )
-
-
-RETRIEVAL_DATASETS = Register[RetrievalDatasetBase]("retrieval_dataset")