{
 "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_count": null,
   "source": [
    "from docs.notebook_support import DATA_DIR, FULL_RUN, mag_lines"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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": null,
   "metadata": {},
   "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": null,
   "metadata": {},
   "outputs": [],
   "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": null,
   "metadata": {},
   "outputs": [],
   "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.pipeline import make_pipeline\n",
    "from sklearn.feature_extraction.text import TfidfVectorizer\n",
    "from sklearn.svm import LinearSVC\n",
    "\n",
    "model = make_pipeline(TfidfVectorizer(), LinearSVC(random_state=42))\n",
    "# todo: hyperparameter-optimisation\n",
    "model.fit(X_train, y_train)"
   ],
   "id": "cell-11"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It's already finished. Let's see how well it performs. But first, we need to [configure matplotlib](https://matplotlib.org/stable/tutorials/introductory/customizing.html) to make the charts more legible. \n",
    "> `ConfusionMatrixDisplay` relies on `matplotlib` for rendering."
   ],
   "id": "cell-12"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "%matplotlib inline\n",
    "plt.rcParams[\"figure.figsize\"] = (30, 15)\n",
    "plt.rcParams[\"figure.facecolor\"] = \"white\"\n",
    "plt.rcParams[\"font.size\"] = 12"
   ],
   "id": "cell-13"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we check the quality of the model on the `test` split. We can use the classification report to check the common metrics, such as the macro F1-score. We also draw the confusion matrix to get some insights into the types of mistakes being made by the model."
   ],
   "id": "cell-14"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn import metrics\n",
    "\n",
    "y_predicted = model.predict(X_test)\n",
    "\n",
    "print(metrics.classification_report(y_test, y_predicted))\n",
    "metrics.ConfusionMatrixDisplay.from_predictions(y_test, y_predicted)\n",
    "None"
   ],
   "id": "cell-15"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The recorded run achieves a test macro F1-score of about **0.69** on the MAG sentence-classification dataset, above the historical SciBERT baseline of **0.6571** associated with the [2019 SciBERT work](https://arxiv.org/abs/1903.10676). This compares a linear SVM with one historical baseline; it does not establish current state-of-the-art performance. A broader claim would require a current benchmark with comparable data splits, preprocessing, and metrics.\n",
    "\n",
    "You might wonder that *\"this is great, but besides some utility functions (`clean`, `simple_parallel_map`, ...) what more value does GreatAI add?\"*. This would be a valid argument because the scope of GreatAI actually only starts here.\n",
    "\n",
    "> Not coincidentally, this is the point where the scope of Data Science ends but it's still a grey zone for software engineering.\n",
    "\n",
    "In order to use this model in production, we have to make it available on some possibly shared infrastructure."
   ],
   "id": "cell-16"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from great_ai import save_model\n",
    "\n",
    "save_model(model, key=\"my-domain-predictor\")"
   ],
   "id": "cell-17"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### [Let's continue by finishing the deployment in the next notebook.](https://great-ai.schmelczer.dev/tutorial/deploy)"
   ],
   "id": "cell-18"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "GreatAI (local)",
   "language": "python",
   "name": "great-ai-local"
  },
  "language_info": {
   "name": "python",
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}