Dataset preparations for training a machine learning model
Before running the training of a model the following things need to be done:
- Features need to be selected: deciding which input variables/attributes will be used to make predictions
- Featuresets need to be labelled: annotating the data with the correct outputs/targets
- Dataset need to be prepared: organizing features into a structured format suitable for training (handling missing values, splitting into train/test sets, normalization, …)
This posts will give a sample on how to prepare the dataset based on the recommendation system project.
Loading the data
One datapoint can be represented by the following ArticleFeaturesSet class:
class ArticleFeatureSet:
def __init__(self,
article: Article,
label = 0):
self.content_embedding = article.content_embedding
self.title_embedding = article.title_embedding
self.url_embedding = article.uri_embedding
self.label = label
It contains the embeddings and the label which is 0 or 1 depending on type of representation.
Article is just the representation of the database object.
The FeatureLoader is responsible for collecting the instances of the ArticleFeatureSet:
class FeatureLoader:
def __init__(self):
self.db_manager = DatabaseManager()
async def extract_features(self) -> tuple[list[ArticleFeatureSet], list[ArticleFeatureSet]]:
articles: list[Article] = await self.db_manager.get_all_articles()
clicked_articles: list[ArticleFeatureSet] = []
favoured_articles: list[ArticleFeatureSet] = []
for article in articles:
clicked_articles.append(ArticleFeatureSet(article, article.is_clicked))
if article.is_clicked or article.is_favoured:
favoured_articles.append(ArticleFeatureSet(article, article.is_favoured))
print("Clicked len", len(clicked_articles))
print("Favoured len", len(favoured_articles))
return clicked_articles, favoured_articles
As there are two models to be trained the data from the table needs to be split up.
All articles are part of the dataset for clicked articles, because the article has been either clicked or not.
The dataset for the favoured articles contains only articles which have been interacted with by either clicking on it or favouring it.
Using the data
The data itself will be used as an input for the model training.
For this it needs to be split into the train and validation size.
The following method instantiates the datasets and returns them:
from sklearn.model_selection import train_test_split
def instantiate_datasets(articles: list[ArticleFeatureSet]):
binary_labels = [1 if a.label else 0 for a in articles]
train, validation = train_test_split(
articles,
test_size=0.2,
random_state=40,
stratify=binary_labels if len(set(binary_labels)) > 1 else None
)
train_dataset = ArticleDataset(train)
validation_dataset = ArticleDataset(validation)
return train_dataset, validation_dataset
train_test_split from the scikit-learn package does the trick here: It takes the list of ArticleFeatureSet and splits it by a 80/20 ratio.
The random state is used
The stratify parameter shall ensure that the test and validation sets don’t have a class imbalance. However, this works only for the simple case and is practically not as true as wished. This will be a topic for a later stage.
Using PyTorch Dataset
The return values of the instantiate_datasets methods are Dataset instances of PyTorch:
import torch
import numpy as np
from torch.utils.data import Dataset
class ArticleDataset(Dataset):
def __init__(self, articles: list[ArticleFeatureSet], normalize: str = 'l2'):
self.articles = [a for a in articles if a.content_embedding is not None
and a.title_embedding is not None
and a.url_embedding is not None] # 1
self.normalize = normalize
self.vectors, self.labels = self._build_vectors()
def _build_vectors(self):
vectors = []
labels = []
for a in self.articles:
v = np.concatenate([
np.nan_to_num(np.asarray(a.url_embedding, dtype=np.float32).flatten()),
np.nan_to_num(np.asarray(a.title_embedding, dtype=np.float32).flatten()),
np.nan_to_num(np.asarray(a.content_embedding, dtype=np.float32).flatten()),
]) # 2
if self.normalize == "l2":
v = v / (np.linalg.norm(v) + 1e-9) # 3
vectors.append(torch.from_numpy(v).float()) # 4
labels.append(torch.tensor([float(a.label)], dtype=torch.float32)) # 4
return vectors, labels
def __len__(self):
return len(self.articles)
def __getitem__(self, idx):
return self.vectors[idx], self.labels[idx]
During initialization the data is normalized:
- Articles are filterd: Only those with all three embeddings are kept.
- Feature vectors are build: The three embeddings are concatenated into one long vector. NaN values are replaced with zeros.
- Normalization: Each vector is scaled to unit length (L2 norm) so no single article has outsized influence due to magnitude.
- Everything is stored as PyTorch tensors.
The purpose is to use the Dataset instances for training as input values. This will be part of the next step.
Conclusion
This post went from using a list of ArticleFeatureSet representations to training and validation DataSets for the PyTorch model.
The split happens by using train_test_split from Pythons scikit-learn package.