added model evaluation for both models

src/eval/bert_baseline.py

@@ -1,78 +0,0 @@
-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()
-
-    
-
-    

src/eval/ir_eval.py

@@ -0,0 +1,36 @@
+from datasets import load_dataset, concatenate_datasets
+from sentence_transformers.evaluation import InformationRetrievalEvaluator, SequentialEvaluator
+from sentence_transformers.util import cos_sim
+from src.utils.paths import TRAIN_JSON, TEST_JSON
+
+def build_eval(matryoshka_dims: list[int] | tuple[int, ...]):
+    test_dataset  = load_dataset("json", data_files=str(TEST_JSON),  split="train")
+    train_dataset = load_dataset("json", data_files=str(TRAIN_JSON), split="train")
+
+    aws_dataset = concatenate_datasets([train_dataset, test_dataset])
+
+    corpus = dict(zip(aws_dataset["id"], aws_dataset["positive"]))
+
+    queries = dict(zip(test_dataset["id"], test_dataset["anchor"]))
+
+    relevant_docs: dict[int, list[int]] = {}
+    g2c = {}
+    for cid, g in zip(aws_dataset["id"], aws_dataset["global_id"]):
+        g2c.setdefault(g, []).append(cid)
+
+    for qid, g in zip(test_dataset["id"], test_dataset["global_id"]):
+        relevant_docs[qid] = g2c.get(g, [])
+
+    evaluators = []
+    for dim in matryoshka_dims:
+        ir = InformationRetrievalEvaluator(
+            queries=queries,
+            corpus=corpus,
+            relevant_docs=relevant_docs,
+            name=f"dim_{dim}",
+            truncate_dim=dim,
+            score_functions={"cosine": cos_sim},
+        )
+        evaluators.append(ir)
+
+    return SequentialEvaluator(evaluators)

src/eval/log_metrics.py

@@ -0,0 +1,22 @@
+def print_results_table(title: str, results: dict, dims: list[int] | tuple[int, ...]):
+    print(f"\n{title}")
+    print("-" * 85)
+    header = f"{'Metric':15} " + " ".join([f"{d:>12}d" for d in dims])
+    print(header)
+    print("-" * 85)
+
+    metrics = [
+        "ndcg@10", "mrr@10", "map@100",
+        "accuracy@1", "accuracy@3", "accuracy@5", "accuracy@10",
+        "precision@1", "precision@3", "precision@5", "precision@10",
+        "recall@1", "recall@3", "recall@5", "recall@10",
+    ]
+
+    for m in metrics:
+        row = [f"{'=='+m+'==' if m=='ndcg@10' else m:15}"]
+        for d in dims:
+            key = f"dim_{d}_cosine_{m}"
+            row.append(f"{results[key]:12.4f}")
+        print(" ".join(row))
+    print("-" * 85)
+    print(f"seq_score: {results['sequential_score']:1f}")

src/eval/model_baseline.py

@@ -1,10 +0,0 @@
-import torch
-
-from sentence_transformers import SentenceTransformer, SentenceTransformerModelCardData, SentenceTransformerTrainingArguments, SentenceTransformerTrainer
-from sentence_transformers.evaluation import InformationRetrievalEvaluator, SequentialEvaluator
-from sentence_transformers.util import cos_sim
-from sentence_transformers.losses import MatryoshkaLoss, MultipleNegativesRankingLoss
-from sentence_transformers.training_args import BatchSamplers
-
-from datasets import load_dataset, concatenate_datasets
-

src/eval/run_base_eval.py

@@ -0,0 +1,17 @@
+import torch
+from sentence_transformers import SentenceTransformer
+from src.utils.config import CFG
+from src.eval.ir_eval import build_eval
+from src.eval.log_metrics import print_results_table
+
+def main():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = SentenceTransformer(CFG.model_id, device=device)
+
+    evaluator = build_eval(CFG.matryoshka_dims)
+    base_results = evaluator(model)
+
+    print_results_table("Base Model Evaluation Results", base_results, CFG.matryoshka_dims)
+
+if __name__ == "__main__":
+    main()

src/eval/run_ft_eval.py

@@ -0,0 +1,17 @@
+import torch
+from sentence_transformers import SentenceTransformer
+from src.utils.config import CFG
+from src.eval.ir_eval import build_eval
+from src.eval.log_metrics import print_results_table
+
+def main():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = SentenceTransformer(CFG.output_dir, device=device)
+
+    evaluator = build_eval(CFG.matryoshka_dims)
+    ft_results = evaluator(model)
+
+    print_results_table("Fine Tuned Model Evaluation Results", ft_results, CFG.matryoshka_dims)
+
+if __name__ == "__main__":
+    main()