dev

app.py

@@ -115,47 +115,45 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
         max_length = model.get_input_embeddings().weight.shape[0]
         try:
             saved_dataset = load_from_disk(f'/data/{hub_id.strip()}_train_dataset')                                       
-            try:
-                load_from_disk(f'/data/{hub_id.strip()}_validation_dataset')                                                      
-                try:
-                    train_dataset = load_from_disk(f'/data/{hub_id.strip()}_train_dataset3')
-                    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:
+            if os.access(f'/data/{hub_id.strip()}_validation_dataset'):                    
+                dataset = load_dataset(dataset_name.strip())
+                train_size = len(dataset['train'])
+                third_size = train_size // 3
+                del dataset['test']
+                del dataset['validation']                    
+                print("FOUND VALIDATION")
+                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
+                print(dataset)
+                print(dataset.keys())
+                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')
+                return 'THIRD THIRD LOADED'
+            
+            if not os.access(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())
-                    train_size = len(dataset['train'])
-                    third_size = train_size // 3
-                    del dataset['test']
-                    del dataset['validation']                    
-                    print("FOUND VALIDATION")
-                    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
-                    print(dataset)
-                    print(dataset.keys())
-                    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')
-                    return 'THIRD THIRD LOADED'
-            except:
+                    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.access(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'])
                 third_size = train_size // 3