45 lines
1.6 KiB
Python
45 lines
1.6 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 Dict, List, Optional
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
DATA_DIR = ROOT / "docs/tutorial/data"
|
|
SURVEY_DIR = ROOT / "docs/surveys"
|
|
FULL_RUN = os.environ.get("GREAT_AI_NOTEBOOK_FULL") == "1"
|
|
|
|
|
|
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: List[Dict[str, int]] = [{}, {}]
|
|
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
|