{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Train your model\n", "\n", "Install the checkout's notebook dependencies first, following the README's\n", "**Run notebooks locally** instructions. Select the `GreatAI (local)` kernel.\n", "The default run uses a balanced subset of the bundled MAG data for speed.\n", "Set `GREAT_AI_NOTEBOOK_FULL=1` before starting Jupyter to use the full splits.\n", "Local subset results differ from the historical full-dataset results below." ], "id": "cell-0" }, { "id": "37cb23a4", "cell_type": "markdown", "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. Run related notebooks from the same working directory\n", "to share the default `.cache/` model store and `tracing_database.json` database.\n" ], "metadata": {} }, { "id": "42ed923c", "cell_type": "code", "metadata": { "execution": { "iopub.execute_input": "2026-10-28T22:47:18.653947Z", "iopub.status.busy": "2026-10-28T22:47:18.653802Z", "iopub.status.idle": "2026-10-28T22:47:18.664778Z", "shell.execute_reply": "2026-10-28T22:47:18.663856Z" } }, "execution_count": 1, "source": [ "from docs.notebook_support import DATA_DIR, FULL_RUN, mag_lines" ], "outputs": [] }, { "cell_type": "code", "execution_count": 2, "metadata": { "execution": { "iopub.execute_input": "2026-10-28T22:47:18.667022Z", "iopub.status.busy": "2026-10-28T22:47:18.666820Z", "iopub.status.idle": "2026-10-28T22:47:18.669716Z", "shell.execute_reply": "2026-10-28T22:47:18.669055Z" } }, "outputs": [], "source": [ "# Dependencies are installed once in the environment; no in-kernel installation needed." ], "id": "cell-3" }, { "cell_type": "markdown", "metadata": {}, "source": [ "First, we have to get some data. After downloading it [from here](https://github.com/allenai/scibert/tree/master/data/text_classification/mag), we might notice that the dataset is in [JSON Lines](https://jsonlines.org/) format (each line is a separate JSON document). \n", "\n", "Let's write a function which takes a single line and returns the sentence and the corresponding label from it. Before returning, the sentence is also [cleaned](https://great-ai.schmelczer.dev/reference/utilities/#great_ai.utilities.clean.clean) to remove any LaTeX, XML, unicode, PDF-extraction artifacts." ], "id": "cell-4" }, { "cell_type": "code", "execution_count": 3, "metadata": { "execution": { "iopub.execute_input": "2026-10-28T22:47:18.671242Z", "iopub.status.busy": "2026-10-28T22:47:18.671089Z", "iopub.status.idle": "2026-10-28T22:47:21.261086Z", "shell.execute_reply": "2026-10-28T22:47:21.260622Z" } }, "outputs": [], "source": [ "import json\n", "from great_ai.utilities import clean\n", "\n", "\n", "def preprocess(line):\n", " data_point = json.loads(line)\n", "\n", " sentence = data_point[\"text\"]\n", " label = data_point[\"label\"]\n", "\n", " return clean(sentence), label" ], "id": "cell-5" }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now, we can load the dataset and extract the *training* samples from it. Since we're impatient, we can do it in parallel using the [`simple_parallel_map`](https://great-ai.schmelczer.dev/reference/utilities/#great_ai.utilities.simple_parallel_map) function.\n", "\n", "> Open files in Python are iterable: in text mode, each iteration returns the next line." ], "id": "cell-6" }, { "cell_type": "code", "execution_count": 4, "metadata": { "execution": { "iopub.execute_input": "2026-10-28T22:47:21.262422Z", "iopub.status.busy": "2026-10-28T22:47:21.262265Z", "iopub.status.idle": "2026-10-28T22:47:22.482973Z", "shell.execute_reply": "2026-10-28T22:47:22.482099Z" } }, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "\r", " 0%| | 0/700 [00:00, ?it/s]" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", " 0%| | 1/700 [00:00<01:30, 7.74it/s]" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", " 10%|█ | 71/700 [00:00<00:01, 352.47it/s]" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", " 20%|██ | 141/700 [00:00<00:01, 467.83it/s]" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", " 27%|██▋ | 190/700 [00:00<00:01, 475.11it/s]" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", " 34%|███▍ | 239/700 [00:00<00:00, 461.41it/s]" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", " 45%|████▌ | 316/700 [00:00<00:00, 538.41it/s]" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", " 53%|█████▎ | 371/700 [00:00<00:00, 523.75it/s]" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", " 61%|██████ | 424/700 [00:00<00:00, 504.26it/s]" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", " 80%|████████ | 561/700 [00:01<00:00, 697.72it/s]" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", " 90%|█████████ | 631/700 [00:01<00:00, 675.11it/s]" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", "100%|██████████| 700/700 [00:01<00:00, 603.58it/s]" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\n" ] } ], "source": [ "from great_ai.utilities import simple_parallel_map\n", "\n", "training_lines = mag_lines(DATA_DIR / \"train.txt\", None if FULL_RUN else 100)\n", "training_data = simple_parallel_map(preprocess, training_lines, concurrency=2)\n", "X_train = [d[0] for d in training_data]\n", "y_train = [d[1] for d in training_data]" ], "id": "cell-7" }, { "cell_type": "markdown", "metadata": {}, "source": [ "Let's do the same for the *test* data." ], "id": "cell-8" }, { "cell_type": "code", "execution_count": 5, "metadata": { "execution": { "iopub.execute_input": "2026-10-28T22:47:22.484582Z", "iopub.status.busy": "2026-10-28T22:47:22.484426Z", "iopub.status.idle": "2026-10-28T22:47:23.612412Z", "shell.execute_reply": "2026-10-28T22:47:23.611412Z" } }, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "\r", " 0%| | 0/210 [00:00, ?it/s]" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", " 0%| | 1/210 [00:00<00:27, 7.72it/s]" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", " 6%|▌ | 12/210 [00:00<00:03, 61.14it/s]" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", " 16%|█▌ | 34/210 [00:00<00:01, 124.86it/s]" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", " 27%|██▋ | 56/210 [00:00<00:00, 156.07it/s]" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", " 37%|███▋ | 78/210 [00:00<00:00, 175.16it/s]" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", " 46%|████▌ | 96/210 [00:00<00:00, 173.89it/s]" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", " 58%|█████▊ | 122/210 [00:00<00:00, 197.51it/s]" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", " 69%|██████▊ | 144/210 [00:00<00:00, 202.19it/s]" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", " 84%|████████▍ | 177/210 [00:00<00:00, 231.71it/s]" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", " 96%|█████████▌| 201/210 [00:01<00:00, 230.89it/s]" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", "100%|██████████| 210/210 [00:01<00:00, 191.75it/s]" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\n" ] } ], "source": [ "test_lines = mag_lines(DATA_DIR / \"test.txt\", None if FULL_RUN else 30)\n", "test_data = simple_parallel_map(preprocess, test_lines, concurrency=2)\n", "X_test = [d[0] for d in test_data]\n", "y_test = [d[1] for d in test_data]" ], "id": "cell-9" }, { "cell_type": "markdown", "metadata": {}, "source": [ "After obtaining some clean data, it's time to create and train a model on it. We are going to use [scikit-learn](https://scikit-learn.org/stable/auto_examples/classification/plot_classifier_comparison.html) for this purpose.\n", "\n", "In this case, we opted for a [linear Support Vector Machine](https://scikit-learn.org/stable/modules/generated/sklearn.svm.LinearSVC.html) because it's known to be an adequate baseline for simple classification tasks." ], "id": "cell-10" }, { "cell_type": "code", "execution_count": 6, "metadata": { "execution": { "iopub.execute_input": "2026-10-28T22:47:23.614164Z", "iopub.status.busy": "2026-10-28T22:47:23.614030Z", "iopub.status.idle": "2026-10-28T22:47:23.726757Z", "shell.execute_reply": "2026-10-28T22:47:23.726161Z" } }, "outputs": [ { "data": { "text/html": [ "
Pipeline(steps=[('tfidfvectorizer', TfidfVectorizer()),\n",
" ('linearsvc', LinearSVC(random_state=42))])In a Jupyter environment, please rerun this cell to show the HTML representation or trust the notebook. | Name | \n", "Type | \n", "Value | \n", "
|---|---|---|
| \n",
" \n",
" classes_\n",
" \n",
" classes_: ndarray of shape (n_classes,) The classes labels. Only exist if the last step of the pipeline is a classifier.\n", " \n", " | \n",
" ndarray[<U10](7,) | \n", "['business','economics','geography',...,'politics','psychology',\n", " 'sociology'] | \n", "\n", "\n", "
| Name | \n", "Type | \n", "Value | \n", "
|---|---|---|
| \n",
" \n",
" fixed_vocabulary_\n",
" \n",
" fixed_vocabulary_: bool True if a fixed vocabulary of term to indices mapping is provided by the user.\n", " \n", " | \n",
" bool | \n", "False | \n", "\n", "\n", "
| \n",
" \n",
" idf_\n",
" \n",
" idf_: array of shape (n_features,) The inverse document frequency (IDF) vector; only defined if ``use_idf`` is True.\n", " \n", " | \n",
" ndarray[float64](3636,) | \n", "[5.25,6.86,6.86,...,6.86,6.86,6.86] | \n", "\n", "\n", "
| \n",
" \n",
" vocabulary_\n",
" \n",
" vocabulary_: dict A mapping of terms to feature indices.\n", " \n", " | \n",
" dict | \n", "{'00': 0, '0864': 1, '0973': 2, '10': 3, ...} | \n", "\n", "\n", "
| 00 | \n", "
| 0864 | \n", "
| 0973 | \n", "
| 10 | \n", "
| 100 | \n", "
| 1120 | \n", "
| 12 | \n", "
| 1258 | \n", "
| 14 | \n", "
| 14614 | \n", "
| 1472454871 | \n", "
| 15 | \n", "
| 1500 | \n", "
| 1538 | \n", "
| 1544 | \n", "
| 1550 | \n", "
| 1558 | \n", "
| 16 | \n", "
| 161pp | \n", "
| 1727 | \n", "
| 1750 | \n", "
| 1769 | \n", "
| 1782 | \n", "
| 18 | \n", "
| 1810 | \n", "
| 1821 | \n", "
| 1846 | \n", "
| 1850 | \n", "
| 1870 | \n", "
| 1872 | \n", "
| 1880 | \n", "
| 1881 | \n", "
| 1890 | \n", "
| 18th | \n", "
| 19 | \n", "
| 1907 | \n", "
| 1909 | \n", "
| 1919 | \n", "
| 1920 | \n", "
| 1924 | \n", "
| 1933 | \n", "
| 1940 | \n", "
| 1950s | \n", "
| 1963 | \n", "
| 1970 | \n", "
| 1971 | \n", "
| 1972 | \n", "
| 1980 | \n", "
| 1984 | \n", "
| 1986 | \n", "
| 1988 | \n", "
| 1989 | \n", "
| 1990 | \n", "
| 1991 | \n", "
| 1992 | \n", "
| 1997 | \n", "
| 1999 | \n", "
| 1st | \n", "
| 20 | \n", "
| 2001 | \n", "
| 2002 | \n", "
| 2003 | \n", "
| 2005 | \n", "
| 2006 | \n", "
| 2007 | \n", "
| 2008 | \n", "
| 2009 | \n", "
| 200pp | \n", "
| 2011 | \n", "
| 2013 | \n", "
| 2014 | \n", "
| 2015 | \n", "
| 2016 | \n", "
| 2017 | \n", "
| 203 | \n", "
| 204 | \n", "
| 20th | \n", "
| 21 | \n", "
| 22 | \n", "
| 234 | \n", "
| 24 | \n", "
| 242 | \n", "
| 247 | \n", "
| 25 | \n", "
| 251 | \n", "
| 267 | \n", "
| 269p | \n", "
| 275 | \n", "
| 278 | \n", "
| 28 | \n", "
| 283 | \n", "
| 28332 | \n", "
| 29 | \n", "
| 295 | \n", "
| 299 | \n", "
| 30 | \n", "
| 300 | \n", "
| 302 | \n", "
| 31st | \n", "
| 338 | \n", "
| Name | \n", "Type | \n", "Value | \n", "
|---|---|---|
| \n",
" \n",
" classes_\n",
" \n",
" classes_: ndarray of shape (n_classes,) The unique classes labels.\n", " \n", " | \n",
" ndarray[<U10](7,) | \n", "['business','economics','geography',...,'politics','psychology',\n", " 'sociology'] | \n", "\n", "\n", "
| \n",
" \n",
" coef_\n",
" \n",
" coef_: ndarray of shape (1, n_features) if n_classes == 2 else (n_classes, n_features) Weights assigned to the features (coefficients in the primal problem). ``coef_`` is a readonly property derived from ``raw_coef_`` that follows the internal memory layout of liblinear.\n", " \n", " | \n",
" ndarray[float64](7, 3636) | \n", "[[ 0.18,-0.01,-0.04,...,-0.1 ,-0.1 ,-0.1 ],\n", " [-0.03,-0. ,-0. ,..., 0.58, 0.58, 0.58],\n", " [-0.13,-0.01,-0.03,...,-0.12,-0.12,-0.12],\n", " ...,\n", " [ 0.21, 0.14,-0.09,...,-0.08,-0.08,-0.08],\n", " [-0.16,-0.04,-0.06,...,-0.09,-0.09,-0.09],\n", " [ 0.13,-0.06, 0.27,...,-0.1 ,-0.1 ,-0.1 ]] | \n", "\n", "\n", "
| \n",
" \n",
" intercept_\n",
" \n",
" intercept_: ndarray of shape (1,) if n_classes == 2 else (n_classes,) Constants in decision function.\n", " \n", " | \n",
" ndarray[float64](7,) | \n", "[-0.7 ,-0.74,-0.65,...,-0.77,-0.72,-0.69] | \n", "\n", "\n", "
| \n",
" \n",
" n_features_in_\n",
" \n",
" n_features_in_: int Number of features seen during :term:`fit`. .. versionadded:: 0.24\n", " \n", " | \n",
" int | \n", "3636 | \n", "\n", "\n", "
| \n",
" \n",
" n_iter_\n",
" \n",
" n_iter_: int Maximum number of iterations run across all classes.\n", " \n", " | \n",
" int | \n", "25 | \n", "\n", "\n", "