fix

app.py

@@ -80,11 +80,11 @@ processor = CLIPProcessor.from_pretrained('openai/clip-vit-base-patch32')
 
 def single_inference_iit(head_img, head_id, tail_img, tail_id, question_txt, question_id):
     # (I, I) -> (T, ?)
-    head_ent_text, tail_ent_text = ent2description[head_id], ent2description[tail_id]
+    ques_ent_text = ent2description[question_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]"]),
+              tokenizer.sep_token.join([analogy_ent2token[head_id] + " ", "[R] ", analogy_ent2token[tail_id] + " "]),
+              tokenizer.sep_token.join([analogy_ent2token[question_id] + " " + ques_ent_text, "[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)
@@ -108,8 +108,37 @@ def single_inference_iit(head_img, head_id, tail_img, tail_id, question_txt, que
     
     return answer
 
+
 def single_inference_tti(head_txt, head_id, tail_txt, tail_id, question_img, question_id):
-    return head_txt
+    # # (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
+    
+    # # image
+    # pixel_values = processor(images=[head_img, tail_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 blended_inference_iti(head_img, head_id, tail_txt, tail_id, question_img, question_id):
     return tail_txt
@@ -143,7 +172,7 @@ def single_tab_iit():
         inputs=[head_image, head_ent, tail_image, tail_ent, question_text, question_ent], 
         outputs=[output_text],
         cache_examples=False, 
-        run_on_click=False
+        run_on_click=True
     )
    
 def single_tab_tti():    
@@ -174,7 +203,7 @@ def single_tab_tti():
         inputs=[head_text, head_ent, tail_text, tail_ent, question_image, question_ent], 
         outputs=[output_text],
         cache_examples=False, 
-        run_on_click=False
+        run_on_click=True
     )
     
 def blended_tab_iti():
@@ -205,7 +234,7 @@ def blended_tab_iti():
         inputs=[head_image, head_ent, tail_txt, tail_ent, question_image, question_ent], 
         outputs=[output_text],
         cache_examples=False, 
-        run_on_click=False
+        run_on_click=True
     )