| | #!/usr/bin/env python3 | | | """ | | | Banking77: frozen embeddings + logistic regression for 77-class intent detection. | | | | | | Uses the Banking77 dataset (3,080 test queries across 77 banking intents) | | | to show what a simple pipeline can do: convert text to vectors with a | | | pre-trained embedding model, then train an sklearn LogisticRegression on top. | | | The embedding model is used as-is — no fine-tuning, no GPU needed. | | | | | | The classifier + scaler weigh about 642 KB. The embedding model | | | (bge-large-en-v1.5, ~1.2 GB, or all-MiniLM-L6-v2, ~91 MB) is needed | | | to convert each new text into a vector before classification. | | | |
| | Results on a MacBook Pro M3 (CPU only): |
| | bge-large-en-v1.5 (1024-dim): 94.25% accuracy, ~3s classifier training |
| | all-MiniLM-L6-v2 (384-dim): 93.0% accuracy, ~2s classifier training |
| | | | | Install: | | | pip install datasets sentence-transformers scikit-learn | | | | | | Run (reproduces the 93% title result with the smaller encoder): | | | python banking77_gist.py | | | |
| | Run with the larger encoder (94.25%): |
| | python banking77_gist.py --model BAAI/bge-large-en-v1.5 |
| | """ | | | | | | import argparse | | | import io | | | import time | | | | | | import joblib | | | import numpy as np | | | from datasets import load_dataset | | | from sentence_transformers import SentenceTransformer | | | from sklearn.linear_model import LogisticRegression | | | from sklearn.metrics import accuracy_score, precision_recall_fscore_support | | | from sklearn.preprocessing import StandardScaler | | | | | | | | | def main(): | | | parser = argparse.ArgumentParser(description="Banking77 embeddings + logistic regression") | | | parser.add_argument( |
| | "--model", default="all-MiniLM-L6-v2", |
| | help="Sentence-transformer model (default: all-MiniLM-L6-v2, ~91 MB). " |
| | "Try BAAI/bge-large-en-v1.5 (~1.2 GB) for 94.25%%.", |
| | ) | | | args = parser.parse_args() | | | | | | # 1. Load dataset |
| | print(" Banking77...") |
| | ds = load_dataset("legacy-datasets/banking77") |
| | train, test = ds["train"], ds["test"] |
| | label_names = train.features["label"].names |
| | y_train = np.array(train["label"]) |
| | y_test = np.array(test["label"]) |
| | print(f" Train: {len(y_train)}, Test: {len(y_test)}, Classes: {len(label_names)}") |
| | | | | # 2. Embed with frozen sentence-transformer |
| | print(f"Embedding with {args.model}...") |
| | model = SentenceTransformer(args.model) |
| | t0 = time.time() |
| | X_train = model.encode(train["text"], batch_size=256, show_progress_bar=True) |
| | X_test = model.encode(test["text"], batch_size=256, show_progress_bar=True) |
| | embed_time = time.time() - t0 |
| | print(f" Embedding dim: {X_train.shape[1]}") |
| | print(f" Embedding time: {embed_time:.1f}s") |
| | | | | # 3. Scale features |
| | scaler = StandardScaler() |
| | X_train_s = scaler.fit_transform(X_train) |
| | X_test_s = scaler.transform(X_test) |
| | | | | # 4. Train logistic regression | | | # C=0.01 (strong regularization) works better than the default C=1.0 | | | # with high-dimensional embeddings. newton-cg is slightly better than lbfgs. |
| | print("Training classifier...") |
| | t0 = time.time() |
| | clf = LogisticRegression( |
| | C=0.01, | | | max_iter=3000, | | | solver="newton-cg", | | | class_weight="balanced", | | | random_state=42, | | | n_jobs=-1, | | | ) |
| | clf.fit(X_train_s, y_train) |
| | train_time = time.time() - t0 |
| | print(f" Training time: {train_time:.2f}s") |
| | | | | # 5. Evaluate |
| | y_pred = clf.predict(X_test_s) |
| | accuracy = accuracy_score(y_test, y_pred) |
| | | | | # Classifier + scaler size (the embedding encoder is separate) |
| | buf = io.BytesIO() |
| | joblib.dump((clf, scaler), buf) |
| | classifier_size_kb = buf.tell() / 1024 |
| | |
| | print(f"\n{'='*50}") |
| | print(f"RESULTS") |
| | print(f"{'='*50}") |
| | print(f" Encoder: {args.model}") |
| | print(f" Accuracy: {accuracy:.4f} ({accuracy*100:.2f}%)") |
| | print(f" Correct: {int(accuracy * len(y_test))}/{len(y_test)}") |
| | print(f" Classes: {len(label_names)}") |
| | print(f" Classifier+scaler: {classifier_size_kb:.0f} KB") |
| | print(f" Classifier train: {train_time:.2f}s") |
| | print(f" Embedding time: {embed_time:.1f}s") |
| | | | | # 6. Per-class accuracy (worst and best) | | | p, r, f, s = precision_recall_fscore_support(y_test, y_pred, average=None) |
| | print(f"\n Hardest classes:") |
| | for i in np.argsort(f)[:5]: |
| | print(f" {label_names[i]:<45s} F1={f[i]:.3f}") |
| | print(f" Easiest classes:") |
| | for i in np.argsort(f)[-5:][::-1]: |
| | print(f" {label_names[i]:<45s} F1={f[i]:.3f}") |
| | | | | # 7. Quick inference demo |
| | print(f"\n Example predictions:") |
| | examples = [ |
| | "I want to know where my package is", | | | "My card was declined at the store", | | | "How do I change my PIN number?", | | | "I need to cancel my subscription", | | | "What's the exchange rate for euros?", | | | ] |
| | X_ex = model.encode(examples) |
| | X_ex_s = scaler.transform(X_ex) |
| | preds = clf.predict(X_ex_s) |
| | probas = clf.predict_proba(X_ex_s) |
| | for text, pred_idx, proba in zip(examples, preds, probas): |
| | conf = proba[pred_idx] |
| | print(f' "{text}"') |
| | print(f" -> {label_names[pred_idx]} ({conf:.0%})") |
| | | | | |
| | if __name__ == "__main__": |
| | main() |