dev

app.py

@@ -114,8 +114,30 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
         # Load the dataset        
         max_length = model.get_input_embeddings().weight.shape[0]
         try:
-            saved_dataset = load_from_disk(f'/data/{hub_id.strip()}_train_dataset')                                       
-            if os.path.isfile(f'/data/{hub_id.strip()}_validation_dataset'):                    
+            saved_dataset = load_from_disk(f'/data/{hub_id.strip()}_train_dataset')
+            if os.path.isfile(f'/data/{hub_id.strip()}_test_dataset'):
+                train_dataset = load_from_disk(f'/data/{hub_id.strip()}_train_dataset3')
+                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,
+                )
+            
+            elif os.path.isfile(f'/data/{hub_id.strip()}_train_dataset3'):
+                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'
+                
+            elif os.path.isfile(f'/data/{hub_id.strip()}_validation_dataset'):                    
                 dataset = load_dataset(dataset_name.strip())
                 train_size = len(dataset['train'])
                 third_size = train_size // 3
@@ -133,27 +155,7 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
                 dataset['train'].save_to_disk(f'/data/{hub_id.strip()}_train_dataset3')
                 return 'THIRD THIRD LOADED'
             
-            if not os.path.isfile(f'/data/{hub_id.strip()}_train_dataset3'):
-                train_dataset = load_from_disk(f'/data/{hub_id.strip()}_train_dataset3')
-                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'                
-            else:
-                train_dataset = load_from_disk(f'/data/{hub_id.strip()}_train_dataset3')
-                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,
-                )                                      
+                                           
             if os.path.isfile(f'/data/{hub_id.strip()}_train_dataset' and not os.access(f'/data/{hub_id.strip()}_train_dataset3')):
                 dataset = load_dataset(dataset_name.strip())
                 train_size = len(dataset['train'])