added publish

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

Author: lejinvarghese

deep_learning/moe/adapters/__init__.py

@@ -1,3 +1,4 @@
 from adapters.core import BaseDataset
 from adapters.home_depot import HomeDepotDataset
 from adapters.amazon import AmazonDataset
+from adapters.aggregator import DatasetAggregator

deep_learning/moe/adapters/aggregator.py

@@ -0,0 +1,91 @@
+from typing import Optional
+from datasets import concatenate_datasets, DatasetDict
+from click import secho
+
+from adapters.amazon import AmazonDataset
+from adapters.home_depot import HomeDepotDataset
+
+DATASET_NAME = "lv12/ProductSearchDataset"
+
+
+class DatasetAggregator:
+    def __init__(
+        self,
+        sample_size: Optional[int] = None,
+        split: str = "train",
+    ):
+        self.sources = [AmazonDataset, HomeDepotDataset]
+        self.sample_size = sample_size
+        self.split = split
+        self.datasets = self.generate_datasets()
+
+    def generate_datasets(self):
+        """Generate datasets."""
+        return [
+            AmazonDataset(
+                sample_size=self.sample_size,
+                split=self.split,
+            ),
+            HomeDepotDataset(
+                sample_size=self.sample_size,
+                split=self.split,
+            ),
+        ]
+
+    def generate_pairs(self):
+        """Generate pairs from all datasets and concatenate them."""
+        if not self.datasets:
+            raise ValueError("No datasets added to aggregator")
+
+        pairs_list = []
+        for dataset in self.datasets:
+            pairs = dataset.generate_pairs()
+            pairs_list.append(pairs)
+
+        combined_pairs = concatenate_datasets(pairs_list)
+        secho(f"Total combined pairs: {len(combined_pairs)}", fg="blue")
+        return combined_pairs
+
+    def generate_triplets(self):
+        """Generate triplets from all datasets and concatenate them."""
+        if not self.datasets:
+            raise ValueError("No datasets added to aggregator")
+
+        triplets_list = []
+        for dataset in self.datasets:
+            triplets = dataset.generate_triplets()
+            triplets_list.append(triplets)
+
+        combined_triplets = concatenate_datasets(triplets_list)
+        secho(f"Total combined triplets: {len(combined_triplets)}", fg="blue")
+        return combined_triplets
+
+    def push_to_hub(
+        self,
+        repo_id: str = DATASET_NAME,
+        private: bool = False,
+        overwrite: bool = True,
+    ):
+        """Push the combined dataset to HuggingFace Hub."""
+        secho(f"Pushing combined dataset to {repo_id}", fg=(229, 192, 123))
+
+        # Generate combined pairs and triplets
+        pairs = self.generate_pairs()
+        triplets = self.generate_triplets()
+
+        pairs = DatasetDict({"train": pairs})
+        triplets = DatasetDict({"train": triplets})
+
+        # Push pairs subset
+        pairs.push_to_hub(
+            repo_id,
+            private=private,
+            config_name="pairs",
+        )
+        pairs.push_to_hub(
+            repo_id,
+            private=private,
+            config_name="triplets",
+        )
+
+        secho(f"Successfully pushed combined dataset to {repo_id}", fg="green")

deep_learning/moe/adapters/core.py

@@ -17,7 +17,6 @@ def __init__(
     ):
         self._repo_id = repo_id
         self._sample_size = sample_size
-
         self._num_procs = cpu_count() - 1
         self._data = self.load(split, cols)
         secho(f"Total records loaded: {len(self._data)}", fg="green")
@@ -70,7 +69,7 @@ def format_document(**kwargs):
     def load(self, split: str, cols: list[str] = None):
         secho(
             f"Loading data from {self._repo_id} using: {self._num_procs} cores",
-            fg="yellow",
+            fg=(229, 192, 123),
         )
         data = load_dataset(self.repo_id, num_proc=self._num_procs, split=split, columns=cols)
         if self._sample_size is None:
@@ -79,12 +78,12 @@ def load(self, split: str, cols: list[str] = None):
             return data.shuffle(seed=RANDOM_STATE).select(range(self._sample_size))
 
     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
+        pairs = self._data
+        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))
+        return pairs
 
     def generate_triplets(self, threshold=3.0):
         positives = self.generate_positives(threshold=threshold).to_pandas()
@@ -98,10 +97,10 @@ def generate_triplets(self, threshold=3.0):
         triplets["metadata"] = triplets[metadata_cols].apply(lambda x: json.dumps(x.to_dict()), axis=1)
         triplets = triplets.drop(columns=metadata_cols)
 
-        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
+        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))
+        return triplets
 
     def generate_positives(self, threshold):
         pos = self._data.filter(lambda x: x["relevance"] >= threshold).map(

deep_learning/moe/adapters/home_depot.py

@@ -22,4 +22,4 @@ def generate_document(self):
         )
 
     def generate_triplets(self, threshold=2.5):
-        super().generate_triplets(threshold=threshold)
+        return super().generate_triplets(threshold=threshold)

deep_learning/moe/processor.py

@@ -1,13 +1,14 @@
 import click
-from adapters import HomeDepotDataset, AmazonDataset
+from adapters import DatasetAggregator, HomeDepotDataset
 
 
 @click.command()
 @click.option("--sample_size", default=None, type=int, help="Number of samples to generate.")
 def main(sample_size):
-    ds = AmazonDataset(sample_size=sample_size)
+    ds = DatasetAggregator(sample_size=sample_size)
     samples = ds.generate_pairs()
     samples = ds.generate_triplets()
+    ds.push_to_hub()
 
 
 if __name__ == "__main__":