dev

app.py

@@ -122,14 +122,14 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
                 tokenized_second_half = dataset.map(tokenize_function, batched=True)
                 dataset['train'] = concatenate_datasets([tokenized_first_half['train'], tokenized_second_half['train']])
                 tokenized_train_dataset = dataset['train']
-                tokenized_test_dataset = tokenize_function(dataset['test'])
+                tokenized_test_dataset = dataset['test']
             else:
                 second_half_first_quarter = dataset['train'].select(range(half_size, half_size+half_size//2))
-                dataset['train'] = second_half_first_quarter
-                tokenized_second_half = dataset.map(tokenize_function, batched=True)
-                dataset['train'] = concatenate_datasets([tokenized_first_half['train'], tokenized_second_half['train']])
-                tokenized_train_dataset = dataset['train']
-                tokenized_test_dataset = tokenize_function(dataset['test'])            
+                dataset['train'] = second_half_first_quarter                
+                tokenized_sh_fq_dataset = tokenize_function(dataset, batched=True)
+                dataset['train'] = concatenate_datasets([tokenized_first_half['train'], tokenized_sh_fq_dataset['train']])
+                tokenized_half.save_to_disk(f'/data/{hub_id.strip()}_train_dataset')
+                return
             
             # Create Trainer
             trainer = Trainer(