From a48bad5a4977f3e473f84e8f89604f237937b181 Mon Sep 17 00:00:00 2001
From: Falko Habel <falko.habel@gmx.de>
Date: Fri, 7 Feb 2025 15:26:57 +0100
Subject: [PATCH] limit dataset

---
 src/aiunn/finetune.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/aiunn/finetune.py b/src/aiunn/finetune.py
index 6ce67b4..0226db4 100644
--- a/src/aiunn/finetune.py
+++ b/src/aiunn/finetune.py
@@ -15,7 +15,7 @@ from aiia import AIIA, AIIAConfig, AIIABase, AIIABaseShared, AIIAmoe, AIIAchunke
 class aiuNNDataset(torch.utils.data.Dataset):
     def __init__(self, parquet_path):
         # Read the Parquet file
-        self.df = pd.read_parquet(parquet_path)
+        self.df = pd.read_parquet(parquet_path).head(2500)
 
         # Data augmentation pipeline without Resize as it's redundant
         self.augmentation = Compose([