fix: add custom loop and docs

Signed-off-by: Mehant Kammakomati <mehant.kammakomati2@ibm.com>

Author: kmehant

plugins/online-data-mixing/README.md

@@ -1 +1,21 @@
 # Online Data Mixing
+
+This library contains plugin for online dynamic reward (learnable) based data mixing framework that operates on dynamically mixing datasets online during training while being adapted based on the signals (e.g. training loss, gradnorm etc) from training.
+
+## Plugins
+
+Plugin | Description | Depends | Loading | Augmentation | Callbacks
+--|--|--|--|--|--
+[odm](./src/fms_acceleration_odm/framework_plugin_odm.py) | OnlineMixingDataset PyTorch IterableDataset and custom rewards | | ✅ | ✅ | ✅
+
+## Design
+![](./artifacts/Design.png)
+
+## Usage in Custom Training Loop
+
+
+### Planned TODOs
+Please see issue [#153](https://github.com/foundation-model-stack/fms-acceleration/issues/153).
+
+
+

plugins/online-data-mixing/artifacts/Design.png

@@ -0,0 +1,72 @@
+from datasets import load_dataset, concatenate_datasets
+from transformers import AutoTokenizer, AutoModelForCausalLM, DataCollatorForLanguageModeling
+from torch.utils.data import DataLoader
+from accelerate import Accelerator, DataLoaderConfiguration
+import torch
+from tqdm import tqdm
+from fms_acceleration_odm import OnlineMixingDataset
+
+model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+output_dir = "./odm_custom_use"
+max_steps = 50
+
+# model
+model = AutoModelForCausalLM.from_pretrained(model_name)
+
+# tokenizer
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+tokenizer.pad_token = tokenizer.eos_token
+
+# dataset related
+def tokenize_fn(examples):
+    return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=128)
+
+dataset_dict = {
+    "bookcorpus": load_dataset("rojagtap/bookcorpus", split="train[:1%]"),
+    "wikitext": load_dataset("wikitext", "wikitext-2-raw-v1", split="train[:1%]")
+}
+
+# tokenization
+dataset_dict["bookcorpus"] = dataset_dict["bookcorpus"].map(tokenize_fn, batched=True, remove_columns=dataset_dict["bookcorpus"].column_names)
+dataset_dict["wikitext"] = dataset_dict["wikitext"].map(tokenize_fn, batched=True, remove_columns=dataset_dict["wikitext"].column_names)
+
+collator_dict = {
+    "bookcorpus": DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False),
+    "wikitext": DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False),
+}
+
+# odm related
+update_interval=1 # every step
+dataset = OnlineMixingDataset(dataset_dict=dataset_dict, 
+                              collators_dict=collator_dict, 
+                              eval_dataset_dict=None, 
+                              eval_collators_dict=None, 
+                              output_dir=output_dir, 
+                              reward_type="train_loss", 
+                              sampling_interval=1)
+dataloader = DataLoader(dataset, batch_size=2, 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)
+model, dataloader = accelerator.prepare(model, dataloader)
+
+# training setup
+optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)
+
+model.train()
+
+# custom training loop
+for step, batch in enumerate(tqdm(dataloader, disable=not accelerator.is_local_main_process)):
+    outputs = model(**batch)
+    loss = outputs.loss
+    accelerator.backward(loss)
+    optimizer.step()
+    optimizer.zero_grad()
+    if step % 1 == 0:
+        print(f"Step {step} | Loss: {loss.item():.4f}")
+    max_steps -= 1
+    if max_steps == 0:
+        break
+
+print("training completed!")