Scikit-LLM and BART: What Future for Text Classification?

The Rise of LLMs in Text Classification

Large Language Models (LLMs) have begun to transform the landscape of text classification, replacing traditional machine learning approaches. However, the question of when to opt for an LLM or stick with classic models remains crucial for developers. The trade-offs include the speed and reliability of conventional models versus the advanced reasoning potential of LLMs.

In this analysis, we examine three approaches to text classification:

• TF-IDF combined with logistic regression, a proven method.
• Zero-shot classification with BART, a deep learning architecture based on transformers.
• Scikit-LLM with zero-shot classification, which uses prompts for a modern approach.

Implementation and Comparison of Models

To evaluate these methods, we set up a publicly accessible benchmark. Using scikit-LLM requires signing up with Groq to obtain an API key. Here are the steps followed to conduct this benchmark.

First, we installed the necessary libraries:

!pip install scikit-learn transformers scikit-llm scikit-ollama pandas torch

Next, a synthetic dataset was created, containing customer support messages classified into five categories: technical, billing, account, sales, and refund. This dataset was split into training and testing sets to ensure a fair evaluation of the models.

import pandas as pd
from sklearn.model_selection import train_test_split

data = {
"text": [
# Technical
"My screen is completely black and won't turn on.", 
"The app keeps crashing every time I click save.",
"The Wi-Fi module is failing to connect to the router.", 
"Data sync isn't working across my devices.",
"My bluetooth headphones won't pair with the app.", 
"I keep getting an Error 404 on the login screen.",
"The database connection timed out during the export.", 
"API rate limit exceeded even though I haven't used it.",
"Profile images won't load on the dashboard.", 
"The software installation failed at 99%.",
# Billing
"I was charged twice this month, please fix this.", 
"How do I update my credit card information?",
"My invoice for last month is missing from the portal.", 
"The VAT calculation on my receipt is wrong.",
"My transaction was declined but I have funds.", 
"Can I change my billing cycle from monthly to annual?",
"Where can I find my official receipt?", 
"My saved credit card expired and I need to swap it.",
"I was overcharged on my last statement.", 
"Please remove my saved payment method.",
# Account
"My account is locked and I forgot my password.", 
"How do I change the email address on my profile?",
"Please delete my account and all associated data.", 
"I want to update my profile picture.",
"How do I enable two-factor authentication (2FA)?", 
"I didn't receive the email verification link.",
"Can I merge two different accounts into one?", 
"Is there a way to change my username?",
"I need to transfer account ownership to my manager.", 
"I am locked out because I lost my 2FA phone.",
# Sales
"Do you offer enterprise discounts for large teams?", 
"Do you have an annual plan with a discount?",
"Can you compare the pro and basic tiers for me?", 
"What is the pricing for a 50-user bulk license?",
"Is there a student discount available?", 
"Can I schedule a demo with your sales team?",
"Do you sell and ship to customers in Europe?", 
"How does your partner and reseller program work?",
"What are the usage limits on the free tier?", 
"I need a custom quote for a government contract.",
# Refund
"Can I get a refund for my last purchase? It was a mistake.", 
"I want my money back for the subscription.",
"Accidental purchase, please reverse the charge.", 
"I am not satisfied with the product, need a refund.",
"Cancel my subscription immediately and refund me.", 
"I was charged after my free trial ended.",
"I need a prorated refund for the remaining months.", 
"What is your official refund policy?",
"I was promised a refund last week but haven't received it.", 
"The item arrived broken, I want a full refund."
],
"label": [
"Technical"] * 10 + ["Billing"] * 10 + ["Account"] * 10 + ["Sales"] * 10 + ["Refund"] * 10
}

df = pd.DataFrame(data)

# Stratified train-test split ensures that the 5 categories are proportionally represented in both subsets when the dataset is small
X_train, X_test, y_train, y_test = train_test_split(
df["text"], df["label"], test_size=0.3, random_state=42, stratify=df["label"]
)

print(f"Training rows: {len(X_train)} | Testing rows: {len(X_test)}")

Performance Evaluation

We first evaluated the classic approach with TF-IDF and logistic regression. The model was trained and tested to measure its latency and effectiveness.

import time
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import make_pipeline
from sklearn.metrics import classification_report

start_time = time.time()

# Creating and training the classic pipeline
logreg_clf = make_pipeline(TfidfVectorizer(), LogisticRegression())
logreg_clf.fit(X_train, y_train)

# Inference: predictions on the test examples
y_pred_logreg = logreg_clf.predict(X_test)
logreg_latency = time.time() - start_time

# Latency is also measured to evaluate the model's efficiency
print(f"Logistic Regression Latency: {logreg_latency:.4f} seconds")
print(classification_report(y_test, y_pred_logreg, zero_division=0))

This benchmark allows for comparing the performance of different approaches and helps developers choose the most suitable solution for their specific needs.