Merge branch 'main' of https://huggingface.co/YiDuo1999/random_models_9

cases_collect.py

@@ -21,12 +21,8 @@ def valid_results_collect(model_path,valid_data,task):
     torch.cuda.ipc_collect()
 #    multiprocessing.set_start_method('spawn')
     trained_model=LLM(model=model_path,gpu_memory_utilization=0.95)
-    
     start_t=time.time()
-    if task=='sql':
-        failed_cases,correct_cases=sql_evaluation(trained_model,valid_data)
-    elif task=='nli':
-        failed_cases,correct_cases=nli_evaluation(trained_model,valid_data)
+    failed_cases,correct_cases=nli_evaluation(trained_model,valid_data)
     del trained_model
     end_t=time.time()
     print('time',start_t-end_t)
@@ -34,9 +30,6 @@ def valid_results_collect(model_path,valid_data,task):
     torch.cuda.empty_cache()
     torch.cuda.ipc_collect()
     torch.cuda.synchronize()
-    #torch.cuda.synchronize()
-    #torch.cuda.empty_cache()
-    #torch.cuda.synchronize()
     time.sleep(10)
     return failed_cases,correct_cases
 def extract_answer_prediction_nli(predicted_output):
@@ -58,7 +51,6 @@ def process_batch(data_batch,trained_model,failed_cases,correct_cases):
     batch_prompts = [data['Input'] for data in data_batch]
     outputs = trained_model.generate(batch_prompts, sampling_params)
     
-    results = []
     labels=['entailment','contradiction','neutral']
     for data, output in zip(data_batch, outputs):
 #        pdb.set_trace()
@@ -70,9 +62,6 @@ def process_batch(data_batch,trained_model,failed_cases,correct_cases):
 #            pdb.set_trace()
             
             predicted_res=predicted_output
-#            print(label,predicted_output) # if 'contradiction #label_transform(data['Output'])
-#        pdb.set_trace()
- #       print(predicted_res,label,'\n')
         non_labels = [lbl for lbl in labels if lbl != label]
         if label not in predicted_res or any(non_label in predicted_res for non_label in non_labels):
             failed_cases.append((data['Input'],predicted_res,label,data))
@@ -80,69 +69,10 @@ def process_batch(data_batch,trained_model,failed_cases,correct_cases):
             correct_cases.append((data['Input'],predicted_res,label,data))
     return failed_cases,correct_cases
 def nli_evaluation(trained_model,valid_data):
-    id=0
     failed_cases=[]
     correct_cases=[]
     batch_size=500
     batched_data = [valid_data[i:i+batch_size] for i in range(0, len(valid_data), batch_size)]
     for batch in batched_data:
         failed_cases,correct_cases=process_batch(batch,trained_model,failed_cases,correct_cases)
-        
-    #for data in valid_data:
-    #    prompt=data['Input']
-    #    output=trained_model.generate(prompt, sampling_params)
-    #    predicted_output=output[0].outputs[0].text
-    #    predicted_res=extract_answer_prediction_nli(predicted_output) #$try:
-    #    #    predicted_res=extract_answer(predicted_output.split('final')[-1].split('is')[1].split('.')[0])
-        #except:
-        #    predicted_res=extract_answer(predicted_output.split('is')[-1])
-   #     label=extract_answer(data['Output'].split('is')[-1])
-   #     print(label,predicted_res)
-   #     if not predicted_res:
-   #         pdb.set_trace()
-   #         predicted_res=''
-       # if 'contradiction #label_transform(data['Output'])
-#        pdb.set_trace()
-   #     if label not in predicted_res:
-   #         failed_cases.append((id,prompt,predicted_res,label,data))
-   #     else:
-   #         correct_cases.append((id,prompt,predicted_res,label,data))
-   #     id+=1
-    #id,prompt,prior_pred+predicted_sql,valid_data[id],ground_truth,predicted_res,ground_truth_res
-    return failed_cases,correct_cases
-def sql_evaluation(trained_model,valid_data):
-    id=0
-    failed_cases=[]
-    correct_cases=[]
-    for triple in valid_data:
-        
-        db_id,prompt,ground_truth=triple
-        prompt=prompt.replace('SELECT','')
-        db_path='/dccstor/obsidian_llm/yiduo/AgentBench/DAMO-ConvAI/bird/data/train/train_databases/{0}/{0}.sqlite'.format(db_id)
-        prompt+=' To generate the SQL query to' #print(db_path) #pdb.set_trace()
-        conn = sqlite3.connect(db_path)
-        output=trained_model.generate(prompt, sampling_params) #pdb.set_trace()
-        predicted_sql = output[0].outputs[0].text
-        #pdb.set_trace()
-        prior_pred=predicted_sql.split('final SQL')[0]
-        try:
-            predicted_sql = predicted_sql.split('final SQL')[1].strip()
-        except:
-            predicted_sql = 'SELECT'+predicted_sql.split('SELECT')[1]
-        predicted_sql=predicted_sql.split(';')[0]
-        predicted_sql=predicted_sql[predicted_sql.find('SELECT'):] #[1:]
-        cursor = conn.cursor()
-    #    pdb.set_trace()
-        try:
-            cursor.execute(predicted_sql)
-            predicted_res = cursor.fetchall()
-            cursor.execute(ground_truth)
-            ground_truth_res = cursor.fetchall()
-    #print('results',predicted_res,'truth',ground_truth_res,'\n')
-            if set(predicted_res) != set(ground_truth_res):
-                failed_cases.append((id,prompt,prior_pred+predicted_sql,valid_data[id],ground_truth,predicted_res,ground_truth_res))
-            else:
-                correct_cases.append((id,prompt,prior_pred+predicted_sql,valid_data[id],ground_truth,predicted_res,ground_truth_res))
-        except Exception as e:
-            failed_cases.append((id,prompt,predicted_sql,valid_data[id],ground_truth,str(Exception)+str(e)))
-        return failed_cases,correct_cases
+    return failed_cases,correct_cases

greedy_search.py

@@ -0,0 +1,97 @@
+import os
+import shutil
+import torch
+import pdb
+from model_soups_utils import average_two_model
+from cases_collect import valid_results_collect
+def remove_folder(path):
+    if os.path.isdir(path):  # Check if the directory exists
+        shutil.rmtree(path)
+        print(f"Directory '{path}' has been removed.")
+    else:
+        print(f"Directory '{path}' does not exist.")
+def score_criteria(x):
+    return x[1]
+def compare_criteria(x,y):
+    return x<=y
+def find_best_combination(model_path,valid_data,test_examples,search_name,iteration=5,seed=True,task='nli'):
+    if seed:
+        if isinstance(model_path,list):
+            paths=[]
+            for m_p in model_path:
+                paths.extend([m_p+'_{0}'.format(seed) for seed in [str(i) for i in range(2020,2030)]])
+        else:
+            paths=[model_path+'_{0}'.format(seed) for seed in [str(i) for i in range(2020,2030)]]
+    else:
+        paths=model_path
+    try:
+        update_scores=torch.load('{0}_score.pt'.format(search_name))
+        del_paths=torch.load('{0}_path.pt'.format(search_name))
+        for path in del_paths:
+            del paths[paths.index(path)]
+        best_path=torch.load('{0}_best_path.pt'.format(search_name))
+        best_score=update_scores[-1]
+    except:
+        del_paths=[]
+        update_scores=[]
+        path_count=[]
+        for path_id,path in enumerate(paths):
+            print(0,path_id,len(paths))
+            f_test,c_test=valid_results_collect(path, test_examples,task) #test_examples, args.task)
+            path_count.append((path,len(c_test)/(len(f_test)+len(c_test)))) #ooa_failed_cases, im_failed_cases, correct_cases=process_nli_validation_batch(path, valid_data,seed=False, iteration=5)
+        print(path_count[-1][1]) #path_count.append((path,len(ooa_failed_cases),len(im_failed_cases),len(correct_cases)))
+        path_count.sort(key=lambda x:score_criteria(x),reverse=True)
+        best_path=path_count[0][0]
+        best_score=score_criteria(path_count[0])
+        update_scores.append(best_score)
+        f_test,c_test=valid_results_collect(best_path, test_examples,'nli') 
+        print(best_score,len(c_test)/(len(f_test)+len(c_test)))
+        torch.save(update_scores,'{0}_score.pt'.format(search_name)) #torch.save(update_scores,'update_scores_backup.pt') 
+        del_paths.append(best_path) 
+        torch.save(del_paths,'{0}_path.pt'.format(search_name)) #torch.save(update_scores,'{0}_score.pt'.format(search_name)) #del_paths=torch.load('{0}_path.pt'.format(search_name)) #pdb.set_trace()
+        del paths[paths.index(best_path)]
+        torch.save(best_path,'{0}_best_path.pt'.format(search_name))
+    while len(paths)>0:
+        path_count=[]
+        for path_id,path in enumerate(paths):
+            print(len(update_scores),path_id,len(paths))
+            average_path="{0}_average".format(best_path+path.split('/')[-1])
+            if not os.path.isdir(average_path):
+                average_path=average_two_model(best_path,path,len(update_scores))
+            f_test,c_test=valid_results_collect(average_path, test_examples, 'nli') #valid_results_collect(path, valid_data,args.task) #f_test,c_test=valid_results_collect(average_path, test_examples$
+            if not path_count: #ooa_failed_cases, im_failed_cases, correct_cases=process_nli_validation_batch(average_path, valid_data,seed=False, iteration=5)
+                path_count.append((path,len(c_test)/(len(f_test)+len(c_test)),average_path))
+            else:
+                score=len(c_test)/(len(f_test)+len(c_test))
+                if score>=path_count[-1][1]:
+                    path_count.append((path,score,average_path))
+                else:
+                    remove_folder(average_path)
+            print(path_count[-1][1]) #len(ooa_failed_cases),len(im_failed_cases),len(correct_cases),average_path))
+        path_count.sort(key=lambda x:score_criteria(x),reverse=True)
+        win_path=path_count[0][0]
+        win_score=score_criteria(path_count[0])
+        #del paths[paths.index(win_path)]
+        if compare_criteria(best_score,win_score):
+            if len(del_paths)>2: 
+                remove_folder(best_path)
+            best_path=path_count[0][2]
+            torch.save(best_path,'{0}_best_path.pt'.format(search_name))
+            best_score=win_score
+            #f_test,c_test=valid_results_collect(best_path, test_examples,args.task) 
+            print(best_score) #,len(c_test)/(len(f_test)+len(c_test)))
+            del paths[paths.index(win_path)] #print(best_score)
+            del_paths.append(win_path) 
+            torch.save(del_paths,'{0}_path.pt'.format(search_name))
+   #         pdb.set_trace()
+            update_scores.append(best_score)
+            torch.save(update_scores,'{0}_score.pt'.format(search_name)) #torch.save(update_scores,'update_scores_backup.pt') 
+        else:
+            while paths:
+                paths.pop()
+            best_path=best_path
+            #break
+#update_scores.append(best_score)
+    return best_path,update_scores
+
+        #ooa_failed_cases, im_failed_cases, correct_cases=process_nli_validation_batch(path, valid_data,seed=False, iteration=100)

model_soups_utils.py

@@ -0,0 +1,25 @@
+import torch
+from transformers import AutoModelForCausalLM,AutoTokenizer
+from transformers import LlamaTokenizer
+from vllm import LLM, SamplingParams
+
+def average_two_model(model_path_1,model_path_2,update_num,base_path='/dccstor/obsidian_llm/yiduo/h100_data/llama-3-8b'):
+
+# Path to save the averaged model and tokenizer
+    averaged_model_path = "{0}".format(model_path_1+model_path_2.split('/')[-1]).replace('00','').replace('random','').replace('naive_3k','').replace('shuffle','').replace('average','')
+    # Load and average the state dicts for each model
+    models=[]
+    model_paths=[model_path_1,model_path_2]
+    for model_path in model_paths:
+        models.append(AutoModelForCausalLM.from_pretrained(model_path))
+    avg_state_dict = {}
+    for key in models[0].state_dict().keys():
+        avg_state_dict[key] = (update_num/(update_num+1))*models[0].state_dict()[key]+(1.0/(update_num+1))*models[1].state_dict()[key]  #sum([model.state_dict()[key] for model in models]) / len(models)
+    base_model = AutoModelForCausalLM.from_pretrained(base_path)  # Load the base model configuration
+    base_model.load_state_dict(avg_state_dict)
+    base_model.save_pretrained(averaged_model_path) # Save the averaged model
+    # Load the tokenizer (assuming all models used the same tokenizer)
+    # If needed, adjust the tokenizer path to match the base LLaMA tokenizer used
+    tokenizer = AutoTokenizer.from_pretrained(model_path_1) #tokenizer = LlamaTokenizer.from_pretrained(model_path+'_{0}'.format(seeds[0]))
+    tokenizer.save_pretrained(averaged_model_path)
+    return averaged_model_path

run_model_soups.py

@@ -0,0 +1,111 @@
+import os
+import torch
+import logging
+from pathlib import Path
+from typing import List, Dict, Tuple
+from datasets import load_dataset
+from greedy_search import find_best_combination
+from cases_collect import valid_results_collect
+
+def setup_logger() -> logging.Logger:
+    """Configure and return logger."""
+    logging.basicConfig(
+        level=logging.INFO,
+        format='%(asctime)s - %(levelname)s - %(message)s'
+    )
+    return logging.getLogger(__name__)
+
+def get_model_paths(model_names: List[str], base_path: str = './') -> List[str]:
+    """Generate model paths from names."""
+    return [os.path.join(base_path, f"{name}_model") for name in model_names]
+
+def load_test_data(dataset_name: str = 'hippocrates/MedNLI_test') -> List[Dict]:
+    """Load and prepare test dataset."""
+    dataset = load_dataset(dataset_name)
+    return [
+        {'Input': item['query'], 'Output': item['answer']} 
+        for item in dataset['test']
+    ]
+
+def calculate_accuracy(correct: List, failed: List) -> float:
+    """Calculate accuracy from correct and failed cases."""
+    total = len(correct) + len(failed)
+    return len(correct) / total if total > 0 else 0.0
+
+def main():
+    """Main execution function."""
+    logger = setup_logger()
+    
+    try:
+        # Configuration
+        config = {
+            'search_name': 'randoms_model',
+            'model_names': ['randoms_data_3k_model'],
+            'base_path': './',
+            'valid_data_path': 'nli_demo.pt',
+            'seed': True,
+            'iteration': 5
+        }
+        
+        # Generate model paths
+        model_paths = get_model_paths(config['model_names'], config['base_path'])
+        logger.info(f"Generated model paths: {model_paths}")
+        
+        # Load datasets
+        logger.info("Loading test data...")
+        test_examples = load_test_data()
+        logger.info(f"Loaded {len(test_examples)} test examples")
+        
+        logger.info("Loading validation data...")
+        try:
+            valid_data = torch.load(config['valid_data_path'])
+            logger.info(f"Loaded validation data from {config['valid_data_path']}")
+        except Exception as e:
+            logger.error(f"Failed to load validation data: {str(e)}")
+            raise
+        
+        # Find best combination
+        logger.info("Finding best model combination...")
+        best_path, update_scores = find_best_combination(
+            model_paths,
+            valid_data,
+            valid_data,
+            config['search_name'],
+            iteration=config['iteration'],
+            seed=config['seed']
+        )
+        logger.info(f"Best path found with scores: {update_scores}")
+        
+        # Evaluate on test set
+        logger.info("Evaluating on test set...")
+        failed_test, correct_test = valid_results_collect(
+            best_path,
+            test_examples,
+            'nli'
+        )
+        
+        # Calculate and log accuracy
+        accuracy = calculate_accuracy(correct_test, failed_test)
+        logger.info(f"Test Accuracy: {accuracy:.4f}")
+        
+        # Save results
+        results = {
+            'best_path': best_path,
+            'update_scores': update_scores,
+            'test_accuracy': accuracy,
+            'test_results': {
+                'correct': len(correct_test),
+                'failed': len(failed_test)
+            }
+        }
+        
+        save_path = Path(f"results_{config['search_name']}.pt")
+        torch.save(results, save_path)
+        logger.info(f"Results saved to {save_path}")
+        
+    except Exception as e:
+        logger.error(f"Error in main execution: {str(e)}", exc_info=True)
+        raise
+
+if __name__ == "__main__":
+    main()