init

parse.py

@@ -123,50 +123,51 @@ def retrieve(model_name, dataset_name, attack_name, prompt_type):
     directory_path = "./db"
     md_dir = os.path.join(directory_path, model_name + "_" + shot + ".md")
     sections_dict = split_markdown_by_title(md_dir)
-
+    results = {}
     for cur_dataset in sections_dict.keys():
         if cur_dataset == dataset_name:
             dataset_dict = sections_dict[cur_dataset]
             for cur_attack in dataset_dict.keys():
+                
                 if cur_attack == attack_name:
-                    pass
 
                     if attack_name == "translation":
-                        results = dataset_dict[attack_name].split("\n")
-
-                        atk_acc = []
+                        prompts_dict = dataset_dict[attack_name].split("\n")
             
-                        for result in results:
-                            if "acc: " not in result:
+                        for prompt_summary in prompts_dict:
+                            if "acc: " not in prompt_summary:
                                 continue
-
+                            
+                            prompt = prompt_summary.split("prompt: ")[1]
+                            
                             import re
                         
                             match_atk = re.search(r'acc: (\d+\.\d+)%', result)
-                            
                             number_atk = float(match_atk.group(1))
-                            atk_acc.append(number_atk)
+                            results[prompt] = number_atk
+  
+                        sorted_results = sorted(results.items(), key=lambda item: item[1])[:6]
+      
 
-                        sorted_atk_acc = sorted(atk_acc)[:6]
+                        return sorted_results
 
-                elif title in ["bertattack", "checklist", "deepwordbug", "stresstest", "textfooler", "textbugger"]:
+                    elif attack_name in ["bertattack", "checklist", "deepwordbug", "stresstest", "textfooler", "textbugger"]:
 
-                    results = sections_dict[dataset][title].split("Original prompt: ")
-                    num = 0
-                    
+                        prompts_dict = dataset_dict[attack_name].split("\n")
+                        num = 0
+                        
 
-                    for result in results:
-                        if "Attacked prompt: " not in result:
-                            continue
-                        num += 1
-                        import re
-                        match_origin = re.search(r'Original acc: (\d+\.\d+)%', result)
-                        match_atk = re.search(r'attacked acc: (\d+\.\d+)%', result)
+                        for prompt_summary in prompts_dict:
+                            if "Attacked prompt: " not in prompt_summary:
+                                continue
+
+                            num += 1
+                            import re
+                            match_origin = re.search(r'Original acc: (\d+\.\d+)%', prompt_summary)
+                            match_atk = re.search(r'attacked acc: (\d+\.\d+)%', prompt_summary)
                         if match_origin and match_atk:
                             number_origin = float(match_origin.group(1))
                             number_atk = float(match_atk.group(1))
-                            summary[title][dataset].append((number_origin - number_atk)/number_origin)
-                            summary[title]["Avg"].append((number_origin - number_atk)/number_origin)
                     
                     # print(model_shot, dataset, title, len(summary[attack][dataset]), num)