Update app.py

app.py

@@ -5,7 +5,6 @@ import jsonlines
 import torch
 from src.modeling_bert import EXBertForMaskedLM
 from higher.patch import monkeypatch as make_functional
-# from src.models.one_shot_learner import OneShotLearner
 
 ### load KGE model
 edit_origin_model = BertForMaskedLM.from_pretrained(pretrained_model_name_or_path="ChancesYuan/KGEditor_Edit_Test")
@@ -23,7 +22,6 @@ id2ent_name = defaultdict(str)
 rel_name2id = defaultdict(str)
 id2ent_text = defaultdict(str)
 id2rel_text = defaultdict(str)
-corrupt_triple = defaultdict(list)
 
 ### init tokenizer
 tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
@@ -34,6 +32,7 @@ def init_triple_input():
     global ent2id
     global id2ent
     global rel2token
+    global rel2id
 
     with open("./dataset/fb15k237/relations.txt", "r") as f:
         lines = f.readlines()
@@ -65,10 +64,10 @@ def init_triple_input():
         ent2id = {ent: i for i, ent in enumerate(entities)}
         id2ent = {i: ent for i, ent in enumerate(entities)}
 
-    with jsonlines.open("./dataset/fb15k237/edit_test.jsonl") as f:
-        lines = []
-        for d in f:
-            corrupt_triple[" ".join(d["ori"])] = d["cor"]
+    rel2id = {
+        w: i + len(entities)
+        for i, w in enumerate(rel2token.keys())
+    }
 
 def solve(triple, alter_label, edit_task):
     print(triple, alter_label)
@@ -77,13 +76,12 @@ def solve(triple, alter_label, edit_task):
         text_a = "[MASK]"
         text_b = id2rel_text[r] + " " + rel2token[r]
         text_c = ent2token[ent_name2id[t]] + " " + id2ent_text[ent_name2id[t]]
-        origin_label = corrupt_triple[" ".join([ent_name2id[alter_label], r, ent_name2id[t]])][0] if edit_task else ent_name2id[alter_label]
+        replace_token = [rel2id[r], ent2id[ent_name2id[t]]]
     else:
         text_a = ent2token[ent_name2id[h]]
-        # text_b = id2rel_text[r] + "[PAD]"
         text_b = id2rel_text[r] + " " + rel2token[r]
         text_c = "[MASK]" + " " + id2ent_text[ent_name2id[h]]
-        origin_label = corrupt_triple[" ".join([ent_name2id[h], r, ent_name2id[alter_label]])][2] if edit_task else ent_name2id[alter_label]
+        replace_token = [ent2id[ent_name2id[h]], rel2id[r]]
 
     if text_a == "[MASK]":
         input_text_a = tokenizer.sep_token.join(["[MASK]", id2rel_text[r] + "[PAD]"])
@@ -91,12 +89,6 @@ def solve(triple, alter_label, edit_task):
     else:
         input_text_a = "[PAD] "
         input_text_b = tokenizer.sep_token.join([id2rel_text[r] + "[PAD]", "[MASK]" + " " + id2ent_text[ent_name2id[h]]])
-    
-    cond_inputs_text = "{} >> {} || {}".format(
-        add_tokenizer.added_tokens_decoder[ent2id[origin_label] + len(tokenizer)],
-        add_tokenizer.added_tokens_decoder[ent2id[ent_name2id[alter_label]] + len(tokenizer)],
-        input_text_a + input_text_b
-    )
 
     inputs = tokenizer(
         f"{text_a} [SEP] {text_b} [SEP] {text_c}",
@@ -115,14 +107,6 @@ def solve(triple, alter_label, edit_task):
         add_special_tokens=True,
     )
 
-    cond_inputs = tokenizer(
-        cond_inputs_text,
-        truncation=True,
-        max_length=64,
-        padding="max_length",
-        add_special_tokens=True,
-    )
-
     inputs = {
         "input_ids": torch.tensor(inputs["input_ids"]).unsqueeze(dim=0),
         "attention_mask": torch.tensor(inputs["attention_mask"]).unsqueeze(dim=0),
@@ -135,13 +119,46 @@ def solve(triple, alter_label, edit_task):
         "token_type_ids": torch.tensor(edit_inputs["token_type_ids"]).unsqueeze(dim=0)
     }
 
+    _, mask_idx = (inputs["input_ids"] == tokenizer.mask_token_id).nonzero(as_tuple=True)
+    logits = edit_origin_model(**inputs).logits[:, :, 30522:45473].squeeze() if edit_task else add_origin_model(**inputs).logits[:, :, 30522:45473].squeeze()
+    logits = logits[mask_idx, :]
+
+    ### origin output
+    _, origin_entity_order = torch.sort(logits, dim=1, descending=True)
+    origin_entity_order = origin_entity_order.squeeze(dim=0)
+    origin_top3 = [id2ent_name[id2ent[origin_entity_order[i].item()]] for i in range(3)]
+
+    origin_label = origin_top3[0] if edit_task else alter_label
+
+    cond_inputs_text = "{} >> {} || {}".format(
+        add_tokenizer.added_tokens_decoder[ent2id[ent_name2id[origin_label]] + len(tokenizer)],
+        add_tokenizer.added_tokens_decoder[ent2id[ent_name2id[alter_label]] + len(tokenizer)],
+        input_text_a + input_text_b
+    )
+
+    cond_inputs = tokenizer(
+        cond_inputs_text,
+        truncation=True,
+        max_length=64,
+        padding="max_length",
+        add_special_tokens=True,
+    )
+
     cond_inputs = {
         "input_ids": torch.tensor(cond_inputs["input_ids"]).unsqueeze(dim=0),
         "attention_mask": torch.tensor(cond_inputs["attention_mask"]).unsqueeze(dim=0),
         "token_type_ids": torch.tensor(cond_inputs["token_type_ids"]).unsqueeze(dim=0)
     }
 
-    return inputs, cond_inputs, edit_inputs
+    flag = 0
+    for idx, i in enumerate(edit_inputs["input_ids"][0, :].tolist()):
+        if i == tokenizer.pad_token_id and flag == 0:
+            edit_inputs["input_ids"][0, idx] = replace_token[0] + 30522
+            flag = 1
+        elif i == tokenizer.pad_token_id and flag != 0:
+            edit_inputs["input_ids"][0, idx] = replace_token[1] + 30522
+
+    return inputs, cond_inputs, edit_inputs, origin_top3
 
 def get_logits_orig_params_dict(inputs, cond_inputs, alter_label, ex_model, learner):
     with torch.enable_grad():
@@ -149,12 +166,7 @@ def get_logits_orig_params_dict(inputs, cond_inputs, alter_label, ex_model, lear
             input_ids=inputs["input_ids"],
             attention_mask=inputs["attention_mask"],
         ).logits
-        # print(logits.shape)
-        # logits_orig, logit_for_grad, _ = logits.split([
-        #     len(inputs["input_ids"]) - 1,
-        #     1,
-        #     0,
-        # ])
+
         input_ids = inputs['input_ids']
         _, mask_idx = (input_ids == tokenizer.mask_token_id).nonzero(as_tuple=True)
         mask_logits = logits[:, mask_idx, 30522:45473].squeeze(dim=0)
@@ -174,7 +186,6 @@ def get_logits_orig_params_dict(inputs, cond_inputs, alter_label, ex_model, lear
         for (name, _), grad in zip(ex_model.named_parameters(), grads)
     }
 
-    # cond_inputs里面有pad
     params_dict = learner(
         cond_inputs["input_ids"][-1:],
         cond_inputs["attention_mask"][-1:],
@@ -184,30 +195,22 @@ def get_logits_orig_params_dict(inputs, cond_inputs, alter_label, ex_model, lear
     return params_dict
 
 def edit_process(edit_input, alter_label):
-    inputs, cond_inputs, edit_inputs = solve(edit_input, alter_label, edit_task=True)
-
-    _, mask_idx = (inputs["input_ids"] == tokenizer.mask_token_id).nonzero(as_tuple=True)
-    logits = edit_origin_model(**inputs).logits[:, :, 30522:45473].squeeze()
-    logits = logits[mask_idx, :]
-
-    ### origin output
-    _, origin_entity_order = torch.sort(logits, dim=1, descending=True)
-    origin_entity_order = origin_entity_order.squeeze(dim=0)
-    origin_top3 = [id2ent_name[id2ent[origin_entity_order[i].item()]] for i in range(3)]
+    _, cond_inputs, edit_inputs, origin_top3 = solve(edit_input, alter_label, edit_task=True)
 
     ### edit output
     fmodel = make_functional(edit_ex_model).eval()
-    params_dict = get_logits_orig_params_dict(inputs, cond_inputs, ent2id[ent_name2id[alter_label]], edit_ex_model, edit_learner)
+    params_dict = get_logits_orig_params_dict(edit_inputs, cond_inputs, ent2id[ent_name2id[alter_label]], edit_ex_model, edit_learner)
     edit_logits = fmodel(
-        input_ids=inputs["input_ids"],
-        attention_mask=inputs["attention_mask"],
+        input_ids=edit_inputs["input_ids"],
+        attention_mask=edit_inputs["attention_mask"],
         # add delta theta
         params=[
             params_dict.get(n, 0) + p
             for n, p in edit_ex_model.named_parameters()
         ],
     ).logits[:, :, 30522:45473].squeeze()
-    
+
+    _, mask_idx = (edit_inputs["input_ids"] == tokenizer.mask_token_id).nonzero(as_tuple=True)
     edit_logits = edit_logits[mask_idx, :]
     _, edit_entity_order = torch.sort(edit_logits, dim=1, descending=True)
     edit_entity_order = edit_entity_order.squeeze(dim=0)
@@ -216,23 +219,14 @@ def edit_process(edit_input, alter_label):
     return "\n".join(origin_top3), "\n".join(edit_top3)
 
 def add_process(edit_input, alter_label):
-    inputs, cond_inputs, add_inputs = solve(edit_input, alter_label, edit_task=False)
-
-    _, mask_idx = (inputs["input_ids"] == tokenizer.mask_token_id).nonzero(as_tuple=True)
-    logits = add_origin_model(**inputs).logits[:, :, 30522:45473].squeeze()
-    logits = logits[mask_idx, :]
-
-    ### origin output
-    _, origin_entity_order = torch.sort(logits, dim=1, descending=True)
-    origin_entity_order = origin_entity_order.squeeze(dim=0)
-    origin_top3 = [id2ent_name[id2ent[origin_entity_order[i].item()]] for i in range(3)]
+    _, cond_inputs, add_inputs, origin_top3 = solve(edit_input, alter_label, edit_task=False)
 
     ### add output
     fmodel = make_functional(add_ex_model).eval()
-    params_dict = get_logits_orig_params_dict(inputs, cond_inputs, ent2id[ent_name2id[alter_label]], add_ex_model, add_learner)
+    params_dict = get_logits_orig_params_dict(add_inputs, cond_inputs, ent2id[ent_name2id[alter_label]], add_ex_model, add_learner)
     add_logits = fmodel(
-        input_ids=inputs["input_ids"],
-        attention_mask=inputs["attention_mask"],
+        input_ids=add_inputs["input_ids"],
+        attention_mask=add_inputs["attention_mask"],
         # add delta theta
         params=[
             params_dict.get(n, 0) + p
@@ -240,6 +234,7 @@ def add_process(edit_input, alter_label):
         ],
     ).logits[:, :, 30522:45473].squeeze()
     
+    _, mask_idx = (add_inputs["input_ids"] == tokenizer.mask_token_id).nonzero(as_tuple=True)
     add_logits = add_logits[mask_idx, :]
     _, add_entity_order = torch.sort(add_logits, dim=1, descending=True)
     add_entity_order = add_entity_order.squeeze(dim=0)
@@ -250,9 +245,6 @@ def add_process(edit_input, alter_label):
 
 with gr.Blocks() as demo:
     init_triple_input()
-    ### example
-    # edit_process("[MASK]|/people/person/profession|Jack Black", "Kellie Martin")
-    add_process("Red Skelton|/people/person/places_lived./people/place_lived/location|[MASK]", "Palm Springs")
     gr.Markdown("# KGE Editing")
 
     # 多个tab
@@ -270,7 +262,12 @@ with gr.Blocks() as demo:
                     edit_output = gr.Textbox(label="After Edit", lines=3, placeholder="")
        
             gr.Examples(
-                examples=[["[MASK]|/people/person/profession|Jack Black", "Kellie Martin"], ["Jay-Z|/people/person/spouse_s./people/marriage/type_of_union|[MASK]", "Sydney Pollack"]],
+                examples=[["[MASK]|/people/person/profession|Jack Black", "Kellie Martin"], 
+                          ["[MASK]|/people/person/nationality|United States of America", "Mark Mothersbaugh"],
+                          ["[MASK]|/people/person/gender|Male", "Iggy Pop"],
+                          ["Rachel Weisz|/people/person/nationality|[MASK]", "J.J. Abrams"],
+                          ["Jeff Goldblum|/people/person/spouse_s./people/marriage/type_of_union|[MASK]", "Sydney Pollack"],
+                          ],
                 inputs=[edit_input, alter_label],
                 outputs=[origin_output, edit_output],
                 fn=edit_process,
@@ -290,7 +287,12 @@ with gr.Blocks() as demo:
                     add_output = gr.Textbox(label="Add Results", lines=3, placeholder="")
 
             gr.Examples(
-                examples=[["Jane Wyman|/people/person/places_lived./people/place_lived/location|[MASK]", "Palm Springs"], ["Red Skelton|/people/person/places_lived./people/place_lived/location|[MASK]", "Palm Springs"]],
+                examples=[["Jane Wyman|/people/person/places_lived./people/place_lived/location|[MASK]", "Palm Springs"], 
+                          ["Darryl F. Zanuck|/people/deceased_person/place_of_death|[MASK]", "Palm Springs"],
+                          ["[MASK]|/location/location/contains|Antigua and Barbuda", "Americas"],
+                          ["Hard rock|/music/genre/artists|[MASK]", "Social Distortion"],
+                          ["[MASK]|/people/person/nationality|United States of America", "Serj Tankian"]
+                          ],
                 inputs=[add_input, inductive_entity],
                 outputs=[add_origin_output, add_output],
                 fn=add_process,