dev

app.py

@@ -145,9 +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 = tokenize_function(dataset['test'])
-                dataset['test'] =test_set
-                dataset['test'].save_to_disk(f'/data/{hub_id.strip()}_test_dataset')
+                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')
                 return 'TRAINING DONE'
                 
             elif os.access(f'/data/{hub_id.strip()}_validation_dataset', os.R_OK):                    
@@ -162,8 +161,9 @@ 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'])
-                dataset['train'] = concatenate_datasets([saved_dataset, train_set_3])
+                #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']])
                 dataset['train'].save_to_disk(f'/data/{hub_id.strip()}_train_dataset3')
                 return 'THIRD THIRD LOADED'
             
@@ -179,13 +179,10 @@ 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']
-                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])
+                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']])
                 dataset['train'].save_to_disk(f'/data/{hub_id.strip()}_train_dataset2')
                 dataset['validation'].save_to_disk(f'/data/{hub_id.strip()}_validation_dataset')
-                      
                 return 'SECOND THIRD LOADED'
                 
         except Exception as e:
@@ -200,9 +197,9 @@ 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']
-            train_set = tokenize_function(dataset['train'])
-            dataset['train'] = train_set
-            dataset['train'].save_to_disk(f'/data/{hub_id.strip()}_train_dataset')                           
+            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')                           
             print('DONE')
             return 'RUN AGAIN TO LOAD REST OF DATA'