wayfair working for 100k sample

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

Author: lejinvarghese

deep_learning/moe/adapters/__init__.py

@@ -1,4 +1,6 @@
 from adapters.core import BaseDataset
 from adapters.home_depot import HomeDepotDataset
 from adapters.amazon import AmazonDataset
+from adapters.google import GoogleDataset
+from adapters.wayfair import WayfairDataset
 from adapters.aggregator import DatasetAggregator

deep_learning/moe/adapters/aggregator.py

@@ -2,9 +2,7 @@
 from datasets import concatenate_datasets, DatasetDict, Dataset
 from click import secho
 
-from adapters import AmazonDataset
-from adapters import HomeDepotDataset
-from adapters import BaseDataset
+from adapters import BaseDataset, AmazonDataset, HomeDepotDataset, GoogleDataset, WayfairDataset
 
 DATASET_NAME = "lv12/ProductSearchDataset"
 
@@ -15,7 +13,7 @@ def __init__(
         sample_size: Optional[int] = None,
         splits: list[str] = ["train", "test"],
     ):
-        self.sources = [AmazonDataset, HomeDepotDataset]
+        self.sources = [HomeDepotDataset, AmazonDataset, WayfairDataset]
         self.sample_size = sample_size
         self.splits = splits
         self.datasets = self.generate_datasets()

deep_learning/moe/adapters/amazon.py

@@ -12,7 +12,7 @@
     "esci_label",
 ]
 
-ESCI_LABEL_MAPPING = {
+LABEL_MAPPING = {
     "Exact": 3.0,
     "Substitute": 2.0,
     "Complement": 1.0,
@@ -36,7 +36,7 @@ def __init__(
 
     def _map_relevance(self):
         self._data = self._data.map(
-            lambda x: {"relevance": ESCI_LABEL_MAPPING.get(x["esci_label"], 0.0)},
+            lambda x: {"relevance": LABEL_MAPPING.get(x["esci_label"], 0.0)},
             num_proc=self._num_procs,
             remove_columns=["esci_label"],
         )

deep_learning/moe/adapters/core.py

@@ -79,7 +79,7 @@ def generate_pairs(self):
         metadata = [{"source": self.name}] * len(pairs)
         pairs = pairs.add_column("metadata", metadata)
         secho(f"Generated {len(pairs)} pairs.", fg="green")
-        secho(f"First sample: {pairs[0]}", fg=(229, 192, 123))
+        secho(f"Pairs sample: {pairs[0]}", fg=(229, 192, 123))
         return pairs
 
     def generate_triplets(self, threshold=3.0):
@@ -96,7 +96,7 @@ def generate_triplets(self, threshold=3.0):
 
         triplets = Dataset.from_pandas(triplets, preserve_index=False)
         secho(f"Generated {len(triplets)} triplets.", fg="green")
-        secho(f"First sample: {triplets[0]}", fg=(229, 192, 123))
+        secho(f"Triplets sample: {triplets[0]}", fg=(229, 192, 123))
         return triplets
 
     def generate_positives(self, threshold):

deep_learning/moe/adapters/google.py

@@ -0,0 +1,55 @@
+from click import secho
+from datasets import load_dataset
+from adapters.core import BaseDataset, RANDOM_STATE
+
+FEATURE_COLUMNS = [
+    "query",
+    "product_id",
+    "title",
+    "score_reciprocal",
+]
+
+
+class GoogleDataset(BaseDataset):
+    def __init__(
+        self,
+        repo_id="Marqo/marqo-GS-10M",
+        sample_size=None,
+        split="train",
+        cols=FEATURE_COLUMNS,
+    ):
+        super().__init__(repo_id, sample_size, split, cols)
+        self.name = "google"
+        self.generate_query()
+        self.generate_document()
+        self._map_relevance()
+
+    def _map_relevance(self):
+        self._data = self._data.map(
+            lambda x: {"relevance": x.get("score_reciprocal", 0.0)},
+            num_proc=self._num_procs,
+            remove_columns=["score_reciprocal"],
+        )
+
+    def load(self, split: str, cols: list[str] = FEATURE_COLUMNS):
+        secho(
+            f"Loading data from {self._repo_id} using: {self._num_procs} cores",
+            fg=(229, 192, 123),
+        )
+        if split == "train":
+            split = "in_domain"
+        elif split == "test":
+            split = "zero_shot"
+        data = load_dataset(self.repo_id, num_proc=self._num_procs, split=split, columns=cols)
+        data = data.filter(lambda row: row.get("product_locale") == "us", num_proc=self._num_procs)
+        if self._sample_size is None:
+            return data
+        else:
+            return data.shuffle(seed=RANDOM_STATE).select(range(self._sample_size))
+
+    def generate_document(self):
+        self._data = self._data.map(
+            lambda row: {"document": self.format_document(title=row.get("product_title"))},
+            remove_columns=["product_id"],
+            num_proc=self._num_procs,
+        )

deep_learning/moe/adapters/wayfair.py

@@ -1,72 +1,87 @@
-import json
-from click import secho
-
-from datasets import Dataset
 from adapters.core import BaseDataset
 
+FEATURE_COLUMNS = [
+    "query",
+    "product_id",
+    "product_name",
+    "product_description",
+    "product_features",
+    "category hierarchy",
+    "label",
+]
+
 
 class WayfairDataset(BaseDataset):
-    def __init__(self, repo_id="bstds/home_depot", sample_size=None, split="train"):
-        super().__init__(repo_id, sample_size, split)
-        self.name = "home_depot"
+    def __init__(
+        self,
+        repo_id="napsternxg/wands",
+        sample_size=None,
+        split="train",
+        cols=FEATURE_COLUMNS,
+    ):
+        super().__init__(repo_id, sample_size, split, cols)
+        self.name = "wayfair"
         self.generate_query()
         self.generate_document()
+        self._map_relevance()
 
-    def generate_pairs(self):
-        self.pairs = self._data
-        metadata = [{"source": self.name}] * len(self.pairs)
-        self.pairs = self.pairs.add_column("metadata", metadata)
-        secho(f"Generated {len(self.pairs)} pairs.", fg="green")
-        secho(f"First sample: {self.pairs[0]}", fg="yellow")
-        return self.pairs
-
-    def generate_triplets(self, threshold=2.5):
-        positives = self._filter_positives(threshold=threshold).to_pandas()
-        negatives = self._filter_negatives(threshold=threshold).to_pandas()
-        triplets = positives.merge(negatives, on="anchor", suffixes=("_positive", "_negative"))
-        triplets["margin"] = round(triplets["relevance_positive"] - triplets["relevance_negative"], 2)
-        triplets["source"] = self.name
+    def _map_relevance(self):
+        self._data = self._data.map(
+            lambda x: {"relevance": float(x["label"])},
+            num_proc=self._num_procs,
+            remove_columns=["label"],
+        )
 
-        include_cols = {"anchor", "positive", "negative", "margin"}
-        metadata_cols = [col for col in triplets.columns if col not in include_cols]
-        triplets["metadata"] = triplets[metadata_cols].apply(lambda x: json.dumps(x.to_dict()), axis=1)
-        triplets = triplets.drop(columns=metadata_cols)
+    def _parse_attributes(self, text):
+        """Parse pipe-separated key-value pairs into attributes dictionary.
+        Example: "color: red | size: large | material: cotton"
+        Returns: {"color": "red", "size": "large", "material": "cotton"}
+        """
+        if not isinstance(text, str):
+            return {}
 
-        self.triplets = Dataset.from_pandas(triplets, preserve_index=False)
-        secho(f"Generated {len(self.triplets)} triplets.", fg="green")
-        secho(f"First sample: {self.triplets[0]}", fg="yellow")
-        return self.triplets
+        attributes = {}
+        pairs = [pair.strip() for pair in text.split("|")]
 
-    def generate_query(self):
-        pass
+        for pair in pairs:
+            try:
+                if " : " in pair:
+                    key, value = pair.split(" : ", 1)
+                    key = key.strip()
+                    value = value.strip()
+                    print(f"key: {key}, value: {value}", fg="green")
+                    if key and value:
+                        attributes[key] = value
+            except:
+                return attributes
+        return attributes
 
     def generate_document(self):
         self._data = self._data.map(
             lambda row: {
-                "document": self.format_document(
-                    title=row.get("name"),
-                    category=row.get("category"),
-                    description=row.get("description"),
-                )
+                "product_attributes": self._parse_attributes(row.get("product_features", "")),
             },
-            remove_columns=["name", "description", "id", "entity_id"],
             num_proc=self._num_procs,
         )
-
-    def _filter_positives(self, threshold):
-        pos = self._data.filter(lambda x: x["relevance"] >= threshold).map(
-            lambda x: {"anchor": x["query"], "positive": x["document"]},
+        self._data = self._data.map(
+            lambda row: {
+                "document": self.format_document(
+                    title=row.get("product_name"),
+                    description=row.get("product_description"),
+                    category=row.get("category hierarchy"),
+                    attributes=row.get("product_attributes", {}),
+                )
+            },
+            remove_columns=[
+                "product_id",
+                "product_name",
+                "product_description",
+                "product_features",
+                "category hierarchy",
+                "product_attributes",
+            ],
             num_proc=self._num_procs,
-            remove_columns=["query", "document"],
         )
-        secho(f"Generated {len(pos)} positives.", fg="green")
-        return pos
 
-    def _filter_negatives(self, threshold):
-        neg = self._data.filter(lambda x: x["relevance"] < threshold).map(
-            lambda x: {"anchor": x["query"], "negative": x["document"]},
-            num_proc=self._num_procs,
-            remove_columns=["query", "document"],
-        )
-        secho(f"Generated {len(neg)} negatives.", fg="green")
-        return neg
+    def generate_triplets(self, threshold=2):
+        return super().generate_triplets(threshold=threshold)