great-ai/docs/examples/simple/data.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cell-0",
   "metadata": {},
   "source": [
    "# Simple example: data engineering\n",
    "\n",
    "Train a scientific domain classifier with an explainable Naive Bayes model.\n",
    "Run `data.ipynb`, `train.ipynb`, then `deploy.ipynb` in that order.\n",
    "\n",
    "The original example downloaded large Semantic Scholar corpus shards.\n",
    "For a reproducible local run, we use the bundled MAG sentences instead,\n",
    "wrapping each domain in a list to keep the same multilabel data interface.\n",
    "By default we take 60 examples per class; `GREAT_AI_NOTEBOOK_FULL=1` uses\n",
    "the entire bundled training split. These results are not the original\n",
    "Semantic Scholar experiment's metrics.\n",
    "\n",
    "### Extract"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0df47c23",
   "metadata": {},
   "source": [
    "## Local setup\n",
    "\n",
    "Use Python 3.12 or 3.13 and install this checkout with `uv pip install --editable '.[notebooks]'`.\n",
    "See **Run notebooks locally** in the repository README for the kernel, execution order,\n",
    "and full-data options. Generated artifacts stay under `.cache/notebooks/`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a800dc4d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[38;5;226mEnvironment variable ENVIRONMENT is not set, defaulting to development mode ‼️\u001b[0m\n",
      "\u001b[38;5;226mThe selected tracing database (ParallelTinyDbDriver) is not recommended for production\u001b[0m\n",
      "\u001b[38;5;39mGreatAI (v0.1.11): configured ✅\u001b[0m\n",
      "\u001b[38;5;39m  🔩 tracing_database: ParallelTinyDbDriver\u001b[0m\n",
      "\u001b[38;5;39m  🔩 large_file_implementation: LargeFileLocal\u001b[0m\n",
      "\u001b[38;5;39m  🔩 is_production: False\u001b[0m\n",
      "\u001b[38;5;39m  🔩 should_log_exception_stack: True\u001b[0m\n",
      "\u001b[38;5;39m  🔩 prediction_cache_size: 512\u001b[0m\n",
      "\u001b[38;5;39m  🔩 dashboard_table_size: 50\u001b[0m\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Notebook artifacts: /Users/schmelczer/Projects/great-ai/.cache/notebooks/simple\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import sys\n",
    "from pathlib import Path\n",
    "\n",
    "# Works when opened from the repository root or any notebook directory.\n",
    "_start = Path(globals().get(\"__file__\", Path.cwd())).resolve()\n",
    "ROOT = next(\n",
    "    p\n",
    "    for p in [_start, *_start.parents]\n",
    "    if (p / \"pyproject.toml\").is_file() and (p / \"great_ai\").is_dir()\n",
    ")\n",
    "sys.path.insert(0, str(ROOT / \"docs\"))\n",
    "from notebook_support import setup_notebook, mag_lines\n",
    "\n",
    "FULL_RUN = os.environ.get(\"GREAT_AI_NOTEBOOK_FULL\") == \"1\"\n",
    "WORK_DIR = setup_notebook(ROOT, \"simple\")\n",
    "DATA_DIR = ROOT / \"docs/tutorial/data\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "cell-3",
   "metadata": {},
   "outputs": [],
   "source": [
    "PER_LABEL = None if FULL_RUN else 60"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "cell-4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded 420 bundled MAG sentences\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "\n",
    "records = [json.loads(line) for line in mag_lines(DATA_DIR / \"train.txt\", PER_LABEL)]\n",
    "print(f\"Loaded {len(records)} bundled MAG sentences\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-5",
   "metadata": {},
   "source": [
    "### Transform\n",
    "\n",
    "Clean the text with `great_ai.utilities.clean` and keep a list of accepted\n",
    "labels for each publication. Duplicate texts are combined before splitting\n",
    "to keep every copy of a paper in the same split."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "cell-6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from great_ai.utilities import clean\n",
    "\n",
    "labels_by_text = {}\n",
    "for record in records:\n",
    "    text = clean(record[\"text\"], convert_to_ascii=True)\n",
    "    labels_by_text.setdefault(text, set()).add(record[\"label\"])\n",
    "preprocessed_data = [(text, sorted(labels)) for text, labels in labels_by_text.items()]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "cell-7",
   "metadata": {},
   "outputs": [],
   "source": [
    "X, y = zip(*preprocessed_data)  # X is the input, y is the expected output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-8",
   "metadata": {},
   "source": [
    "### Load\n",
    "\n",
    "Store the dataset in the local TinyDB tracing database. All notebooks in\n",
    "this example share `.cache/notebooks/simple/` (or `simple-full/`). No S3 or\n",
    "MongoDB credentials are needed. A dataset tag makes rerunning this cell\n",
    "replace just this example's ground truth instead of duplicating it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "cell-9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "from great_ai import add_ground_truth, delete_ground_truth\n",
    "\n",
    "random.seed(42)\n",
    "delete_ground_truth(\"local-domain-example\")\n",
    "add_ground_truth(\n",
    "    X, y, tags=\"local-domain-example\", train_split_ratio=0.8, test_split_ratio=0.2\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-10",
   "metadata": {},
   "source": [
    "### Next: [Part 2](https://great-ai.scoutinscience.com/examples/simple/train)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv (3.13.5)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}