deve

app.py

@@ -32,33 +32,33 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
         torch.cuda.empty_cache()
         torch.nn.CrossEntropyLoss()
         rouge_metric = evaluate.load("rouge", cache_dir='/data/cache')
-        def compute_metrics(eval_preds):
-            preds, labels = eval_preds
-            if isinstance(preds, tuple):
-                preds = preds[0]
-            from pprint import pprint as pp
-            pp(preds)
-            # Replace -100s used for padding as we can't decode them
-            preds = np.where(preds != -100, preds, tokenizer.pad_token_id)
-            labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
+        #def compute_metrics(eval_preds):
+            #preds, labels = eval_preds
+            #if isinstance(preds, tuple):
+                #preds = preds[0]
+            #from pprint import pprint as pp
+            #pp(preds)
+            ## Replace -100s used for padding as we can't decode them
+            #preds = np.where(preds != -100, preds, tokenizer.pad_token_id)
+            #labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
         
-            # Decode predictions and labels
-            decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
-            decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
+            ## Decode predictions and labels
+            #decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
+            #decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
         
-            # Compute ROUGE metrics
-            result = rouge_metric.compute(predictions=decoded_preds, references=decoded_labels)
-            result = {k: round(v * 100, 4) for k, v in result.items()}
+            ## Compute ROUGE metrics
+            #result = rouge_metric.compute(predictions=decoded_preds, references=decoded_labels)
+            #result = {k: round(v * 100, 4) for k, v in result.items()}
         
-            # Calculate accuracy
-            accuracy = accuracy_score(decoded_labels, decoded_preds)
-            result["eval_accuracy"] = round(accuracy * 100, 4)
+            ## Calculate accuracy
+            #accuracy = accuracy_score(decoded_labels, decoded_preds)
+            #result["eval_accuracy"] = round(accuracy * 100, 4)
         
-            # Calculate average generation length
-            prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in preds]
-            result["gen_len"] = np.mean(prediction_lens)
+            ## Calculate average generation length
+            #prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in preds]
+            #result["gen_len"] = np.mean(prediction_lens)
         
-            return result
+            #return result
         
         login(api_key.strip())
    
@@ -135,118 +135,112 @@ 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'],
+        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,
                     #compute_metrics=compute_metrics,
                     #data_collator=data_collator,
-                    ##processing_class=tokenizer,
-                #)
+                    #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)
+                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(15000))
+                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)
+                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(15000))
+                dataset['validation'] = dataset['validation'].select(range(2000))             
+                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,)
+                saved_dataset = load_from_disk(f'/data/{hub_id.strip()}_train_dataset') 
+                dataset['train'] = concatenate_datasets([saved_dataset, 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(15000))
+            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,)
             
-            #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(4000))
-        dataset['validation'] = dataset['validation'].select(range(200))
-        train_set = dataset.map(tokenize_function, batched=True)
-        #valid_set = dataset['validation'].map(tokenize_function, batched=True)
+        #dataset['train'] = dataset['train'].select(range(4000))
+        #dataset['validation'] = dataset['validation'].select(range(200))
+        #train_set = dataset.map(tokenize_function, batched=True)
         
 
         #print(train_set.keys())
         print('DONE')
 
         
-        data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)
-        trainer = Trainer(
-            model=model,
-            args=training_args,
-            train_dataset=train_set['train'], 
-            eval_dataset=train_set['validation'], 
-            #compute_metrics=compute_metrics,
-            #data_collator=data_collator,
-            #processing_class=tokenizer,
-        )    
+        #data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)
+        #trainer = Trainer(
+            #model=model,
+            #args=training_args,
+            #train_dataset=train_set['train'], 
+            #eval_dataset=train_set['validation'], 
+            ##compute_metrics=compute_metrics,
+            ##data_collator=data_collator,
+            ##processing_class=tokenizer,
+        #)    
     
         # Fine-tune the model
-        trainer.train()
-        #if os.path.exists(training_args.output_dir) and os.listdir(training_args.output_dir):
-            #train_result = trainer.train(resume_from_checkpoint=True)
-        #else:
-            #train_result = trainer.train()
+        if os.path.exists(training_args.output_dir) and os.listdir(training_args.output_dir):
+            train_result = trainer.train(resume_from_checkpoint=True)
+        else:
+            train_result = trainer.train()
         trainer.push_to_hub(commit_message="Training complete!")
     except Exception as e:
         return f"An error occurred: {str(e)}, TB: {traceback.format_exc()}"