Spaces:
Sleeping
Sleeping
File size: 3,868 Bytes
79b7d56 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | """Sample/mock data for testing the Search RL Environment.
This module provides sample data that matches the exact format produced by
the data generators in data/generator/. Use this for testing and development.
Production file structure (one JSON per seed):
sample/
├── instagram.json # Tech: Instagram task (level 0)
├── whatsapp.json # Tech: WhatsApp task (level 1)
├── facebook_acquisitions.json # Tech: Multi-source task (level 2)
├── curie.json # Science: Marie Curie task (level 0)
└── berlin_wall.json # History: Berlin Wall task (level 0)
Each file contains:
{
"seed": "topic_name",
"domain": "tech|science|history",
"tasks": [
{"level": 0, "truth": "...", "supporting_items": [...], ...},
{"level": 1, ...}, // extension tasks (if any)
]
}
Usage:
from sample import get_sample_tasks, get_sample_tasks_by_level
tasks = get_sample_tasks()
level_0_tasks = get_sample_tasks_by_level(0)
"""
from __future__ import annotations
import json
from pathlib import Path
from typing import Any
try:
from searcharena.models import SearchTask
except ImportError:
from models import SearchTask
_SAMPLE_DIR = Path(__file__).parent
_cached_tasks: list[SearchTask] | None = None
def _load_json(file_path: Path) -> dict[str, Any]:
with open(file_path, "r", encoding="utf-8") as f:
return json.load(f)
def _load_tasks_from_files() -> list[SearchTask]:
"""Load all tasks from per-seed JSON files in sample/."""
all_tasks: list[SearchTask] = []
# Get all JSON files except __pycache__ etc
task_files = list(_SAMPLE_DIR.glob("*.json"))
for task_file in task_files:
try:
data = _load_json(task_file)
except (json.JSONDecodeError, OSError) as e:
print(f"Warning: Skipping {task_file}: {e}")
continue
tasks_data = data.get("tasks", [])
for task_data in tasks_data:
try:
task = SearchTask(**task_data)
all_tasks.append(task)
except Exception as e:
print(f"Warning: Skipping invalid task in {task_file}: {e}")
continue
return all_tasks
def get_sample_tasks() -> list[SearchTask]:
"""Load all sample tasks from per-seed JSON files (cached after first call).
Tasks match the format produced by the data generators:
- level: int (0, 1, 2, ...)
- truth: str (the ground truth answer)
- supporting_items: list of SupportingItem
- items_and_contents: dict mapping chunk IDs to content
- valid_distractors: list of DistractorItem
- distractors_and_contents: dict mapping distractor IDs to content
- clues: str
- truth_type: str
- passed_verification: bool
"""
global _cached_tasks
if _cached_tasks is not None:
return list(_cached_tasks)
result = _load_tasks_from_files()
_cached_tasks = result
return list(result)
def get_sample_tasks_by_level(level: int) -> list[SearchTask]:
"""Get sample tasks filtered by level."""
return [t for t in get_sample_tasks() if t.level == level]
def get_sample_tasks_by_domain(domain: str) -> list[SearchTask]:
"""Get sample tasks filtered by domain."""
return [t for t in get_sample_tasks() if t.domain == domain]
def get_sample_statistics() -> dict[str, Any]:
"""Get statistics about sample tasks."""
all_tasks = get_sample_tasks()
by_level: dict[int, int] = {}
for task in all_tasks:
by_level[task.level] = by_level.get(task.level, 0) + 1
return {
"total_tasks": len(all_tasks),
"by_level": by_level,
"domains": list(set(t.domain for t in all_tasks)),
"files": [f.name for f in _SAMPLE_DIR.glob("*.json")],
}
|