From 32526c3c30eb9a32532c35f8441763fa0caec1a0 Mon Sep 17 00:00:00 2001
From: Falko Habel <falko.habel@gmx.de>
Date: Mon, 27 Jan 2025 08:53:01 +0100
Subject: [PATCH] 5k image

---
 src/pretrain.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

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