61 lines
2.3 KiB
Python
61 lines
2.3 KiB
Python
"""Shared local paths and small datasets for the documentation notebooks."""
|
|
|
|
import json
|
|
import os
|
|
from collections import Counter
|
|
from pathlib import Path
|
|
from typing import List, Optional
|
|
|
|
|
|
def setup_notebook(root: Path, 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 os.environ.get("GREAT_AI_NOTEBOOK_FULL") == "1" 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()
|
|
lines = []
|
|
with path.open(encoding="utf-8") as stream:
|
|
for line in stream:
|
|
label = json.loads(line)["label"]
|
|
if per_label is None or counts[label] < per_label:
|
|
counts[label] += 1
|
|
lines.append(line)
|
|
return lines
|
|
|
|
|
|
def demo_annotations() -> List[dict]:
|
|
"""Synthetic annotations for exercising the pipeline, not measuring quality."""
|
|
annotations = [{}, {}]
|
|
for index, topic in enumerate(
|
|
["protein", "cancer", "solar", "quantum", "neural", "genome"]
|
|
):
|
|
for sample in range(10):
|
|
positive = sample % 2
|
|
sentence = (
|
|
f"Our {topic} method improves prediction accuracy in experiment {sample}."
|
|
if positive
|
|
else f"The {topic} dataset contains observations from experiment {sample}."
|
|
)
|
|
annotations[0][sentence] = positive
|
|
annotations[1][sentence] = (
|
|
1 - positive if (index * 10 + sample) % 11 == 0 else positive
|
|
)
|
|
return annotations
|