Optimise and train a model¶

The blue boxes show the steps implemented in this notebook.

Run data.ipynb first. We read its ground truth from the shared local TinyDB database. The default uses a compact parameter search suitable for the bundled sample; GREAT_AI_NOTEBOOK_FULL=1 enables the original larger search on the full MAG training data.

In this part, we hyperparameter-optimise and train a simple Naive Bayes classifier which we then export for deployment using great_ai.save_model.

Load data that has been extracted in part 1¶

Local setup¶

Use Python 3.12 or 3.13 and install this checkout with uv pip install --editable '.[notebooks]'. See Run notebooks locally in the repository README for the kernel, execution order, and full-data options. Run related notebooks from the same working directory to share the default .cache/ model store and tracing_database.json database.

In [1]:

from docs.notebook_support import FULL_RUN

from docs.notebook_support import FULL_RUN

In [2]:

from great_ai import query_ground_truth

data = query_ground_truth(["train", "local-domain-example"])
if not data:
    raise ValueError("Run data.ipynb first in the same local/full mode.")

X = [d.input for d in data for domain in d.feedback]
y = [domain for d in data for domain in d.feedback]

from great_ai import query_ground_truth data = query_ground_truth(["train", "local-domain-example"]) if not data: raise ValueError("Run data.ipynb first in the same local/full mode.") X = [d.input for d in data for domain in d.feedback] y = [domain for d in data for domain in d.feedback]

Environment variable ENVIRONMENT is not set, defaulting to development mode ‼️

Cannot find credentials files, defaulting to using ParallelTinyDbDriver

The selected tracing database (ParallelTinyDbDriver) is not recommended for production

Cannot find credentials files, defaulting to using LargeFileLocal

GreatAI (v0.1.11): configured ✅

  🔩 tracing_database: ParallelTinyDbDriver

  🔩 large_file_implementation: LargeFileLocal

  🔩 is_production: False

  🔩 should_log_exception_stack: True

  🔩 prediction_cache_size: 512

  🔩 dashboard_table_size: 50

You still need to check whether you follow all best practices before trusting your deployment.

> Find out more at https://se-ml.github.io/practices

In [3]:

from collections import Counter
import matplotlib.pyplot as plt

domains, counts = zip(*Counter(y).most_common())

# Configure matplotlib to have nice, high-resolution charts
%matplotlib inline
plt.rcParams["figure.figsize"] = (20, 5)
plt.rcParams["figure.facecolor"] = "white"
plt.rcParams["font.size"] = 12

plt.xticks(rotation=90)
plt.bar(domains, counts)
None

from collections import Counter import matplotlib.pyplot as plt domains, counts = zip(*Counter(y).most_common()) # Configure matplotlib to have nice, high-resolution charts %matplotlib inline plt.rcParams["figure.figsize"] = (20, 5) plt.rcParams["figure.facecolor"] = "white" plt.rcParams["font.size"] = 12 plt.xticks(rotation=90) plt.bar(domains, counts) None

Optimise and train Multinomial Naive Bayes classifier¶

Each paper may have several accepted labels. We expand it into one training row per label for this single-label classifier, but keep all rows with the same paper text in one cross-validation fold using GroupKFold. The final train/test split was already made by paper in Part 1, before expansion.

The search uses macro F1 on the expanded single-label validation rows as a training surrogate. This differs from the top-1 hit rate and fixed-target multilabel metrics used to evaluate the deployed label sets in Part 3. Rerun the search to obtain results for these grouped folds; the previous ungrouped outputs have been removed.

In [4]:

from sklearn.naive_bayes import MultinomialNB
from sklearn.pipeline import Pipeline
from sklearn.feature_extraction.text import TfidfVectorizer


def create_pipeline() -> Pipeline:
    return Pipeline(
        steps=[
            ("vectorizer", TfidfVectorizer(sublinear_tf=True)),
            ("classifier", MultinomialNB()),
        ]
    )

from sklearn.naive_bayes import MultinomialNB from sklearn.pipeline import Pipeline from sklearn.feature_extraction.text import TfidfVectorizer def create_pipeline() -> Pipeline: return Pipeline( steps=[ ("vectorizer", TfidfVectorizer(sublinear_tf=True)), ("classifier", MultinomialNB()), ] )

In [5]:

from sklearn.model_selection import GridSearchCV, GroupKFold
import pandas as pd

optimisation_pipeline = GridSearchCV(
    create_pipeline(),
    {
        "vectorizer__min_df": ([5, 20, 100] if FULL_RUN else [1, 3]),
        "vectorizer__max_df": ([0.05, 0.1] if FULL_RUN else [0.9]),
        "classifier__alpha": [0.5, 1],
        "classifier__fit_prior": [True, False],
    },
    scoring="f1_macro",
    cv=GroupKFold(n_splits=3),
    n_jobs=2,
    verbose=1,
)
# Group by paper text so that expanded labels and duplicate texts stay together.
optimisation_pipeline.fit(X, y, groups=X)

results = pd.DataFrame(optimisation_pipeline.cv_results_)
results.sort_values("rank_test_score")

from sklearn.model_selection import GridSearchCV, GroupKFold import pandas as pd optimisation_pipeline = GridSearchCV( create_pipeline(), { "vectorizer__min_df": ([5, 20, 100] if FULL_RUN else [1, 3]), "vectorizer__max_df": ([0.05, 0.1] if FULL_RUN else [0.9]), "classifier__alpha": [0.5, 1], "classifier__fit_prior": [True, False], }, scoring="f1_macro", cv=GroupKFold(n_splits=3), n_jobs=2, verbose=1, ) # Group by paper text so that expanded labels and duplicate texts stay together. optimisation_pipeline.fit(X, y, groups=X) results = pd.DataFrame(optimisation_pipeline.cv_results_) results.sort_values("rank_test_score")

Fitting 3 folds for each of 8 candidates, totalling 24 fits

Out[5]:

	mean_fit_time	std_fit_time	mean_score_time	std_score_time	param_classifier__alpha	param_classifier__fit_prior	param_vectorizer__max_df	param_vectorizer__min_df	params	split0_test_score	split1_test_score	split2_test_score	mean_test_score	std_test_score	rank_test_score
2	0.002220	0.000173	0.001228	0.000017	0.5	False	0.9	1	{'classifier__alpha': 0.5, 'classifier__fit_pr...	0.455431	0.434571	0.423547	0.437850	0.013221	1
6	0.002214	0.000084	0.001206	0.000011	1.0	False	0.9	1	{'classifier__alpha': 1, 'classifier__fit_prio...	0.458154	0.415951	0.406097	0.426734	0.022578	2
0	0.008167	0.004108	0.001306	0.000036	0.5	True	0.9	1	{'classifier__alpha': 0.5, 'classifier__fit_pr...	0.422790	0.389784	0.352330	0.388301	0.028784	3
4	0.002201	0.000114	0.001167	0.000094	1.0	True	0.9	1	{'classifier__alpha': 1, 'classifier__fit_prio...	0.381607	0.382587	0.285005	0.349733	0.045772	4
3	0.001697	0.000039	0.001383	0.000154	0.5	False	0.9	3	{'classifier__alpha': 0.5, 'classifier__fit_pr...	0.348260	0.358595	0.329416	0.345424	0.012080	5
7	0.001766	0.000113	0.001158	0.000053	1.0	False	0.9	3	{'classifier__alpha': 1, 'classifier__fit_prio...	0.344794	0.332797	0.323581	0.333724	0.008685	6
1	0.001895	0.000174	0.001272	0.000105	0.5	True	0.9	3	{'classifier__alpha': 0.5, 'classifier__fit_pr...	0.349248	0.328047	0.277855	0.318383	0.029937	7
5	0.001901	0.000075	0.001156	0.000107	1.0	True	0.9	3	{'classifier__alpha': 1, 'classifier__fit_prio...	0.337362	0.306776	0.292058	0.312065	0.018870	8

In [6]:

from sklearn import set_config

set_config(display="diagram")

classifier = create_pipeline()
classifier.set_params(**optimisation_pipeline.best_params_)
classifier.fit(X, y)

from sklearn import set_config set_config(display="diagram") classifier = create_pipeline() classifier.set_params(**optimisation_pipeline.best_params_) classifier.fit(X, y)

Out[6]:

Pipeline(steps=[('vectorizer', TfidfVectorizer(max_df=0.9, sublinear_tf=True)),
                ('classifier', MultinomialNB(alpha=0.5, fit_prior=False))])

In a Jupyter environment, please rerun this cell to show the HTML representation or trust the notebook.
On GitHub, the HTML representation is unable to render, please try loading this page with nbviewer.org.

Export the model using GreatAI¶

In [7]:

from great_ai import save_model

save_model(classifier, key="small-domain-prediction", keep_last_n=5)

from great_ai import save_model save_model(classifier, key="small-domain-prediction", keep_last_n=5)

Fetching cached versions of small-domain-prediction

Fetching cached versions of small-domain-prediction

Fetching cached versions of small-domain-prediction

Model small-domain-prediction uploaded with version 0

Out[7]:

'small-domain-prediction:0'

Next: Part 3¶