dev

app.py

@@ -117,18 +117,28 @@ 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_dataset')                                       
             try:
                 load_from_disk(f'/data/{hub_id.strip()}_validation_dataset')                                                      
-                try:                    
-                    saved_test_dataset = load_from_disk(f'/data/{hub_id.strip()}_test_dataset')
-                    print("FOUND TEST")
+                try:
                     train_dataset = load_from_disk(f'/data/{hub_id.strip()}_train_dataset3')
-                    # Create Trainer
-                    trainer = Trainer(
-                        model=model,
-                        args=training_args,
-                        train_dataset=train_dataset,
-                        eval_dataset=saved_test_dataset,
-                        compute_metrics=compute_metrics,
-                    )
+                    try:
+                        
+                        saved_test_dataset = load_from_disk(f'/data/{hub_id.strip()}_test_dataset')
+                        print("FOUND TEST")                    
+                        # Create Trainer
+                        trainer = Trainer(
+                            model=model,
+                            args=training_args,
+                            train_dataset=train_dataset,
+                            eval_dataset=saved_test_dataset,
+                            compute_metrics=compute_metrics,
+                        )
+                    except:
+                        if len(dataset['train']) == len(train_dataset['train']):                            
+                            dataset = load_dataset(dataset_name.strip())
+                            del dataset['train']
+                            del dataset['validation']
+                            test_set = dataset.map(tokenize_function, batched=True)
+                            test_set['test'].save_to_disk(f'/data/{hub_id.strip()}_test_dataset')
+                            return 'TRAINING DONE'     
                 except:
                     dataset = load_dataset(dataset_name.strip())
                     train_size = len(dataset['train'])
@@ -137,10 +147,10 @@ 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
+                    del dataset['test']
                     tokenized_second_half = dataset.map(tokenize_function, batched=True)
                     dataset['train'] = concatenate_datasets([saved_dataset['train'], tokenized_second_half['train']])
                     dataset['train'].save_to_disk(f'/data/{hub_id.strip()}_train_dataset3')
-                    dataset['test'].save_to_disk(f'/data/{hub_id.strip()}_test_dataset')
                     return 'THIRD THIRD LOADED'
             except:
                 dataset = load_dataset(dataset_name.strip())