fix

app.py

@@ -110,38 +110,66 @@ def single_inference_iit(head_img, head_id, tail_img, tail_id, question_txt, que
 
 
 def single_inference_tti(head_txt, head_id, tail_txt, tail_id, question_img, question_id):
-    # # (T, T) -> (I, ?)
-    # head_ent_text, tail_ent_text = ent2description[head_id], ent2description[tail_id]
+    # (T, T) -> (I, ?)
+    head_ent_text, tail_ent_text = ent2description[head_id], ent2description[tail_id]
 
-    # inputs = tokenizer(
-    #           tokenizer.sep_token.join([analogy_ent2token[head_id] + " " + head_ent_text, "[R] ", analogy_ent2token[tail_id] + " " + tail_ent_text]),
-    #           tokenizer.sep_token.join([analogy_ent2token[question_id] + " ", "[R] ", "[MASK]"]),
-    #           truncation="longest_first", max_length=128, padding="longest", return_tensors='pt', add_special_tokens=True)
-    # sep_idx = [[i for i, ids in enumerate(input_ids) if ids == tokenizer.sep_token_id] for input_ids in inputs['input_ids']]
-    # inputs['sep_idx'] = torch.tensor(sep_idx)
-    # inputs['attention_mask'] = inputs['attention_mask'].unsqueeze(1).expand([inputs['input_ids'].size(0), inputs['input_ids'].size(1), inputs['input_ids'].size(1)]).clone()
-    # for i, idx in enumerate(sep_idx):
-    #     inputs['attention_mask'][i, :idx[2], idx[2]:] = 0
+    inputs = tokenizer(
+              tokenizer.sep_token.join([analogy_ent2token[head_id] + " " + head_ent_text, "[R] ", analogy_ent2token[tail_id] + " " + tail_ent_text]),
+              tokenizer.sep_token.join([analogy_ent2token[question_id] + " ", "[R] ", "[MASK]"]),
+              truncation="longest_first", max_length=128, padding="longest", return_tensors='pt', add_special_tokens=True)
+    sep_idx = [[i for i, ids in enumerate(input_ids) if ids == tokenizer.sep_token_id] for input_ids in inputs['input_ids']]
+    inputs['sep_idx'] = torch.tensor(sep_idx)
+    inputs['attention_mask'] = inputs['attention_mask'].unsqueeze(1).expand([inputs['input_ids'].size(0), inputs['input_ids'].size(1), inputs['input_ids'].size(1)]).clone()
+    for i, idx in enumerate(sep_idx):
+        inputs['attention_mask'][i, :idx[2], idx[2]:] = 0
     
-    # # image
-    # pixel_values = processor(images=[head_img, tail_img], return_tensors='pt')['pixel_values'].squeeze()
-    # inputs['pixel_values'] = pixel_values.unsqueeze(0)
+    # image
+    pixel_values = processor(images=[question_img], return_tensors='pt')['pixel_values'].squeeze()
+    inputs['pixel_values'] = pixel_values.unsqueeze(0)
     
-    # input_ids = inputs['input_ids']
+    input_ids = inputs['input_ids']
 
-    # model_output = mkgformer.model(**inputs, return_dict=True)
-    # logits = model_output[0].logits
-    # bsz = input_ids.shape[0]
+    model_output = mkgformer.model(**inputs, return_dict=True)
+    logits = model_output[0].logits
+    bsz = input_ids.shape[0]
     
-    # _, mask_idx = (input_ids == tokenizer.mask_token_id).nonzero(as_tuple=True)    # bsz
-    # mask_logits = logits[torch.arange(bsz), mask_idx][:, analogy_entity_ids]    # bsz, 1, entity
-    # answer = ent2text[list(analogy_ent2token.keys())[mask_logits.argmax().item()]]
+    _, mask_idx = (input_ids == tokenizer.mask_token_id).nonzero(as_tuple=True)    # bsz
+    mask_logits = logits[torch.arange(bsz), mask_idx][:, analogy_entity_ids]    # bsz, 1, entity
+    answer = ent2text[list(analogy_ent2token.keys())[mask_logits.argmax().item()]]
     
     return answer
 
 
 def blended_inference_iti(head_img, head_id, tail_txt, tail_id, question_img, question_id):
-    return tail_txt
+    # (I, T) -> (I, ?)
+    head_ent_text, tail_ent_text = ent2description[head_id], ent2description[tail_id]
+
+    inputs = tokenizer(
+              tokenizer.sep_token.join([analogy_ent2token[head_id], "[R] ", analogy_ent2token[tail_id] + " " + tail_ent_text]),
+              tokenizer.sep_token.join([analogy_ent2token[question_id] + " ", "[R] ", "[MASK]"]),
+              truncation="longest_first", max_length=128, padding="longest", return_tensors='pt', add_special_tokens=True)
+    sep_idx = [[i for i, ids in enumerate(input_ids) if ids == tokenizer.sep_token_id] for input_ids in inputs['input_ids']]
+    inputs['sep_idx'] = torch.tensor(sep_idx)
+    inputs['attention_mask'] = inputs['attention_mask'].unsqueeze(1).expand([inputs['input_ids'].size(0), inputs['input_ids'].size(1), inputs['input_ids'].size(1)]).clone()
+    for i, idx in enumerate(sep_idx):
+        inputs['attention_mask'][i, :idx[2], idx[2]:] = 0
+    
+    # image
+    pixel_values = processor(images=[head_img, question_img], return_tensors='pt')['pixel_values'].squeeze()
+    inputs['pixel_values'] = pixel_values.unsqueeze(0)
+    
+    input_ids = inputs['input_ids']
+
+    model_output = mkgformer.model(**inputs, return_dict=True)
+    logits = model_output[0].logits
+    bsz = input_ids.shape[0]
+    
+    _, mask_idx = (input_ids == tokenizer.mask_token_id).nonzero(as_tuple=True)    # bsz
+    mask_logits = logits[torch.arange(bsz), mask_idx][:, analogy_entity_ids]    # bsz, 1, entity
+    answer = ent2text[list(analogy_ent2token.keys())[mask_logits.argmax().item()]]
+    
+    return answer
+
 
 def single_tab_iit():
     with gr.Column():
@@ -192,7 +220,7 @@ def single_tab_tti():
     submit_btn = gr.Button("Submit")
     output_text = gr.Textbox(label="Output")
     
-    submit_btn.click(fn=single_inference_iit, 
+    submit_btn.click(fn=single_inference_tti, 
         inputs=[head_text, head_ent, tail_text, tail_ent, question_image, question_ent], 
         outputs=[output_text])
         
@@ -223,7 +251,7 @@ def blended_tab_iti():
     submit_btn = gr.Button("Submit")
     output_text = gr.Textbox(label="Output")
     
-    submit_btn.click(fn=single_inference_iit, 
+    submit_btn.click(fn=blended_inference_iti, 
         inputs=[head_image, head_ent, tail_txt, tail_ent, question_image, question_ent], 
         outputs=[output_text])