sever

server.py

@@ -11,6 +11,7 @@ import networkx as nx
 import spacy
 # os.system("python -m spacy download en-core-web-sm")
 import pickle as pkl
+from tqdm import tqdm
 #%%
 # please use torch.load with map_location=torch.device('cpu') to map your storages to the CPU.
 # torch.loa
@@ -37,13 +38,13 @@ parser.add_argument('--init-mode', type = str, default='single', help = 'How to
 args = parser.parse_args()
 args = utils.set_hyperparams(args)
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-# device = torch.device("cpu")    
+# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+device = torch.device("cpu")    
 args.device = device
 args.device1 = device
-if torch.cuda.device_count() >= 2:
-    args.device = "cuda:0"
-    args.device1 = "cuda:1"
+# if torch.cuda.device_count() >= 2:
+#     args.device = "cuda:0"
+#     args.device1 = "cuda:1"
 
 utils.seed_all(args.seed)
 np.set_printoptions(precision=5)
@@ -77,6 +78,43 @@ with open(Parameters.GNBRfile+'raw_text_of_each_sentence', 'rb') as fl:
 with open(Parameters.UMLSfile+'drug_term', 'rb') as fl:
     drug_term = pkl.load(fl)
 
+gallery_specific_target_path = os.path.join(data_path, 'DD_target_distmult_GNBR_random_50_exists:False_single.txt')
+gallery_specific_link_path = 'DiseaseSpecific/attack_results/GNBR/cos_distmult_random_50_exists:False_20_quadratic_single_0.5.txt' 
+gallery_specific_text_path = 'DiseaseSpecific/generate_abstract/random_0.5_bioBART_finetune.json'
+gallery_agnostic_target_path = 'DiseaseAgnostic/processed_data/target_0.7random.pkl'
+gallery_agnostic_link_path = 'DiseaseAgnostic/processed_data/attack_edge_distmult_0.7random.pkl'
+gallery_agnostic_text_path = 'DiseaseAgnostic/generate_abstract/random0.7_bioBART_finetune.json'
+gallery_specific_target = utils.load_data(gallery_specific_target_path, drop=False)
+gallery_specific_link = utils.load_data(gallery_specific_link_path, drop=False)
+with open(gallery_specific_text_path, 'r') as fl:
+    gallery_specific_text = json.load(fl)
+with open(gallery_agnostic_target_path, 'rb') as fl:
+    gallery_agnostic_target = pkl.load(fl)
+with open(gallery_agnostic_link_path, 'rb') as fl:
+    gallery_agnostic_link = pkl.load(fl)
+with open(gallery_agnostic_text_path, 'r') as fl:
+    gallery_agnostic_text = json.load(fl)
+
+gallery_specific_list = []
+gallery_specific_target_dict = {}
+for i, (s, r, o) in enumerate(gallery_specific_target):
+    s = id_to_meshid[str(s)]
+    o = id_to_meshid[str(o)]
+    target_name = f'{capitalize_the_first_letter(entity_raw_name[s])} - {capitalize_the_first_letter(entity_raw_name[o])}'
+    if target_name not in gallery_specific_target_dict:
+        gallery_specific_target_dict[target_name] = i
+        gallery_specific_list.append(target_name)
+gallery_specific_list.sort()
+
+gallery_agnostic_list = []
+gallery_agnostic_target_dict = {}
+
+for i, iid in enumerate(gallery_agnostic_target):
+    target_name = capitalize_the_first_letter(entity_raw_name[id_to_meshid[str(iid)]])
+    if target_name not in gallery_agnostic_target_dict:
+        gallery_agnostic_target_dict[target_name] = i
+        gallery_agnostic_list.append(target_name)
+gallery_agnostic_list.sort()
 drug_dict = {}
 disease_dict = {}
 for k, v in entity_raw_name.items():
@@ -193,7 +231,7 @@ def tune_chatgpt(draft, attack_data, dpath):
 
         batch_size = 8
         Outs = []
-        for l in range(0, len(Text), batch_size):
+        for l in tqdm(range(0, len(Text), batch_size)):
             R = min(len(Text), l + batch_size)
             A = bart_tokenizer(Text[l:R],
             truncation = True,
@@ -211,8 +249,8 @@ def tune_chatgpt(draft, attack_data, dpath):
 def score_and_select(s, r, o, span , prompt , sen_list, BART_in, Assist, dpath, v):
 
     criterion = CrossEntropyLoss(reduction="none")
-    text_s = entity_raw_name[id_to_meshid[s]]
-    text_o = entity_raw_name[id_to_meshid[o]]
+    text_s = entity_raw_name[id_to_meshid[str(s)]]
+    text_o = entity_raw_name[id_to_meshid[str(o)]]
 
     sen_list = [server_utils.process(text) for text in sen_list]
     path_text = dpath[0].replace('\n', '')
@@ -290,7 +328,7 @@ def score_and_select(s, r, o, span , prompt , sen_list, BART_in, Assist, dpath,
     attention_mask = tokens['attention_mask'].to(args.device1)
     L = len(sen_list)
     ret_log_L = []
-    for l in range(0, L, 5):
+    for l in tqdm(range(0, L, 5)):
         R = min(L, l + 5)
         target = target_ids[l:R, :]
         attention = attention_mask[l:R, :]
@@ -380,7 +418,7 @@ def generate_template_for_triplet(attack_data):
     L = len(candidate_text_sen)
     assert L > 0
     ret_log_L = []
-    for l in range(0, L, GPT_batch_size):
+    for l in tqdm(range(0, L, GPT_batch_size)):
         R = min(L, l + GPT_batch_size)
         target = target_ids[l:R, :]
         attention = attention_mask[l:R, :]
@@ -399,10 +437,14 @@ def generate_template_for_triplet(attack_data):
     ret_log_L = list(torch.cat(ret_log_L, -1).cpu().numpy())
     sen_score = list(zip(candidate_text_sen, ret_log_L, candidate_ori_sen, Dp_path, candidate_parse_sen))
     sen_score.sort(key = lambda x: x[1])
-    test_text.append(sen_score[0][2])
-    test_dp.append(sen_score[0][3])
-    test_parse.append(sen_score[0][4])
-    single_sentence.append(sen_score[0][0])
+    Len = len(sen_score)
+    p = 0
+    if Len > 10:
+        p = np.random.choice(np.array(range(Len // 10)), 1)[0]
+    test_text.append(sen_score[p][2])
+    test_dp.append(sen_score[p][3])
+    test_parse.append(sen_score[p][4])
+    single_sentence.append(sen_score[p][0])
 
     gpt_model.to('cpu')
     return single_sentence, test_text, test_dp, test_parse
@@ -478,7 +520,7 @@ sorted_rank['merged'] = sorted_rank['merged'][llen * 3 // 4 : ]
 
 def generate_specific_attack_edge(start_entity, end_entity):
 
-    if not torch.cuda.is_available():
+    if device == torch.device('cpu'):
         print('We can just set the malicious link equals to the target link, since the generation of malicious link is too slow on cpu')
         return entity_to_id[drug_dict[start_entity]], '10', entity_to_id[disease_dict[end_entity]]
     global specific_model
@@ -649,6 +691,26 @@ def agnostic_func(agnostic_entity):
     text = score_and_select(s, r, o, span , prompt , sen_list, BART_in, Assist, dpath, {'in':single_sentence[0], 'out': draft})
     return f'{capitalize_the_first_letter(s_name)} - {capitalize_the_first_letter(r_name)} - {capitalize_the_first_letter(o_name)}', server_utils.process(text)
 
+def gallery_specific_func(specific_target):
+    index = gallery_specific_target_dict[specific_target]
+    s, r, o = gallery_specific_link[index]
+    s_name = entity_raw_name[id_to_entity[str(s)]]
+    r_name = Parameters.edge_id_to_type[int(r)].split(':')[1]
+    o_name = entity_raw_name[id_to_entity[str(o)]]
+
+    k = f'{s}_{r}_{o}_{index}'
+    text = gallery_specific_text[k]['out']
+    return f'{capitalize_the_first_letter(s_name)} - {capitalize_the_first_letter(r_name)} - {capitalize_the_first_letter(o_name)}', server_utils.process(text)
+
+def gallery_agnostic_func(agnostic_target):
+    index = gallery_agnostic_target_dict[agnostic_target]
+    s, r, o = gallery_agnostic_link[index]
+    s_name = entity_raw_name[id_to_entity[str(s)]]
+    r_name = Parameters.edge_id_to_type[int(r)].split(':')[1]
+    o_name = entity_raw_name[id_to_entity[str(o)]]
+    k = f'{s}_{r}_{o}_{index}'
+    text = gallery_agnostic_text[k]['out']
+    return f'{capitalize_the_first_letter(s_name)} - {capitalize_the_first_letter(r_name)} - {capitalize_the_first_letter(o_name)}', server_utils.process(text)
 #%%
 with gr.Blocks() as demo:
 
@@ -660,15 +722,26 @@ with gr.Blocks() as demo:
             # Center
             with gr.Column():
                 gr.Markdown("Select your poison target")
-                with gr.Tab('Target specific'):
-                    with gr.Column():
-                        with gr.Row():
-                            start_entity = gr.Dropdown(drug_list, label="Promoting drug")
-                            end_entity = gr.Dropdown(disease_list, label="Target disease")
-                        specific_generation_button = gr.Button('Poison!')
-                with gr.Tab('Target agnostic'):
-                    agnostic_entity = gr.Dropdown(drug_list, label="Promoting drug")
-                    agnostic_generation_button = gr.Button('Poison!')
+                with gr.Tab('Gallery'):
+                    with gr.Tab('Target specific'):
+                        specific_target = gr.Dropdown(gallery_specific_list, label="Poisonging target")
+                        gallery_specific_generation_button = gr.Button('Poison!')
+                    with gr.Tab('Target agnostic'):
+                        agnostic_target = gr.Dropdown(gallery_agnostic_list, label="Poisonging target")
+                        gallery_agnostic_generation_button = gr.Button('Poison!')
+
+                with gr.Tab('Poison'):
+                    with gr.Tab('Target specific'):
+                        with gr.Column():
+                            with gr.Row():
+                                start_entity = gr.Dropdown(drug_list, label="Promoting drug")
+                                end_entity = gr.Dropdown(disease_list, label="Target disease")
+                            if device == torch.device('cpu'):
+                                gr.Markdown("Since the project is currently running on the CPU, we directly treat the malicious link as equivalent to the poisoning target, to accelerate the generation process.")
+                            specific_generation_button = gr.Button('Poison!')
+                    with gr.Tab('Target agnostic'):
+                        agnostic_entity = gr.Dropdown(drug_list, label="Promoting drug")
+                        agnostic_generation_button = gr.Button('Poison!')
             with gr.Column():
                 gr.Markdown("Malicious link")
                 malicisous_link = gr.Textbox(lines=1, label="Malicious link")
@@ -676,6 +749,8 @@ with gr.Blocks() as demo:
                 malicious_text = gr.Textbox(label="Malicious text", lines=5)
     specific_generation_button.click(specific_func, inputs=[start_entity, end_entity], outputs=[malicisous_link, malicious_text])
     agnostic_generation_button.click(agnostic_func, inputs=[agnostic_entity], outputs=[malicisous_link, malicious_text])
+    gallery_specific_generation_button.click(gallery_specific_func, inputs=[specific_target], outputs=[malicisous_link, malicious_text])
+    gallery_agnostic_generation_button.click(gallery_agnostic_func, inputs=[agnostic_target], outputs=[malicisous_link, malicious_text])
 
 # demo.launch(server_name="0.0.0.0", server_port=8000, debug=False)
 demo.launch()