---
title: "XGBoost Training"
description: "Part 3: Learn how to train XGBoost models, make predictions, and
evaluate model performance using xorq's deferred pipelines"
---

## Overview

In this tutorial (Part 3 of our series), you'll learn how to:
- Define deferred model training and prediction operations
- Split data into train and test sets
- Train an XGBoost model with TF-IDF
- Make predictions on both training and test data
- Evaluate model performance 

### Prerequisites

<Accordion title="Prerequisites">
- Completed Part 1 (Data Ingestion and Model-Assisted Labeling)
- Completed Part 2 (Feature Engineering with TF-IDF)
- Python 3.8+ installed on your system
- Basic understanding of machine learning concepts
</Accordion>

### Installation and Imports

First, ensure you have the required packages:

```bash
pip install xorq pandas scikit-learn xgboost
```

Then import the necessary modules:

```python
import pandas as pd
import toolz
import xgboost as xgb
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics import mean_absolute_error, confusion_matrix

import xorq as xo
import xorq.vendor.ibis.expr.datatypes as dt
from xorq.caching import ParquetStorage
from xorq.common.utils.defer_utils import deferred_read_parquet
from xorq.common.utils.import_utils import import_python
from xorq.expr.ml import (
    deferred_fit_predict,
    deferred_fit_transform_series_sklearn,
    train_test_splits,
)

# Import the helper modules we used in previous parts
m = import_python(xo.options.pins.get_path("hackernews_lib"))
o = import_python(xo.options.pins.get_path("openai_lib"))
```

## Setting Up Deferred Operations
### Defining Model Training and Prediction Functions

Let's define functions for training and making predictions with XGBoost:

```python
@toolz.curry
def fit_xgboost_model(feature_df, target_series, seed=0):
    xgb_r = xgb.XGBRegressor(
        eval_metric=mean_absolute_error,
        max_depth=6,
        seed=seed,
    )
    X = pd.DataFrame(feature_df.squeeze().tolist())
    xgb_r.fit(X, target_series)
    return xgb_r

@toolz.curry
def predict_xgboost_model(model, df):
    return model.predict(df.squeeze().tolist())
```

<Info>
The `fit_xgboost_model` function trains an XGBoost model on the provided
features and target. The `predict_xgboost_model` function applies the trained
model to new data to generate predictions.

Note that we're using `multi:softmax` as the objective function since we have
three sentiment classes (POSITIVE=2, NEUTRAL=1, NEGATIVE=0).
</Info>


Now, let's set up our deferred operations for both the TF-IDF transformation
and XGBoost prediction:

```python
# Define column names
transform_col = "title"
features = (transformed_col,) = (f"{transform_col}_transformed",)
target = "sentiment_int"
target_predicted = f"{target}_predicted"

# Create a deferred TF-IDF transformer (same as in Part 2)
deferred_fit_transform_tfidf = deferred_fit_transform_series_sklearn(
    col=transform_col,
    cls=TfidfVectorizer,
    return_type=dt.Array(dt.float64),
)

# Create a deferred XGBoost model
deferred_fit_predict_xgb = deferred_fit_predict(
    target=target,
    features=list(features),
    fit=fit_xgboost_model,
    predict=predict_xgboost_model,
    return_type=dt.float32,
)
```

<Note>
The `deferred_fit_predict` function creates a deferred operation that will:
1. Fit a model using the specified `fit` function on the training data
2. Create a prediction operation that can be applied to any dataset

Unlike the TF-IDF transformation (which we covered in detail in Part 2), model
training is implemented as an aggregate function rather than a UDXF function.
This is because training involves aggregating across the entire dataset to
learn patterns, while transformation is applied row by row.
</Note>

## Loading and Preparing the Data

Let's load and prepare our data, similar to what we did in the previous parts:

```python
# Initialize the backend
con = xo.connect()
storage = ParquetStorage(source=con)

# Define the input dataset name
name = "hn-fetcher-input-large"

# Load and process the data (similar to Parts 1 and 2)
raw_expr = (
    deferred_read_parquet(
        con,
        xo.options.pins.get_path(name),
        name,
    )
    .pipe(m.do_hackernews_fetcher_udxf)
)

t = (
    raw_expr
    .filter(xo._.text.notnull())
    .pipe(o.do_hackernews_sentiment_udxf, con=con)
    .cache(storage=ParquetStorage(con))
    .filter(~xo._.sentiment.contains("ERROR"))
    .mutate(
        sentiment_int=xo._.sentiment.cases(
            {"POSITIVE": 2, "NEUTRAL": 1, "NEGATIVE": 0}.items()
        ).cast(int)
    )
)
```

### Splitting the Data into Train and Test Sets

Before training our model, we'll split the data into training and testing sets:

```python
# Split into train (60%) and test (40%) sets
(train_expr, test_expr) = t.pipe(
    train_test_splits,
    unique_key="id",
    test_sizes=(0.6, 0.4),
    random_seed=42,
)
```

<Info>
The `train_test_splits` function in xorq ensures a proper and deterministic
split of your data. It works by using a hashing function to convert the unique
key (`id` in our case) into an integer, then applies a modulo operation to
split the data into buckets.

Having a unique key field is essential as it allows xorq to deterministically
order the table and assign records to either the training or test set. This
approach ensures that:
1. The same record will always end up in the same split when using the same random seed
2. The splitting is distributed evenly across the dataset
3. Records are not duplicated across splits
</Info>

### Applying TF-IDF Transformation

Let's apply the TF-IDF transformation to our training data:

```python
# Create the deferred TF-IDF transformation
(deferred_tfidf_model, tfidf_udaf, deferred_tfidf_transform) = (
    deferred_fit_transform_tfidf(
        train_expr,
        storage=storage,
    )
)

# Apply the transformation to the training data
train_tfidf_transformed = train_expr.mutate(
    **{transformed_col: deferred_tfidf_transform.on_expr}
)
```

<Note>
We're using the same TF-IDF approach we explored in [Part
2](/tutorials/hn_transform_tfidf), fitting on the
training data to create a vocabulary and then transforming the documents into
numerical feature vectors. This step is necessary to convert the text into a
format that our XGBoost model can process.
</Note>

## Training the XGBoost Model

Now, let's train our XGBoost model on the transformed training data:

```python
# Create the deferred XGBoost model and prediction operation
(deferred_xgb_model, xgb_udaf, deferred_xgb_predict) = deferred_fit_predict_xgb(
    train_tfidf_transformed,
    storage=storage,
)

# Apply predictions to the training data
train_xgb_predicted = (
    train_tfidf_transformed
    .mutate(**{target_predicted: deferred_xgb_predict.on_expr})
)
```
<Info>
Unlike the transformation step, model training is implemented as an aggregate
function (`xgb_udaf`). This is an important distinction:

1. **Transformation (UDF)**: Operates row by row, applying the same function to
each record independently
2. **Training (UDAF)**: Aggregates across the entire dataset, learning patterns
from all records collectively

The `deferred_fit_predict_xgb` function returns three key components:
- `deferred_xgb_model`: an Expr that returns a trained model. 
- `xgb_udaf`: The User-Defined Aggregate Function that performs the training
- `deferred_xgb_predict`: The scalar UDF that takes Expr as an input i.e.
  `ExprScalarUDF`
</Info>

### Making Predictions on Test Data

Similarly, we'll apply both the TF-IDF transformation and XGBoost prediction to
our test data:

```python
# Apply TF-IDF transformation and XGBoost prediction to test data
train_xgb_predicted = (
    train_tfidf_transformed
    .into_backend(xo.connect()) # extra into backend (see warning below)
    .mutate(**{target_predicted: deferred_xgb_predict.on_expr})
)
```
<Warning>
Note the use of superfluous `.into_backend(xo.connect())`. This is necessary to ensure
proper handling of the data types during the prediction process and should be
fixed. See the GitHub issue for more information.
</Warning>

## Evaluating Model Performance

Let's execute our pipeline and evaluate the model's performance:

```python
# Execute the training and test predictions
train_results = train_xgb_predicted.execute()
test_results = test_xgb_predicted.execute()

# Print model performance statistics by sentiment class
print("Training Set Performance:")
print(train_results.groupby("sentiment_int").sentiment_int_predicted.describe().T)

print("\nTest Set Performance:")
print(test_results.groupby("sentiment_int").sentiment_int_predicted.describe().T)

# Calculate overall accuracy
train_accuracy = (train_results[target_predicted] == train_results[target]).mean()
test_accuracy = (test_results[target_predicted] == test_results[target]).mean()

print(f"\nTraining Accuracy: {train_accuracy:.4f}")
print(f"Test Accuracy: {test_accuracy:.4f}")
```

## Summary and Next Steps

Congratulations! In this third part of our tutorial series, you've:
1. Created deferred operations for model training and prediction
2. Split data into training and testing sets
3. Applied TF-IDF transformation to convert text to features
4. Trained an XGBoost model for sentiment classification
5. Made predictions on both training and test data
6. Evaluated model performance using various metrics
7. Applied the model to make predictions on new data

In the next tutorial (Part 4), we'll explore how to deploy our trained model
for real-time predictions using xorq's Flight serving capabilities.

## Further Reading

- [XGBoost Documentation](https://xgboost.readthedocs.io/)
- [xorq Documentation](https://docs.xorq.io)
- [Model Evaluation Metrics](https://scikit-learn.org/stable/modules/model_evaluation.html)

## Appendix

### Complete Example Code

<Accordion title="Complete Code">
```python
import pandas as pd
import toolz
import xgboost as xgb
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics import mean_absolute_error, confusion_matrix

import xorq as xo
import xorq.vendor.ibis.expr.datatypes as dt
from xorq.caching import ParquetStorage
from xorq.common.utils.defer_utils import deferred_read_parquet
from xorq.common.utils.import_utils import import_python
from xorq.expr.ml import (
    deferred_fit_predict,
    deferred_fit_transform_series_sklearn,
    train_test_splits,
)

# Import the helper modules
m = import_python(xo.options.pins.get_path("hackernews_lib"))
o = import_python(xo.options.pins.get_path("openai_lib"))

# Define model training and prediction functions
@toolz.curry
def fit_xgboost_model(feature_df, target_series, seed=0):
    xgb_r = xgb.XGBRegressor(
        objective="multi:softmax",
        num_class=3,
        eval_metric=mean_absolute_error,
        max_depth=6,
        n_estimators=10,
        seed=seed,
    )
    X = pd.DataFrame(feature_df.squeeze().tolist())
    xgb_r.fit(X, target_series)
    return xgb_r

@toolz.curry
def predict_xgboost_model(model, df):
    return model.predict(df.squeeze().tolist())

# Define column names
transform_col = "title"
features = (transformed_col,) = (f"{transform_col}_transformed",)
target = "sentiment_int"
target_predicted = f"{target}_predicted"

# Create deferred operations
deferred_fit_transform_tfidf = deferred_fit_transform_series_sklearn(
    col=transform_col,
    cls=TfidfVectorizer,
    return_type=dt.Array(dt.float64),
)

deferred_fit_predict_xgb = deferred_fit_predict(
    target=target,
    features=list(features),
    fit=fit_xgboost_model,
    predict=predict_xgboost_model,
    return_type=dt.float32,
)

# Initialize the backend
con = xo.connect()
storage = ParquetStorage(source=con)

# Load and process data
name = "hn-fetcher-input-large"
raw_expr = (
    deferred_read_parquet(
        con,
        xo.options.pins.get_path(name),
        name,
    )
    .pipe(m.do_hackernews_fetcher_udxf)
)

t = (
    raw_expr
    .filter(xo._.text.notnull())
    .pipe(o.do_hackernews_sentiment_udxf, con=con)
    .cache(storage=ParquetStorage(con))
    .filter(~xo._.sentiment.contains("ERROR"))
    .mutate(
        sentiment_int=xo._.sentiment.cases(
            {"POSITIVE": 2, "NEUTRAL": 1, "NEGATIVE": 0}.items()
        ).cast(int)
    )
)

# Split into train and test sets
(train_expr, test_expr) = t.pipe(
    train_test_splits,
    unique_key="id",
    test_sizes=(0.6, 0.4),
    random_seed=42,
)

# Apply TF-IDF transformation
(deferred_tfidf_model, tfidf_udaf, deferred_tfidf_transform) = (
    deferred_fit_transform_tfidf(
        train_expr,
        storage=storage,
    )
)

train_tfidf_transformed = train_expr.mutate(
    **{transformed_col: deferred_tfidf_transform.on_expr}
)

# Train XGBoost model
(deferred_xgb_model, xgb_udaf, deferred_xgb_predict) = deferred_fit_predict_xgb(
    train_tfidf_transformed,
    storage=storage,
)

train_xgb_predicted = (
    train_tfidf_transformed
    .into_backend(xo.connect())
    .mutate(**{target_predicted: deferred_xgb_predict.on_expr})
)

# Make predictions on test data
test_xgb_predicted = (
    test_expr
    .mutate(**{transformed_col: deferred_tfidf_transform.on_expr})
    .into_backend(xo.connect(), name="stable-name")
    .mutate(**{target_predicted: deferred_xgb_predict.on_expr})
)

# Execute and evaluate
train_results = train_xgb_predicted.execute()
test_results = test_xgb_predicted.execute()

# Calculate overall accuracy
train_accuracy = (train_results[target_predicted] == train_results[target]).mean()
test_accuracy = (test_results[target_predicted] == test_results[target]).mean()

print(f"Training Accuracy: {train_accuracy:.4f}")
print(f"Test Accuracy: {test_accuracy:.4f}")
```
</Accordion>
