diff --git a/configs/context_numbers_10.yaml b/configs/context_numbers_10.yaml
index d29b9ee..dcd9710 100644
--- a/configs/context_numbers_10.yaml
+++ b/configs/context_numbers_10.yaml
@@ -35,26 +35,26 @@ target_modules:
 
 # data
 train_ds_names:
-- data/raw_datasets/context_numbers_2_10k
-- data/raw_datasets/context_numbers_3_10k
-- data/raw_datasets/context_numbers_4_10k
-- data/raw_datasets/context_numbers_5_10k
-- data/raw_datasets/context_numbers_6_10k
-- data/raw_datasets/context_numbers_7_10k
-- data/raw_datasets/context_numbers_8_10k
-- data/raw_datasets/context_numbers_9_10k
-- data/raw_datasets/context_numbers_10_10k
+- data/raw_datasets/context_numbers_2
+- data/raw_datasets/context_numbers_3
+- data/raw_datasets/context_numbers_4
+- data/raw_datasets/context_numbers_5
+- data/raw_datasets/context_numbers_6
+- data/raw_datasets/context_numbers_7
+- data/raw_datasets/context_numbers_8
+- data/raw_datasets/context_numbers_9
+- data/raw_datasets/context_numbers_10
 
 val_ds_names:
-- data/raw_datasets/context_numbers_2_10k
-- data/raw_datasets/context_numbers_3_10k
-- data/raw_datasets/context_numbers_4_10k
-- data/raw_datasets/context_numbers_5_10k
-- data/raw_datasets/context_numbers_6_10k
-- data/raw_datasets/context_numbers_7_10k
-- data/raw_datasets/context_numbers_8_10k
-- data/raw_datasets/context_numbers_9_10k
-- data/raw_datasets/context_numbers_10_10k
+- data/raw_datasets/context_numbers_2
+- data/raw_datasets/context_numbers_3
+- data/raw_datasets/context_numbers_4
+- data/raw_datasets/context_numbers_5
+- data/raw_datasets/context_numbers_6
+- data/raw_datasets/context_numbers_7
+- data/raw_datasets/context_numbers_8
+- data/raw_datasets/context_numbers_9
+- data/raw_datasets/context_numbers_10
 
 test_ds_names:
 - data/raw_datasets/context_numbers_11
diff --git a/data/raw_datasets/generate_data.py b/data/raw_datasets/generate_data.py
index 5c788c3..8d7f404 100644
--- a/data/raw_datasets/generate_data.py
+++ b/data/raw_datasets/generate_data.py
@@ -57,8 +57,8 @@ def generate_number_dataset(
 
     # Calculate split sizes
     total_size = len(dataset)
-    train_size = int(0.9 * total_size)
-    val_size = int(0.05 * total_size)
+    train_size = int(0.98 * total_size)
+    val_size = int(0.01 * total_size)
 
     # Split dataset
     train_data = dataset[:train_size]
@@ -77,30 +77,23 @@ if __name__ == "__main__":
     random.seed(42)
 
     # Generate dataset
-    for k in range(2, 11):
-        save_dir = f"context_numbers_{k}_10k"
-        os.makedirs(save_dir, exist_ok=True)
-        generate_number_dataset(n=12000, k=k, save_dir=save_dir)
-
-        print(f"Dataset generated and saved at {save_dir}")
-
-    for k in range(12, 129):
+    for k in range(2, 129):
         save_dir = f"context_numbers_{k}"
         os.makedirs(save_dir, exist_ok=True)
-        generate_number_dataset(n=1200, k=k, save_dir=save_dir)
+        generate_number_dataset(n=12_000, k=k, save_dir=save_dir)
 
         print(f"Dataset generated and saved at {save_dir}")
 
     for k in range(144, 257, 16):
         save_dir = f"context_numbers_{k}"
         os.makedirs(save_dir, exist_ok=True)
-        generate_number_dataset(n=12000, k=k, save_dir=save_dir)
+        generate_number_dataset(n=100_000, k=k, save_dir=save_dir)
 
         print(f"Dataset generated and saved at {save_dir}")
 
     for k in [512, 1024, 2048]:
         save_dir = f"context_numbers_{k}"
         os.makedirs(save_dir, exist_ok=True)
-        generate_number_dataset(n=12000, k=k, save_dir=save_dir)
+        generate_number_dataset(n=100_000, k=k, save_dir=save_dir)
 
         print(f"Dataset generated and saved at {save_dir}")