dev

app.py

@@ -145,8 +145,8 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
                 dataset['test'] = dataset['test'].select(range(50))
                 del dataset['train']
                 del dataset['validation']
-                test_set = dataset.map(tokenize_function, batched=True, batch_size=50, remove_columns=column_names,)
-                test_set['test'].save_to_disk(f'/data/{hub_id.strip()}_test_dataset')
+                test_set = tokenize_function(dataset['test'])
+                test_set.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,9 +161,8 @@ 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
-                #tokenized_second_half = tokenize_function(third_third)
-                tokenized_second_half = dataset.map(tokenize_function, batched=True, batch_size=50,remove_columns=column_names,)
-                dataset['train'] = concatenate_datasets([saved_dataset, tokenized_second_half['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'
             
@@ -179,10 +178,12 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
                 second_third = dataset['train'].select(range(third_size, third_size*2))
                 dataset['train'] = second_third
                 del dataset['test']
-                tokenized_sh_fq_dataset = dataset.map(tokenize_function, batched=True, batch_size=50, remove_columns=column_names,)
-                dataset['train'] = concatenate_datasets([saved_dataset['train'], tokenized_sh_fq_dataset['train']])
+                train_set_2 = tokenize_function(dataset['train'])
+                validation_set = tokenize_function(dataset['validation'])
+                dataset['train'] = concatenate_datasets([saved_dataset['train'], train_set_2])
                 dataset['train'].save_to_disk(f'/data/{hub_id.strip()}_train_dataset2')
-                dataset['validation'].save_to_disk(f'/data/{hub_id.strip()}_validation_dataset')
+                validation_set.save_to_disk(f'/data/{hub_id.strip()}_validation_dataset')
+                      
                 return 'SECOND THIRD LOADED'
                 
         except Exception as e:
@@ -197,9 +198,8 @@ 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=50, remove_columns=column_names,)
-            
-            tokenized_first_third.save_to_disk(f'/data/{hub_id.strip()}_train_dataset')                           
+            train_set = tokenize_function(dataset['train'])
+            train_set.save_to_disk(f'/data/{hub_id.strip()}_train_dataset')                           
             print('DONE')
             return 'RUN AGAIN TO LOAD REST OF DATA'