Home Tech & ScienceArtificial Intelligence (AI)How to Combine LLM Embeddings + TF-IDF + Metadata in One Scikit-learn Pipeline

How to Combine LLM Embeddings + TF-IDF + Metadata in One Scikit-learn Pipeline

by Delarno
0 comments
How to Combine LLM Embeddings + TF-IDF + Metadata in One Scikit-learn Pipeline


In this article, you will learn how to fuse dense LLM sentence embeddings, sparse TF-IDF features, and structured metadata into a single scikit-learn pipeline for text classification.

Topics we will cover include:

  • Loading and preparing a text dataset alongside synthetic metadata features.
  • Building parallel feature pipelines for TF-IDF, LLM embeddings, and numeric metadata.
  • Fusing all feature branches with ColumnTransformer and training an end-to-end classifier.

Let’s break it down.

Combine LLM Embeddings TF-IDF Metadata Scikit-learn Pipeline

How to Combine LLM Embeddings + TF-IDF + Metadata in One Scikit-learn Pipeline (click to enlarge)
Image by Editor

Introduction

Data fusion, or combining diverse pieces of data into a single pipeline, sounds ambitious enough. If we talk not just about two, but about three complementary feature sources, then the challenge — and the potential payoff — goes to the next level. The most exciting part is that scikit-learn allows us to unify all of them cleanly within a single, end-to-end workflow. Do you want to see how? This article walks you step by step through building a complete fusion pipeline from scratch for a downstream text classification task, combining dense semantic information from LLM-generated embeddings, sparse lexical features from TF-IDF, and structured metadata signals. Interested? Keep reading.

Step-by-Step Pipeline Building Process

First, we will make all the necessary imports for the pipeline-building process. If you are working in a local environment, you might need to pip install some of them first:

Let’s look closely at this — almost endless! — list of imports. I bet one element has caught your attention: fetch_20newsgroups. This is a freely available text dataset in scikit-learn that we will use throughout this article: it contains text extracted from news articles belonging to a wide variety of categories.

To keep our dataset manageable in practice, we will pick the news articles belonging to a subset of categories specified by us. The following code does the trick:

We called this freshly created dataset X_raw to emphasize that this is a raw, far-from-final version of the dataset we will gradually construct for downstream tasks like using machine learning models for predictive purposes. It is fair to say that the “raw” suffix is also used because here we have the raw text, from which three different data components (or streams) will be generated and later merged.

For the structured metadata associated with the news articles obtained, in real-world contexts, this metadata might already be available or provided by the dataset owner. That’s not the case with this publicly available dataset, so we will synthetically create some simple metadata features based on the text, including features describing character length, word count, average word length, uppercase ratio, and digit ratio.

Before getting fully into the pipeline-building process, we will split the data into train and test subsets:

Very important: splitting the data into training and test sets must be done before extracting the LLM embeddings and TF-IDF features. Why? Because these two extraction processes become part of the pipeline, and they involve fitting transformations with scikit-learn, which are learning processes — for example, learning the TF-IDF vocabulary and inverse document frequency (IDF) statistics. The scikit-learn logic to implement this is as follows: any data transformations must be fitted (learn the transformation logic) only on the training data and then applied to the test data using the learned logic. This way, no information from the test set will influence or bias feature construction or downstream model training.

Now comes a key stage: defining a class that encapsulates a pre-trained sentence transformer (a language model like all-MiniLM-L6-v2 capable of generating text embeddings from raw text) to produce our custom LLM embeddings.

Now we are building the three main data branches (or parallel pipelines) we are interested in, one by one. First, the pipeline for TF-IDF feature extraction, in which we will use scikit-learn’s TfidfVectorizer class to extract these features seamlessly:

Next comes the LLM embeddings pipeline, aided by the custom class we defined earlier:

Last, we define the branch pipeline for the metadata features, in which we aim to standardize these attributes due to their disparate ranges:

Now we have three parallel pipelines, but nothing to connect them — at least not yet. Here comes the main, overarching pipeline that will orchestrate the fusion process among all three data branches, by using a very useful and versatile scikit-learn artifact for the fusion of heterogeneous data flows: a ColumnTransformer pipeline.

And the icing on the cake: a full, end-to-end pipeline that will combine the fusion pipeline with an example of a machine learning-driven downstream task. In particular, here’s how to combine the entire data fusion pipeline we have just architected with the training of a logistic regression classifier to predict the news category:

The following instruction will do all the heavy lifting we have been designing so far. The LLM embeddings part will particularly take a few minutes (especially if the model needs to be downloaded), so be patient. This step will undertake the whole threefold process of data preprocessing, fusion, and model training:

To finalize, we can make predictions on the test set and see how our fusion-driven classifier performs.

And for a visual wrap-up, here’s what the entire pipeline looks like:

Text data fusion pipeline with scikit-learn

Wrapping Up

This article guided you through the process of building an entire machine learning-oriented workflow that focuses on the fusion of several information sources derived from raw text data, so that everything can be put together in downstream predictive tasks like text classification. We have seen how scikit-learn provides a set of useful classes and methods to make the process easier and more intuitive.



Source link

You may also like

Leave a Comment