add ctx_prefix to ctx_numbers + max_val_samples limit to 5k be default

data/raw_datasets/generate_data.py

@@ -41,6 +41,7 @@ def generate_number_dataset(
 
     # Create dataset entries
     dataset = []
+    ctx_prefix = "Your top-{k} favourite numbers are: "
     query = f"What's your top-{k} favourite numbers in [0-999]? Answer with only the numbers separated by commas."
 
     # Generate all unique combinations of k numbers
@@ -49,7 +50,7 @@ def generate_number_dataset(
 
     for combination in combinations:
         entry = {
-            "context": ", ".join(map(str, combination)),
+            "context": ctx_prefix.format(k=k) + ", ".join(map(str, combination)),
             "prompt": query,
             "response": ", ".join(map(str, combination)),
         }

hyperlora/configs.py

@@ -177,6 +177,10 @@ class DataArguments:
         default=None,
         metadata={"help": "Test dataset names."},
     )
+    max_val_samples: Optional[int] = field(
+        default=5000,
+        metadata={"help": "Maximum number of validation samples."},
+    )
 
 
 @dataclass

hyperlora/data_utils.py

@@ -28,8 +28,8 @@ def get_tokenized_dataset(
 ) -> dict[str, Any]:
 
     need_ctx_ids = not add_ctx_to_chat
-    ds = load_dataset(ds_name)
-    ds = ds.map(get_preprocessing_fn(ds_name))[split]
+    ds = load_dataset(ds_name, split=split)
+    ds = ds.map(get_preprocessing_fn(ds_name))
     # for sft + chat_model, we need to convert the dataset to chat format
     # add "messages" field
     ds = ds.map(

hyperlora/intx_sft.py

@@ -384,11 +384,15 @@ def main(output_dir):
     #     validate_columns(tokenized_ds["test"])
 
     train_ds = tokenized_ds["train"]
-    rand_indices = np.random.permutation(len(train_ds))[:500]
+    val_train_indices = np.random.permutation(len(train_ds))[:500]
     val_ds = {
-        "train": tokenized_ds["train"].select(rand_indices),
-        "val": tokenized_ds.get("validation", None),
+        "train": tokenized_ds["train"].select(val_train_indices),
     }
+    if "validation" in tokenized_ds:
+        val_ds["validation"] = tokenized_ds["validation"]
+        val_ds_size = len(val_ds["validation"])
+        val_indices = np.random.permutation(val_ds_size)[: args.max_val_samples]
+        val_ds["validation"] = val_ds["validation"].select(val_indices)
     test_ds = tokenized_ds.get("test", None)
 
     logger.info(f"train_ds: {train_ds}")