{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Train domain classifier on the [semantic scholar dataset](https://api.semanticscholar.org/corpus)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import json\n", "import matplotlib.pyplot as plt\n", "from pathlib import Path\n", "import pandas as pd\n", "from sklearn import metrics, set_config\n", "from sklearn.model_selection import train_test_split\n", "from sklearn.naive_bayes import MultinomialNB\n", "from sklearn.pipeline import Pipeline\n", "from sklearn.feature_extraction.text import TfidfVectorizer\n", "from sklearn.model_selection import GridSearchCV\n", "\n", "from great_ai.utilities.parallel_map import parallel_map\n", "from great_ai.utilities.language import is_english, predict_language\n", "from great_ai import save_model, configure, LargeFile\n", "\n", "from helpers import preprocess, lemmatize" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Configuration" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "PREFIX = \"domain-\"\n", "DATASET_KEY = \"data\"\n", "MAX_FILE_COUNT = 5\n", "MODEL_KEY = \"small-domain-prediction-v2\"" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "configure()\n", "corpus_path = LargeFile(DATASET_KEY).get()\n", "\n", "set_config(display=\"diagram\")\n", "plt.rcParams[\"figure.figsize\"] = (30, 15)\n", "plt.rcParams[\"figure.facecolor\"] = \"white\"\n", "plt.rcParams[\"font.size\"] = 12\n", "plt.rcParams[\"axes.xmargin\"] = 0" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Preprocessing" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def clean_file(p: Path) -> None:\n", " try:\n", " processed_path = p.with_name(f\"{PREFIX}{p.stem}{p.suffix}\")\n", "\n", " if processed_path.exists():\n", " return\n", "\n", " with open(p) as f:\n", " content = json.load(f)\n", "\n", " result = {\n", " lemmatize(preprocess(f'{c[\"title\"]} {c[\"abstract\"]}')): c[\"domain\"]\n", " for c in content\n", " if (\n", " c[\"domain\"]\n", " and c[\"abstract\"]\n", " and is_english(predict_language(c[\"abstract\"]))\n", " )\n", " }\n", "\n", " with open(processed_path, \"w\") as f:\n", " json.dump(result, f)\n", " except Exception as e:\n", " print(f\"Error ({e}) processing {p}\")\n", "\n", "\n", "parallel_map(\n", " clean_file,\n", " list(corpus_path.glob(\"s2-corpus-*.json\"))[:MAX_FILE_COUNT],\n", " chunk_size=1,\n", ")\n", "None" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "corpora = list(corpus_path.glob(f\"{PREFIX}*.json\"))[:MAX_FILE_COUNT]\n", "print(f\"Found {len(corpora)} files\")\n", "\n", "data = []\n", "for p in corpora:\n", " with open(p) as f:\n", " data.extend(json.load(f).items())\n", "\n", "print(f\"Found {len(data)} documents\")\n", "\n", "X_train, X_test, y_train, y_test = train_test_split(\n", " [d[0] for d in data], [d[1] for d in data], test_size=0.1, random_state=1\n", ")\n", "\n", "X_train = [x for x, y in zip(X_train, y_train) for domain in y]\n", "y_train = [domain for x, y in zip(X_train, y_train) for domain in y]" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Optimise and train Multinomial Naive Bayes classifier" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def create_pipeline() -> Pipeline:\n", " return Pipeline(\n", " steps=[\n", " (\"vectorizer\", TfidfVectorizer()),\n", " (\"classifier\", MultinomialNB()),\n", " ]\n", " )" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "optimisation_pipeline = GridSearchCV(\n", " create_pipeline(),\n", " {\n", " \"vectorizer__max_df\": [0.05, 0.1],\n", " \"vectorizer__min_df\": [5, 20],\n", " \"vectorizer__sublinear_tf\": [True, False],\n", " \"classifier__alpha\": [0.5, 1],\n", " \"classifier__fit_prior\": [True, False],\n", " },\n", " scoring=\"f1_macro\",\n", " cv=3,\n", " n_jobs=8,\n", " verbose=1,\n", ")\n", "optimisation_pipeline.fit(X_train, y_train)\n", "\n", "results = pd.DataFrame(optimisation_pipeline.cv_results_)\n", "results.sort_values(\"rank_test_score\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "classifier = create_pipeline()\n", "classifier.set_params(**optimisation_pipeline.best_params_)\n", "classifier.fit(X_train, y_train)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Check accuracy on the test split" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "predicted = classifier.predict(X_test)\n", "\n", "y_test_aligned = [p if p in y else y[0] for p, y in zip(predicted, y_test)]\n", "\n", "print(metrics.classification_report(y_test_aligned, predicted))\n", "metrics.ConfusionMatrixDisplay.from_predictions(\n", " y_true=y_test_aligned,\n", " y_pred=predicted,\n", " xticks_rotation=\"vertical\",\n", " normalize=\"pred\",\n", " values_format=\".2f\",\n", ")\n", "None" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Export the model using GreatAI" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "save_model(classifier, key=MODEL_KEY, keep_last_n=5)" ] } ], "metadata": { "interpreter": { "hash": "acc70e949538f42041ccc57bc4df2261507e3fd7d6b9ce5dcc28e3bcf9d48274" }, "kernelspec": { "display_name": "Python 3.8.5 ('.env': venv)", "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.8.5" }, "orig_nbformat": 4 }, "nbformat": 4, "nbformat_minor": 2 }