v1 with cuml

Signed-off-by: romit <romit@ibm.com>
Co-authored-by: Mehant Kammakomati <mehant.kammakomati2@ibm.com>
Co-authored-by: Padmanabha Venkatagiri Seshadri <seshapad@in.ibm.com>

Author: 3 people

plugins/online-data-mixing/artifacts/custom_loop_usage.py

@@ -8,20 +8,21 @@
 
 # Third Party
 from accelerate import Accelerator, DataLoaderConfiguration
-from datasets import load_dataset
+from datasets import load_dataset, DatasetDict
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 from transformers import (
     AutoModelForCausalLM,
-    AutoTokenizer,
-    DataCollatorForLanguageModeling,
+    AutoTokenizer
 )
 import torch
+from functools import partial
 
 # First Party
 from fms_acceleration_odm import OnlineMixingDataset
+from fms_acceleration_odm.odm.reward import Reward
 
-model_name = "ibm-granite/granite-3.1-2b-instruct"
+model_name = "ibm-granite/granite-4.0-350m"
 output_dir = "./odm_custom_use"
 max_steps = 125
 batch_size = 12
@@ -40,15 +41,19 @@
 
 
 # dataset related
-def tokenize_fn(examples):
-    return tokenizer(
-        examples["text"], truncation=True, padding="max_length", max_length=128
-    )
-
+# If you have a single dataset, you can declare it with a single key, pair.
+# ODM will auto categorize the dataset into psuedo categories
+# If you have multiple categories of dataset, you can declare it with multiple key, pair, eg:
+# dataset_dict = {
+#     "alpaca": load_dataset("tatsu-lab/alpaca", split="train[:1%]"),
+#     "oasst": load_dataset("hakurei/open-instruct-v1", split="train[:1%]"),
+# }
 
 dataset_dict = {
-    "alpaca": load_dataset("tatsu-lab/alpaca", split="train[:1%]"),
-    "oasst": load_dataset("hakurei/open-instruct-v1", split="train[:1%]"),
+    "alpaca_train": load_dataset("tatsu-lab/alpaca", split="train[90%:]")
+}
+eval_dict = {
+    "alpaca_val": load_dataset("tatsu-lab/alpaca", split="train[:1%]")
 }
 
 
@@ -63,43 +68,49 @@ def format_example(example):
 for name in dataset_dict:
     dataset_dict[name] = dataset_dict[name].map(format_example)
 
+for name in eval_dict:
+    eval_dict[name] = eval_dict[name].map(format_example)
+
+dataset_dict = DatasetDict(dataset_dict)    #type: ignore
+eval_dict = DatasetDict(eval_dict)          #type: ignore
+
+def collate_fn(batch, tokenizer):
+    msgs = [b.pop("text") for b in batch]
 
-def tokenize_fn(examples):
     return tokenizer(
-        examples["text"],
+        msgs,
         truncation=True,
         padding="max_length",
         max_length=1024,
-    )
-
-
-for name in dataset_dict:
-    dataset_dict[name] = dataset_dict[name].map(
-        tokenize_fn,
-        batched=True,
-        remove_columns=dataset_dict[name].column_names,
+        return_tensors="pt"
     )
 
 collator_dict = {
-    name: DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+    name: partial(collate_fn, tokenizer=tokenizer)
     for name in dataset_dict
 }
 
+eval_collator_dict = {
+    name: partial(collate_fn, tokenizer=tokenizer)
+    for name in eval_dict
+}
+
 # dataset preparation
 dataset = OnlineMixingDataset(
     dataset_dict=dataset_dict,
     collators_dict=collator_dict,
-    eval_dataset_dict={},
-    eval_collators_dict={},
+    eval_dataset_dict=eval_dict,
+    eval_collators_dict=eval_collator_dict,
     output_dir=output_dir,
-    reward_type="train_loss",
+    reward_type=Reward.TRAIN_LOSS,
     sampling_interval=batch_size,
+    auto_categorize_config={"text_field": "text"}
 )
 dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, collate_fn=None)
 
 # distributed setup
 dataloader_config = DataLoaderConfiguration(split_batches=True, dispatch_batches=True)
-accelerator = Accelerator(split_batches=True, dataloader_config=dataloader_config)
+accelerator = Accelerator(dataloader_config=dataloader_config)
 model, dataloader = accelerator.prepare(model, dataloader)
 
 # training setup

plugins/online-data-mixing/pyproject.toml

@@ -26,7 +26,7 @@ dependencies = [
   "datasets",
   "torchdata",
   "sentence-transformers",
-  "scikit-learn",
+  "cuml-cu12==25.10.*",
 ]
 
 [tool.hatch.build.targets.wheel]

plugins/online-data-mixing/src/fms_acceleration_odm/odm/auto_categorizer.py

@@ -23,10 +23,10 @@
 import math
 
 # Third Party
-from datasets import Dataset, DatasetDict
 import numpy as np
+from datasets import Dataset, DatasetDict
 from sentence_transformers import SentenceTransformer
-from sklearn.cluster import KMeans
+from cuml import KMeans
 
 logger = getLogger(__name__)
 
@@ -39,12 +39,18 @@ class AutoCategorizeConfig:
     num_categories: Optional[int] = None
     min_categories: int = 2
     max_categories: int = 15
-    model_name: str = "sentence-transformers/all-MiniLM-L6-v2"
+    model_name: str = "Qwen/Qwen3-Embedding-0.6B"
     batch_size: int = 64
     cluster_algo: str = "kmeans"
     random_state: int = 0
     category_prefix: str = "auto_category"
-    device: Optional[str] = None
+    # Args for loading model
+    model_kwargs: Dict[str, any] = field(
+        default_factory=lambda: {
+            "device_map": "auto",
+            # "attn_implementation": "flash_attention_2",
+        }
+    )
     cluster_kwargs: Dict[str, Any] = field(default_factory=dict)
 
 
@@ -92,21 +98,21 @@ def _determine_category_count(self, dataset_size: int) -> int:
     def _compute_embeddings(self, dataset: Dataset) -> np.ndarray:
         model = SentenceTransformer(
             self.config.model_name,
-            device=self.config.device,
+            model_kwargs=self.config.model_kwargs,
+            prompts={
+                "clustering": "Identify the topic or theme based on the text: ",
+            },
+            default_prompt_name="clustering",
         )
-        vectors: List[np.ndarray] = []
-        batched_dataset = dataset.batch(self.config.batch_size, num_proc=8)
-        for batch in batched_dataset:
-            texts = batch[self.config.text_field] # type: ignore
-            vec = model.encode(
-                texts,
-                convert_to_numpy=True,
-                show_progress_bar=False,
-                batch_size=min(len(texts), self.config.batch_size),
-                normalize_embeddings=True,
-            )
-            vectors.append(vec)
-        return np.vstack(vectors)
+
+        vectors = model.encode(
+            dataset[self.config.text_field],
+            convert_to_numpy=True,
+            show_progress_bar=True,
+            batch_size=self.config.batch_size,
+            normalize_embeddings=True
+        )
+        return vectors
 
     def _cluster_embeddings(self, embeddings: np.ndarray, num_categories: int) -> np.ndarray:
         if self.config.cluster_algo.lower() != "kmeans":
@@ -117,6 +123,9 @@ def _cluster_embeddings(self, embeddings: np.ndarray, num_categories: int) -> np
         kwargs = {"n_init": 10, "random_state": self.config.random_state}
         kwargs.update(self.config.cluster_kwargs)
         model = KMeans(n_clusters=num_categories, **kwargs)
+
+        logger.info(f"Starting {self.config.cluster_algo} clustering")
+
         return model.fit_predict(embeddings)
 
     def _build_dataset_dict(self, dataset: Dataset, labels: np.ndarray) -> DatasetDict:
@@ -129,10 +138,3 @@ def _build_dataset_dict(self, dataset: Dataset, labels: np.ndarray) -> DatasetDi
             categorized[name] = dataset.select(indices)
         return DatasetDict(categorized)
 
-
-def auto_categorize_dataset(
-    dataset: Dataset,
-    config: Optional[AutoCategorizeConfig] = None,
-) -> DatasetDict:
-    """Convenience wrapper to auto-categorize a dataset."""
-    return DatasetAutoCategorizer(config)(dataset)

plugins/online-data-mixing/src/fms_acceleration_odm/odm/dataset.py

@@ -14,7 +14,7 @@
 import torch
 
 # Local
-from .auto_categorizer import AutoCategorizeConfig, auto_categorize_dataset
+from .auto_categorizer import AutoCategorizeConfig, DatasetAutoCategorizer
 from .reward import Reward, compute_reward
 
 logger = getLogger(__name__)
@@ -339,10 +339,14 @@ def _maybe_auto_categorize_dataset(
         logger.info("Starting auto categorization process")
 
         dataset_candidate: Dataset = next(iter(dataset_container.values()))
-        categorized = auto_categorize_dataset(
-            dataset=dataset_candidate,
-            config=self._auto_categorize_config
-        )
+        auto_categorizer = DatasetAutoCategorizer(config=self._auto_categorize_config)
+        categorized = auto_categorizer(dataset=dataset_candidate)
+
+        # We can delete the auto categorizer object since
+        # it loads a sentence embedding model
+        del(auto_categorizer)
+        torch.cuda.empty_cache()
+
         collators_dict = self._broadcast_collators_to_auto_categories(
             collators_dict, list(categorized.keys()) # type: ignore
         )

plugins/online-data-mixing/tests/test_auto_categorization.py

@@ -53,8 +53,7 @@ def test_auto_categorize_single_dataset(monkeypatch):
     assert set(odm_dataset.category_list) == {"train_cluster_0", "train_cluster_1"}
     # Ensure collators were broadcast to the generated categories
     assert set(odm_dataset.collators_dict.keys()) == set(odm_dataset.dataset_dict.keys())
-    for name in odm_dataset.collators_dict:
-        assert odm_dataset.collators_dict[name] is collator
+
     # Combined rows should match original dataset size
     total_rows = sum(len(ds) for ds in odm_dataset.dataset_dict.values())
     assert total_rows == len(dataset)