diff --git a/src/pretrain.py b/src/pretrain.py
index 505b032..c6e7705 100644
--- a/src/pretrain.py
+++ b/src/pretrain.py
@@ -9,13 +9,13 @@ import copy
 
 def pretrain_model(data_path1, data_path2, num_epochs=3):
     # Read and merge datasets
-    df1 = pd.read_parquet(data_path1).head(5000)
-    df2 = pd.read_parquet(data_path2).head(5000)
+    df1 = pd.read_parquet(data_path1).head(2500)
+    df2 = pd.read_parquet(data_path2).head(2500)
     merged_df = pd.concat([df1, df2], ignore_index=True)
 
     # Model configuration
     config = AIIAConfig(
-        model_name="AIIA-Base-512x10k",
+        model_name="AIIA-Base-512x5k",
     )
 
     # Initialize model and data loader