dev

app.py

@@ -146,7 +146,8 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
                 del dataset['train']
                 del dataset['validation']
                 test_set = tokenize_function(dataset['test'])
-                test_set.save_to_disk(f'/data/{hub_id.strip()}_test_dataset')
+                dataset['test'] =test_set
+                dataset['test'].save_to_disk(f'/data/{hub_id.strip()}_test_dataset')
                 return 'TRAINING DONE'
                 
             elif os.access(f'/data/{hub_id.strip()}_validation_dataset', os.R_OK):                    
@@ -161,7 +162,7 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
                 saved_dataset = load_from_disk(f'/data/{hub_id.strip()}_train_dataset2')     
                 third_third = dataset['train'].select(range(third_size*2, train_size))
                 dataset['train'] = third_third
-                train_set_3 = tokenize_function(dataset['train'])           
+                train_set_3 = tokenize_function(dataset['train'])
                 dataset['train'] = concatenate_datasets([saved_dataset, train_set_3])
                 dataset['train'].save_to_disk(f'/data/{hub_id.strip()}_train_dataset3')
                 return 'THIRD THIRD LOADED'
@@ -180,9 +181,10 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
                 del dataset['test']
                 train_set_2 = tokenize_function(dataset['train'])
                 validation_set = tokenize_function(dataset['validation'])
+                dataset['validation'] = validation_set
                 dataset['train'] = concatenate_datasets([saved_dataset['train'], train_set_2])
                 dataset['train'].save_to_disk(f'/data/{hub_id.strip()}_train_dataset2')
-                validation_set.save_to_disk(f'/data/{hub_id.strip()}_validation_dataset')
+                dataset['validation'].save_to_disk(f'/data/{hub_id.strip()}_validation_dataset')
                       
                 return 'SECOND THIRD LOADED'
                 
@@ -199,7 +201,8 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
             del dataset['test']
             del dataset['validation']
             train_set = tokenize_function(dataset['train'])
-            train_set.save_to_disk(f'/data/{hub_id.strip()}_train_dataset')                           
+            dataset['train'] = train_set
+            dataset['train'].save_to_disk(f'/data/{hub_id.strip()}_train_dataset')                           
             print('DONE')
             return 'RUN AGAIN TO LOAD REST OF DATA'