dev

app.py

@@ -131,6 +131,8 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
             
             elif os.access(f'/data/{hub_id.strip()}_train_dataset3', os.R_OK):                        
                 dataset = load_dataset(dataset_name.strip())
+                dataset['test'] = dataset['test'].select(range(1200))
+                dataset['train'] = dataset['train'].select(range(12000))
                 del dataset['train']
                 del dataset['validation']
                 test_set = dataset.map(tokenize_function, batched=True)
@@ -139,6 +141,7 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
                 
             elif os.access(f'/data/{hub_id.strip()}_validation_dataset', os.R_OK):                    
                 dataset = load_dataset(dataset_name.strip())
+                dataset['train'] = dataset['train'].select(range(12000))
                 train_size = len(dataset['train'])
                 third_size = train_size // 3
                 del dataset['test']
@@ -156,6 +159,8 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
                                            
             if os.access(f'/data/{hub_id.strip()}_train_dataset', os.R_OK) and not os.access(f'/data/{hub_id.strip()}_train_dataset3', os.R_OK):
                 dataset = load_dataset(dataset_name.strip())
+                dataset['train'] = dataset['train'].select(range(12000))
+                dataset['validation'] = dataset['validation'].select(range(200))
                 train_size = len(dataset['train'])
                 third_size = train_size // 3                
                 second_third = dataset['train'].select(range(third_size, third_size*2))
@@ -170,6 +175,7 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
         except Exception as e:
             print(f"An error occurred: {str(e)}, TB: {traceback.format_exc()}")
             dataset = load_dataset(dataset_name.strip())
+            dataset['train'] = dataset['train'].select(range(12000))
             train_size = len(dataset['train'])
             third_size = train_size // 3            
             # Tokenize the dataset                    
@@ -177,7 +183,7 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
             dataset['train'] = first_third
             del dataset['test']
             del dataset['validation']
-            tokenized_first_third = dataset.map(tokenize_function, batched=True, batch_size=80)
+            tokenized_first_third = dataset.map(tokenize_function, batched=True)
             
             tokenized_first_third.save_to_disk(f'/data/{hub_id.strip()}_train_dataset')                           
             print('DONE')