Coverage for src / local_deep_research / benchmarks / datasets / __init__.py: 83%

1""" 

2Dataset handling for benchmarks. 

3 

4This module provides functions and classes for loading and processing 

5benchmark datasets in a maintainable, extensible way. 

6""" 

7 

8from .base import DatasetRegistry 

9from .browsecomp import BrowseCompDataset 

10from .simpleqa import SimpleQADataset 

11from .xbench_deepsearch import XBenchDeepSearchDataset 

12 

13# Register datasets 

14DatasetRegistry.register(SimpleQADataset) 

15DatasetRegistry.register(BrowseCompDataset) 

16DatasetRegistry.register(XBenchDeepSearchDataset) 

17 

18# Add convenience functions mirroring the old interface 

19DEFAULT_DATASET_URLS = { 

20 "simpleqa": SimpleQADataset.get_default_dataset_path(), 

21 "browsecomp": BrowseCompDataset.get_default_dataset_path(), 

22 "xbench_deepsearch": XBenchDeepSearchDataset.get_default_dataset_path(), 

23} 

24 

25 

26def get_available_datasets(): 

27 """Get information about all registered datasets.""" 

28 return DatasetRegistry.get_available_datasets() 

29 

30 

31def load_dataset( 

32 dataset_type, 

33 dataset_path=None, 

34 num_examples=None, 

35 seed=42, 

36): 

37 """Load a dataset by name. 

38 

39 This is a backwards-compatible function that uses the new dataset classes 

40 under the hood. It exists to maintain compatibility with existing code. 

41 

42 Args: 

43 dataset_type: Type of dataset ('simpleqa' or 'browsecomp') 

44 dataset_path: Path or URL to dataset (if None, uses default URL) 

45 num_examples: Optional number of examples to sample 

46 seed: Random seed for sampling 

47 

48 Returns: 

49 List of processed examples 

50 """ 

51 return DatasetRegistry.load_dataset( 

52 dataset_id=dataset_type.lower(), 

53 dataset_path=dataset_path, 

54 num_examples=num_examples, 

55 seed=seed, 

56 )