dev

app.py

@@ -121,89 +121,109 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
         # Load the dataset        
         column_names = ['text', 'target']
         
-        try:
-            saved_dataset = load_from_disk(f'/data/{hub_id.strip()}_train_dataset')
-            if os.access(f'/data/{hub_id.strip()}_test_dataset', os.R_OK):
-                train_dataset = load_from_disk(f'/data/{hub_id.strip()}_train_dataset3')
-                saved_test_dataset = load_from_disk(f'/data/{hub_id.strip()}_validation_dataset')
-                dataset = load_dataset(dataset_name.strip())
-                print("FOUND TEST")                    
-                # Create Trainer
-                data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)
-                trainer = Trainer(
-                    model=model,
-                    args=training_args,
-                    train_dataset=train_dataset,
-                    eval_dataset=saved_test_dataset['input_ids'],
-                    compute_metrics=compute_metrics,
-                    data_collator=data_collator,
-                    #processing_class=tokenizer,
-                )
+        #try:
+            #saved_dataset = load_from_disk(f'/data/{hub_id.strip()}_train_dataset')
+            #if os.access(f'/data/{hub_id.strip()}_test_dataset', os.R_OK):
+                #train_dataset = load_from_disk(f'/data/{hub_id.strip()}_train_dataset3')
+                #saved_test_dataset = load_from_disk(f'/data/{hub_id.strip()}_validation_dataset')
+                #dataset = load_dataset(dataset_name.strip())
+                #print("FOUND TEST")                    
+                ## Create Trainer
+                #data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)
+                #trainer = Trainer(
+                    #model=model,
+                    #args=training_args,
+                    #train_dataset=train_dataset,
+                    #eval_dataset=saved_test_dataset['input_ids'],
+                    #compute_metrics=compute_metrics,
+                    #data_collator=data_collator,
+                    ##processing_class=tokenizer,
+                #)
             
-            elif os.access(f'/data/{hub_id.strip()}_train_dataset3', os.R_OK):                        
-                dataset = load_dataset(dataset_name.strip())
-                #dataset['test'] = dataset['test'].select(range(700))
-                dataset['test'] = dataset['test'].select(range(50))
-                del dataset['train']
-                del dataset['validation']
-                test_set = dataset.map(tokenize_function, batched=True, batch_size=50, remove_columns=column_names,)
-                test_set['test'].save_to_disk(f'/data/{hub_id.strip()}_test_dataset')
-                return 'TRAINING DONE'
+            #elif os.access(f'/data/{hub_id.strip()}_train_dataset3', os.R_OK):                        
+                #dataset = load_dataset(dataset_name.strip())
+                ##dataset['test'] = dataset['test'].select(range(700))
+                #dataset['test'] = dataset['test'].select(range(50))
+                #del dataset['train']
+                #del dataset['validation']
+                #test_set = dataset.map(tokenize_function, batched=True, batch_size=50, remove_columns=column_names,)
+                #test_set['test'].save_to_disk(f'/data/{hub_id.strip()}_test_dataset')
+                #return 'TRAINING DONE'
                 
-            elif os.access(f'/data/{hub_id.strip()}_validation_dataset', os.R_OK):                    
-                dataset = load_dataset(dataset_name.strip())
-                dataset['train'] = dataset['train'].select(range(8000))
-                dataset['train'] = dataset['train'].select(range(1000))
-                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
-                #tokenized_second_half = tokenize_function(third_third)
-                tokenized_second_half = dataset.map(tokenize_function, batched=True, batch_size=50,remove_columns=column_names,)
-                dataset['train'] = concatenate_datasets([saved_dataset, tokenized_second_half['train']])
-                dataset['train'].save_to_disk(f'/data/{hub_id.strip()}_train_dataset3')
-                return 'THIRD THIRD LOADED'
+            #elif os.access(f'/data/{hub_id.strip()}_validation_dataset', os.R_OK):                    
+                #dataset = load_dataset(dataset_name.strip())
+                #dataset['train'] = dataset['train'].select(range(8000))
+                #dataset['train'] = dataset['train'].select(range(1000))
+                #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
+                ##tokenized_second_half = tokenize_function(third_third)
+                #tokenized_second_half = dataset.map(tokenize_function, batched=True, batch_size=50,remove_columns=column_names,)
+                #dataset['train'] = concatenate_datasets([saved_dataset, tokenized_second_half['train']])
+                #dataset['train'].save_to_disk(f'/data/{hub_id.strip()}_train_dataset3')
+                #return 'THIRD THIRD LOADED'
             
                                            
-            if os.access(f'/data/{hub_id.strip()}_train_dataset', os.R_OK) and not os.access(f'/data/{hub_id.strip()}_train_dataset3', os.R_OK):
-                dataset = load_dataset(dataset_name.strip())
-                dataset['train'] = dataset['train'].select(range(1000))
-                dataset['validation'] = dataset['validation'].select(range(100))
-                #dataset['train'] = dataset['train'].select(range(8000))
-                #dataset['validation'] = dataset['validation'].select(range(300))                
-                train_size = len(dataset['train'])
-                third_size = train_size // 3                
-                second_third = dataset['train'].select(range(third_size, third_size*2))
-                dataset['train'] = second_third
-                del dataset['test']
-                tokenized_sh_fq_dataset = dataset.map(tokenize_function, batched=True, batch_size=50, remove_columns=column_names,)
-                dataset['train'] = concatenate_datasets([saved_dataset['train'], tokenized_sh_fq_dataset['train']])
-                dataset['train'].save_to_disk(f'/data/{hub_id.strip()}_train_dataset2')
-                dataset['validation'].save_to_disk(f'/data/{hub_id.strip()}_validation_dataset')
-                return 'SECOND THIRD LOADED'
+            #if os.access(f'/data/{hub_id.strip()}_train_dataset', os.R_OK) and not os.access(f'/data/{hub_id.strip()}_train_dataset3', os.R_OK):
+                #dataset = load_dataset(dataset_name.strip())
+                #dataset['train'] = dataset['train'].select(range(1000))
+                #dataset['validation'] = dataset['validation'].select(range(100))
+                ##dataset['train'] = dataset['train'].select(range(8000))
+                ##dataset['validation'] = dataset['validation'].select(range(300))                
+                #train_size = len(dataset['train'])
+                #third_size = train_size // 3                
+                #second_third = dataset['train'].select(range(third_size, third_size*2))
+                #dataset['train'] = second_third
+                #del dataset['test']
+                #tokenized_sh_fq_dataset = dataset.map(tokenize_function, batched=True, batch_size=50, remove_columns=column_names,)
+                #dataset['train'] = concatenate_datasets([saved_dataset['train'], tokenized_sh_fq_dataset['train']])
+                #dataset['train'].save_to_disk(f'/data/{hub_id.strip()}_train_dataset2')
+                #dataset['validation'].save_to_disk(f'/data/{hub_id.strip()}_validation_dataset')
+                #return 'SECOND THIRD LOADED'
                 
-        except Exception as e:
-            print(f"An error occurred: {str(e)}, TB: {traceback.format_exc()}")
-            dataset = load_dataset(dataset_name.strip())
-            #dataset['train'] = dataset['train'].select(range(8000))
-            dataset['train'] = dataset['train'].select(range(1000))
-            train_size = len(dataset['train'])
-            third_size = train_size // 3            
-            # Tokenize the dataset                    
-            first_third = dataset['train'].select(range(third_size))
-            dataset['train'] = first_third
-            del dataset['test']
-            del dataset['validation']
-            tokenized_first_third = dataset.map(tokenize_function, batched=True, batch_size=50, remove_columns=column_names,)
+        #except Exception as e:
+            #print(f"An error occurred: {str(e)}, TB: {traceback.format_exc()}")
+            #dataset = load_dataset(dataset_name.strip())
+            ##dataset['train'] = dataset['train'].select(range(8000))
+            #dataset['train'] = dataset['train'].select(range(1000))
+            #train_size = len(dataset['train'])
+            #third_size = train_size // 3            
+            ## Tokenize the dataset                    
+            #first_third = dataset['train'].select(range(third_size))
+            #dataset['train'] = first_third
+            #del dataset['test']
+            #del dataset['validation']
+            #tokenized_first_third = dataset.map(tokenize_function, batched=True, batch_size=50, remove_columns=column_names,)
             
-            tokenized_first_third.save_to_disk(f'/data/{hub_id.strip()}_train_dataset')                           
-            print('DONE')
-            return 'RUN AGAIN TO LOAD REST OF DATA'
+            #tokenized_first_third.save_to_disk(f'/data/{hub_id.strip()}_train_dataset')                           
+            #print('DONE')
+            #return 'RUN AGAIN TO LOAD REST OF DATA'
+        dataset = load_dataset(dataset_name.strip())
+        #dataset['train'] = dataset['train'].select(range(8000))
+        dataset['train'] = dataset['train'].select(range(1000))
+        dataset['validation'] = dataset['validatin'].select(range(100))
+        tokenized_first_third = dataset.map(tokenize_function, batched=True, batch_size=50, remove_columns=column_names,)
+        
 
+        print('DONE')
+
+        
+        data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)
+        trainer = Trainer(
+            model=model,
+            args=training_args,
+            train_dataset=dataset['train'],
+            eval_dataset=dataset['validation'],
+            compute_metrics=compute_metrics,
+            data_collator=data_collator,
+            #processing_class=tokenizer,
+        )    
+    
         # Fine-tune the model
         trainer.evaluate()
         #if os.path.exists(training_args.output_dir) and os.listdir(training_args.output_dir):