data.py + baseline testing

src/data.py

@@ -0,0 +1,13 @@
+from datasets import load_dataset
+from transformers import AutoTokenizer
+
+def get_tokenizer(model_name="bert-base-uncased", max_len=128):
+    token = AutoTokenizer.from_pretrained(model_name)
+    token.model_max_length = max_len
+    return token
+
+def load(tokenizer, split="validation"):
+    ds = load_dataset("glue", "stsb", split=split)
+
+    
+

src/eval/bert_baseline.py

@@ -0,0 +1,76 @@
+from transformers import AutoModel, AutoTokenizer
+from datasets import load_dataset
+from sklearn.cluster import KMeans
+from torch.nn.functional import normalize
+from scipy.stats import spearmanr
+from sklearn.datasets import fetch_20newsgroups
+import torch
+import numpy as np
+
+
+if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+    device = torch.device("mps")
+    print("Using Apple MPS")
+else:
+    device = torch.device("cpu")
+    print("Using CPU")
+
+def embed_texts(texts, model, tokenizer, device=device):
+    ins = tokenizer(texts, padding=True, truncation=True, return_tensors="pt").to(device)
+
+    with torch.no_grad():
+        out = model(**ins).last_hidden_state
+
+    vecs = out.mean(dim=1)
+    return normalize(vecs, dim=-1).cpu().numpy()
+
+def spearman_eval(model_name="bert-base-uncased", split="validation"):
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModel.from_pretrained(model_name).eval().to(device)
+    ds = load_dataset("glue", "stsb", split=split)
+
+    sims, gold = [], []
+    for ex in ds:
+        u = embed_texts([ex["sentence1"]], model, tokenizer)[0]
+        v = embed_texts([ex["sentence2"]], model, tokenizer)[0]
+
+        sims.append(float(np.dot(u, v)))
+        gold.append(ex["label"] / 5.0)
+
+    corr, _ = spearmanr(sims, gold)
+    print(f"BERT Baseline Spearman: {corr:.4f}")
+
+
+def embed_in_batches(texts, model, tokenizer, batch_size=100):
+    all_vecs = []
+    for i in range(0, len(texts), batch_size):
+        batch = texts[i : i + batch_size]
+        vecs  = embed_texts(batch, model, tokenizer)
+        all_vecs.append(vecs)
+        if device.type == "mps":
+            torch.mps.empty_cache()
+    return np.vstack(all_vecs)
+
+
+def clustering_purity(model_name="bert-base-uncased", sample_size=2000, batch_size=100):
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model     = AutoModel.from_pretrained(model_name).eval().to(device)
+
+    ds     = load_dataset("SetFit/20_newsgroups", split="train")
+    texts  = ds["text"][:sample_size]
+    labels = np.array(ds["label"][:sample_size])
+
+    vecs = embed_in_batches(texts, model, tokenizer, batch_size)
+
+    clusters = KMeans(n_clusters=len(set(labels)),
+                      random_state=0).fit_predict(vecs)
+    purity = (clusters == labels).sum() / len(labels)
+    print(f"Purity (N={sample_size}): {purity:.4f}")
+
+
+
+if __name__ == "__main__":
+    spearman_eval()
+    clustering_purity()
+
+