Simplify notebook setup

This commit is contained in:
Andras Schmelczer 2026-10-28 22:34:22 +00:00
commit 13dadc74ed

View file

@ -12,27 +12,6 @@ SURVEY_DIR = ROOT / "docs/surveys"
FULL_RUN = os.environ.get("GREAT_AI_NOTEBOOK_FULL") == "1"
def setup_notebook(workflow: str) -> Path:
"""Keep related notebooks' models and traces together, outside source files."""
from great_ai import ParallelTinyDbDriver, configure
from great_ai.large_file import LargeFileLocal
base = Path(os.environ.get("GREAT_AI_NOTEBOOK_WORKDIR", ROOT / ".cache/notebooks"))
suffix = "-full" if FULL_RUN else ""
work_dir = (base / (workflow + suffix)).resolve()
work_dir.mkdir(parents=True, exist_ok=True)
LargeFileLocal.cache_path = work_dir / "models-cache"
ParallelTinyDbDriver.path_to_db = work_dir / "tracing_database.json"
configure(
large_file_implementation=LargeFileLocal,
tracing_database_factory=ParallelTinyDbDriver,
disable_se4ml_banner=True,
seed=42,
)
print(f"Notebook artifacts: {work_dir}")
return work_dir
def mag_lines(path: Path, per_label: Optional[int] = None) -> List[str]:
"""Read the bundled MAG data, optionally limiting each class for a quick run."""
counts: Counter = Counter()