Spaces:
Sleeping
Sleeping
File size: 5,366 Bytes
32d14f4 | 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 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 | #!/usr/bin/env python
"""Prompt-based LLM/static-prior runner for BrainRL.
Runs one (or more) full episodes through the OpenEnv environment with the
selected condition / subject context. With ``--use-llm`` it asks a chat model
for one JSON action per OpenEnv step; without, it uses the deterministic
prompt fallback so the demo always works.
"""
from __future__ import annotations
import argparse
from pathlib import Path
from data_split import DEFAULT_PARTICIPANT_INFO, build_condition_split
from prompts import build_action_prompt, llm_prompt_policy_action, static_prompt_policy_action
from server.brain_environment import BrainRegionSelectionEnvironment
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description="Run prompt-based inference on BrainRL")
parser.add_argument("--use-llm", action="store_true", help="Ask an LLM for each action")
parser.add_argument("--show-prompts", action="store_true", help="Print prompt every step")
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--episodes", type=int, default=1, help="How many episodes to roll out")
parser.add_argument(
"--condition",
choices=("single_m", "single_f", "mixed_m", "mixed_f"),
default=None,
)
parser.add_argument(
"--participant-info",
type=str,
default=str(DEFAULT_PARTICIPANT_INFO),
)
parser.add_argument("--train-subjects", type=str, default=None)
parser.add_argument("--test-subjects", type=str, default=None)
parser.add_argument(
"--exclude-subjects",
type=str,
default=None,
help="Comma list / range of subjects to drop (e.g. corrupted recordings).",
)
parser.add_argument(
"--split",
choices=("train", "test", "all"),
default="all",
)
parser.add_argument("--subject-id", type=str, default=None, help="Override single subject id")
parser.add_argument("--run-id", type=str, default=None, help="Override single run id (run1..run4)")
return parser
def _episode_pairs(args: argparse.Namespace) -> list[dict[str, str | None]]:
if args.subject_id:
return [
{
"subject_id": args.subject_id,
"run_id": args.run_id,
"condition": args.condition,
}
]
if args.condition is None:
return [{"subject_id": None, "run_id": None, "condition": None}]
split = build_condition_split(
condition=args.condition,
participant_info_path=args.participant_info,
train_subjects_spec=args.train_subjects,
test_subjects_spec=args.test_subjects,
exclude_subjects=args.exclude_subjects,
)
pairs = split.pairs_for(args.split)
if not pairs:
raise SystemExit(f"No pairs for {args.condition}/{args.split}.")
return [pair.as_dict() for pair in pairs]
def run_one_episode(env: BrainRegionSelectionEnvironment, args: argparse.Namespace) -> None:
stim = env._stimulus_features
stim_label = (
f" stimulus=window={stim['window_index'] + 1}/{stim['n_windows']} "
f"dominant_pos={stim.get('dominant_pos')} n_words={stim.get('n_words')} "
f"density={stim.get('speech_density'):.2f}"
if stim
else " stimulus=unavailable"
)
print(
f"[INIT] subject={env._subject_id} run={env._run_id} condition={env._condition} "
f"candidates={env._subset.n_regions} budget={env._subset.selection_budget} "
f"prompt_top_k={env._subset.prompt_top_k}" + stim_label
)
rewards: list[float] = []
done = False
while not done:
state = env._build_selection_state()
if args.show_prompts:
print("[PROMPT]")
print(build_action_prompt(state))
if args.use_llm:
region_id, raw_text = llm_prompt_policy_action(state)
else:
region_id = static_prompt_policy_action(state)
raw_text = f'{{"region_id": "{region_id}"}}'
result = env._process_action(region_id)
rewards.append(float(result["reward"]))
print(
f"[STEP {env._timestep}/{env._subset.selection_budget}] region={region_id} "
f"reward={float(result['reward']):.4f} r2={float(result['current_r2']):.4f} "
f"done={result['done']}"
)
if args.use_llm:
print(f"[MODEL] {raw_text}")
done = bool(result["done"])
print(
f"[END] final_r2={env._current_r2:.4f} total_reward={sum(rewards):.4f} "
f"selected={len(env._selected_region_ids)} "
f"order={' -> '.join(env._selected_region_ids[:10])}"
+ (" ..." if len(env._selected_region_ids) > 10 else "")
)
def main() -> None:
args = build_parser().parse_args()
pairs = _episode_pairs(args)
env = BrainRegionSelectionEnvironment()
for episode_idx in range(int(args.episodes)):
pair = pairs[episode_idx % len(pairs)]
env.reset(
seed=int(args.seed) + episode_idx,
subject_id=pair.get("subject_id"),
run_id=pair.get("run_id"),
condition=pair.get("condition"),
)
if int(args.episodes) > 1:
print(f"\n=== Episode {episode_idx + 1}/{args.episodes} ===")
run_one_episode(env, args)
if __name__ == "__main__":
main()
|