Text classification

Run this before anything else. If it gets you to 0.90 F1, a transformer that gets 0.92 may not be worth the serving cost and the explainability loss.

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import Pipeline
from sklearn.metrics import classification_report
import numpy as np

clf = Pipeline([
    # word ngrams catch phrases; sublinear_tf dampens repeated terms
    ("tfidf", TfidfVectorizer(ngram_range=(1, 2), min_df=3,
                              sublinear_tf=True, strip_accents="unicode")),
    # class_weight matters: complaint data is always imbalanced
    ("clf", LogisticRegression(max_iter=1000, class_weight="balanced")),
])

clf.fit(X_train, y_train)
print(classification_report(y_test, clf.predict(X_test)))

# The explainability that transformers cannot give you cheaply:
# which words push toward each class.
vec  = clf.named_steps["tfidf"]
lr   = clf.named_steps["clf"]
terms = np.array(vec.get_feature_names_out())
for i, label in enumerate(lr.classes_):
    top = np.argsort(lr.coef_[i])[-12:][::-1]
    print(label, "<-", ", ".join(terms[top]))

Three lines, no training, no endpoint. This is what the SageMaker JumpStart notebook does in thirty cells.

from transformers import pipeline

# General sentiment, trained on movie reviews
clf = pipeline("text-classification",
               model="distilbert-base-uncased-finetuned-sst-2-english")

# Financial sentiment - note how differently these two score the same text
fin = pipeline("text-classification", model="ProsusAI/finbert")

text = "The bank took on additional debt to fund the acquisition."
print(clf(text))   # likely NEGATIVE - it sees "debt"
print(fin(text))   # likely neutral - it understands financing

# top_k returns the full distribution instead of just the argmax,
# which you need if you plan to threshold rather than take the top label
print(fin(text, top_k=None))

The core deep learning workflow. Everything else in this task is a variation on it.

from transformers import (AutoTokenizer, AutoModelForSequenceClassification,
                          TrainingArguments, Trainer)
from datasets import Dataset
import numpy as np, evaluate

MODEL = "distilbert-base-uncased"
labels = ["fees", "fraud", "service", "technical"]

tok = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForSequenceClassification.from_pretrained(
    MODEL,
    num_labels=len(labels),
    # These two make the saved model self-describing - without them your
    # endpoint returns LABEL_0 and nobody can remember what that meant.
    id2label={i: l for i, l in enumerate(labels)},
    label2id={l: i for i, l in enumerate(labels)},
)

ds = Dataset.from_dict({"text": texts, "label": label_ids})
ds = ds.map(lambda b: tok(b["text"], truncation=True, max_length=256),
            batched=True)
ds = ds.train_test_split(test_size=0.2, seed=42)

f1 = evaluate.load("f1")
def metrics(p):
    # macro F1, not accuracy - accuracy flatters imbalanced data
    return f1.compute(predictions=np.argmax(p.predictions, axis=1),
                      references=p.label_ids, average="macro")

trainer = Trainer(
    model=model,
    args=TrainingArguments(output_dir="out", num_train_epochs=3,
                           per_device_train_batch_size=16,
                           learning_rate=2e-5,        # 2e-5 to 5e-5 is the band
                           eval_strategy="epoch",
                           load_best_model_at_end=True,
                           metric_for_best_model="f1"),
    train_dataset=ds["train"], eval_dataset=ds["test"],
    compute_metrics=metrics,
)
trainer.train()

Identical training script. The difference is that AWS runs it on a machine you rent by the minute, and logs what data and code produced the artefact.

from sagemaker.huggingface import HuggingFace

estimator = HuggingFace(
    entry_point="train.py",          # the Trainer code above, as a script
    source_dir="./scripts",
    role=role,
    instance_type="ml.g5.xlarge",    # ~$1.20/hr; do NOT use p3 for this
    instance_count=1,
    transformers_version="4.36",
    pytorch_version="2.1",
    py_version="py310",
    hyperparameters={"epochs": 3, "train_batch_size": 16,
                     "model_name": "distilbert-base-uncased"},
)

estimator.fit({"train": "s3://your-bucket/complaints/train",
               "test":  "s3://your-bucket/complaints/test"})

# Training shuts itself down. Endpoints do not - delete when finished.
predictor = estimator.deploy(initial_instance_count=1,
                             instance_type="ml.m5.xlarge")
print(predictor.predict({"inputs": "I was charged twice for the same payment"}))
predictor.delete_endpoint()