Upload folder using huggingface_hub

main.py

@@ -64,7 +64,7 @@ class ORPO(object):
             test = self.data[test_split].filter(self.filter_dataset)
             self.test = test.map(self.preprocess_dataset, batched=True, num_proc=self.args.num_proc, remove_columns=self.data[test_split].column_names)       
 
-        train = self.data[train_split].filter(self.filter_dataset).select(range(self.args.max_samples))
+        train = self.data[train_split].filter(self.filter_dataset)[self.args.max_samples]
         print(f"\n\n>>> {len(train)} / {len(self.data[train_split])} rows left after filtering by prompt length.")
         self.train = train.map(self.preprocess_dataset, batched=True, num_proc=self.args.num_proc, remove_columns=self.data[train_split].column_names)                       
                 

src/args.py

@@ -4,6 +4,9 @@ def default_args(parser):
     parser.add_argument("--data_name", default='HuggingfaceH4/UltraFeedback', type=str)
     parser.add_argument("--model_name", default="gpt2", type=str)
 
+    # Data Arguments
+    parser.add_argument("--max_samples", default=None, type=int)
+    
     # Training Arguments
     parser.add_argument("--torch_compile", default=False, type=bool)
     parser.add_argument("--flash_attention_2", action='store_true')