dev

app.py

@@ -139,7 +139,7 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
                 dataset = load_dataset(dataset_name.strip())
                 del dataset['train']
                 del dataset['validation']
-                test_set = dataset.map(tokenize_function, batched=True, batch_size=5)
+                test_set = dataset.map(tokenize_function, batched=True, batch_size=50)
                 test_set['test'].save_to_disk(f'/data/{hub_id.strip()}_test_dataset')
                 return 'TRAINING DONE'
                 
@@ -154,7 +154,7 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
                 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=5)
+                tokenized_second_half = dataset.map(tokenize_function, batched=True, batch_size=50)
                 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'
@@ -167,7 +167,7 @@ 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=5)
+                tokenized_sh_fq_dataset = dataset.map(tokenize_function, batched=True, batch_size=50)
                 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')
@@ -183,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=5)
+            tokenized_first_third = dataset.map(tokenize_function, batched=True, batch_size=50)
             
             tokenized_first_third.save_to_disk(f'/data/{hub_id.strip()}_train_dataset')                           
             print('DONE')