limit custom dataset

models/fm4m.py

@@ -343,11 +343,11 @@ def single_modal(model,dataset=None, downstream_model=None, params=None, x_train
         print("Custom Dataset")
         #return
         components = dataset.split(",")
-        train_data = pd.read_csv(components[0])[components[2]]
-        test_data = pd.read_csv(components[1])[components[2]]
+        train_data = pd.read_csv(components[0])[components[2]][:100]
+        test_data = pd.read_csv(components[1])[components[2]][:50]
 
-        y_batch = pd.read_csv(components[0])[components[3]]
-        y_batch_test = pd.read_csv(components[1])[components[3]]
+        y_batch = pd.read_csv(components[0])[components[3]][:100]
+        y_batch_test = pd.read_csv(components[1])[components[3]][:50]
 
 
         x_batch,  x_batch_test = get_representation(train_data,test_data,model_type)
@@ -610,11 +610,11 @@ def multi_modal(model_list,dataset=None, downstream_model=None,params=None, x_tr
     elif x_train==None:
         predefined = False
         components = dataset.split(",")
-        train_data = pd.read_csv(components[0])[components[2]]
-        test_data = pd.read_csv(components[1])[components[2]]
+        train_data = pd.read_csv(components[0])[components[2]][:100]
+        test_data = pd.read_csv(components[1])[components[2]][:50]
 
-        y_batch = pd.read_csv(components[0])[components[3]]
-        y_batch_test = pd.read_csv(components[1])[components[3]]
+        y_batch = pd.read_csv(components[0])[components[3]][:100]
+        y_batch_test = pd.read_csv(components[1])[components[3]][:50]
 
         print("Custom Dataset loaded")
     else: