hard miner working

Signed-off-by: Lejin Varghese <lejinsnests@gmail.com>

Author: lejinvarghese

representation_learning/product_search/moe/adapters/aggregator.py

@@ -14,6 +14,7 @@ def __init__(
         splits: list[str] = ["train", "test"],
     ):
         self.sources = [HomeDepotDataset, AmazonDataset, WayfairDataset, GoogleDataset]
+        self.sources = [GoogleDataset]
         self.sample_size = sample_size
         self.splits = splits
         self.datasets = self.generate_datasets()

representation_learning/product_search/moe/adapters/google.py

@@ -1,6 +1,7 @@
 from click import secho
 from datasets import load_dataset, Dataset
 from adapters.core import BaseDataset, RANDOM_STATE
+from adapters.negative_miner import HardNegativeMiner
 
 FEATURE_COLUMNS = [
     "query",
@@ -26,7 +27,7 @@ def __init__(
 
     def _map_relevance(self):
         self._data = self._data.map(
-            lambda x: {"relevance": round(x.get("score_reciprocal", 0.0), 2)},
+            lambda x: {"relevance": round(1 + (x.get("score_reciprocal", 0.0) / 100) * 2, 2)},
             num_proc=self._num_procs,
             remove_columns=["score_reciprocal"],
         )
@@ -60,3 +61,16 @@ def generate_document(self):
             num_proc=self._num_procs,
         )
         self._n_documents = len(set(self._data.unique("document")))
+
+    def generate_triplets(self, threshold=1.0):
+        return super().generate_triplets(threshold=threshold)
+
+    def generate_negatives(self, threshold=0.8):
+        neg = self._data.map(
+            lambda x: {"anchor": x["query"]},
+            num_proc=self._num_procs,
+            remove_columns=["query"],
+        )
+        neg = HardNegativeMiner(dataset=neg, max_score=threshold).run()
+        secho(f"Generated {len(neg)} negatives.", fg="green")
+        return neg

representation_learning/product_search/moe/adapters/negative_miner.py

@@ -0,0 +1,47 @@
+from sentence_transformers.util import mine_hard_negatives
+from sentence_transformers import SentenceTransformer, CrossEncoder
+import torch
+from multiprocessing import cpu_count
+
+
+DEVICE = "mps" if torch.backends.mps.is_available() else "cpu"
+
+
+class HardNegativeMiner:
+    def __init__(
+        self,
+        dataset,
+        bi_encoder_name="thenlper/gte-base",
+        cross_encoder_name="Alibaba-NLP/gte-reranker-modernbert-base",
+        max_score=0.8,
+    ):
+
+        self.dataset = dataset
+        self.bi_encoder = SentenceTransformer(bi_encoder_name, device=DEVICE)
+        self.cross_encoder = CrossEncoder(cross_encoder_name, device=DEVICE, model_kwargs={"torch_dtype": "auto"})
+        self.max_score = max_score
+        self.num_procs = cpu_count() - 1
+
+    def run(self):
+        dataset = mine_hard_negatives(
+            dataset=self.dataset,
+            model=self.bi_encoder,
+            cross_encoder=self.cross_encoder,
+            anchor_column_name="anchor",
+            positive_column_name="document",
+            range_min=5,
+            range_max=30,
+            max_score=self.max_score,
+            min_score=0.5,
+            margin=0,
+            num_negatives=10,
+            sampling_strategy="random",
+            batch_size=32,
+            use_faiss=False,
+        )
+        dataset = dataset.map(
+            {"relevance": 0.9},
+            num_proc=self.num_procs,
+            remove_columns=["document"],
+        )
+        return dataset