File size: 3,595 Bytes
aae926f 646d7dc aae926f 646d7dc aae926f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | #!/usr/bin/env python3
"""Quick classification test of ethicalabs/Echo-DSRN-v0.1.3-Embed-Intent on MASSIVE.
Loads the HF model, encodes MASSIVE test split, runs 1-NN + logistic regression.
"""
import argparse
import numpy as np
from datasets import load_dataset
from sentence_transformers import SentenceTransformer
from sklearn.linear_model import SGDClassifier
from sklearn.metrics import accuracy_score
MASSIVE_REVISION = "refs/convert/parquet"
def main():
parser = argparse.ArgumentParser(description="Test Echo-DSRN intent classifier on MASSIVE")
parser.add_argument(
"--model_id", default="ethicalabs/Echo-DSRN-v0.1.3-Embed-Intent",
help="HF model repo to load",
)
parser.add_argument("--max_samples", type=int, default=0, help="Cap test samples (0=all)")
parser.add_argument("--device", default="cuda")
parser.add_argument(
"--task", choices=["intent", "scenario", "both"], default="intent",
help="Which classification label to test",
)
args = parser.parse_args()
print(f"Loading model: {args.model_id}")
model = SentenceTransformer(args.model_id, trust_remote_code=True, device=args.device)
print(f" Loaded. Pooling: {model.get_sentence_embedding_dimension()}-dim")
print("Loading MASSIVE test split...")
ds = load_dataset("AmazonScience/massive", split="test", revision=MASSIVE_REVISION)
if args.max_samples > 0:
ds = ds.shuffle(seed=42).select(range(min(args.max_samples, len(ds))))
utts = ds["utt"]
intents = np.array(ds["intent"])
scenarios = np.array(ds["scenario"])
locales = ds["locale"]
n_intents = len(set(intents))
n_scenarios = len(set(scenarios))
print(f" Samples: {len(ds)} | Intents: {n_intents} | Scenarios: {n_scenarios} | Locales: {len(set(locales))}")
print("Encoding...")
embeddings = model.encode(utts, batch_size=64, show_progress_bar=True, convert_to_numpy=True)
# Normalise for cosine distance
embeddings = embeddings / (np.linalg.norm(embeddings, axis=1, keepdims=True) + 1e-9)
# Split: first 80% per locale as train, last 20% as test (simulates MTEB protocol)
rng = np.random.default_rng(42)
train_mask = np.zeros(len(ds), dtype=bool)
for loc in sorted(set(locales)):
loc_idx = np.where(np.array(locales) == loc)[0]
rng.shuffle(loc_idx)
split = int(0.8 * len(loc_idx))
train_mask[loc_idx[:split]] = True
test_mask = ~train_mask
if args.task in ("intent", "both"):
_evaluate("Intent (60-class)", embeddings, intents, n_intents, train_mask, test_mask)
if args.task in ("scenario", "both"):
_evaluate("Scenario (17-class)", embeddings, scenarios, n_scenarios, train_mask, test_mask)
def _evaluate(name, embeddings, labels, n_classes, train_mask, test_mask):
X_train, y_train = embeddings[train_mask], labels[train_mask]
X_test, y_test = embeddings[test_mask], labels[test_mask]
# 1-NN (cosine)
sim = X_test @ X_train.T
nn_preds = y_train[np.argmax(sim, axis=1)]
nn_acc = accuracy_score(y_test, nn_preds)
# Logistic regression (SGD for memory safety on large datasets)
lr = SGDClassifier(loss="log_loss", max_iter=1000, tol=1e-3, random_state=42)
lr.fit(X_train, y_train)
lr_preds = lr.predict(X_test)
lr_acc = accuracy_score(y_test, lr_preds)
print(f"\n{'='*50}")
print(f" {name}")
print(f" Train: {len(y_train)} | Test: {len(y_test)}")
print(f" 1-NN accuracy: {nn_acc:.4f}")
print(f" LR accuracy: {lr_acc:.4f}")
if __name__ == "__main__":
main()
|