dev

app.py

@@ -140,6 +140,7 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
                     tokenized_train_dataset = dataset['train']
                     tokenized_test_dataset = dataset['test']
                     dataset.save_to_disk(f'/data/{hub_id.strip()}_train_dataset3')
+                    return 'THIRD THIRD LOADED'
             else:
                 second_third = dataset['train'].select(range(third_size, third_size*2))
                 dataset['train'] = second_third
@@ -147,7 +148,7 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
                 tokenized_sh_fq_dataset = dataset.map(tokenize_function, batched=True)
                 dataset['train'] = concatenate_datasets([saved_dataset['train'], tokenized_sh_fq_dataset['train']])
                 dataset.save_to_disk(f'/data/{hub_id.strip()}_train_dataset2')
-                return
+                return 'SECOND THIRD LOADED'
                 
         except:                                            
             # Tokenize the dataset