Spaces:
Running
Running
| import gradio as gr | |
| import torch | |
| from torch import nn | |
| from huggingface_hub import hf_hub_download | |
| from transformers import BertModel, BertTokenizer, CLIPModel, BertConfig, CLIPConfig, CLIPProcessor | |
| from modeling_unimo import UnimoForMaskedLM | |
| def load_dict_text(path): | |
| with open(path, 'r') as f: | |
| load_data = {} | |
| lines = f.readlines() | |
| for line in lines: | |
| key, value = line.split('\t') | |
| load_data[key] = value.replace('\n', '') | |
| return load_data | |
| def load_text(path): | |
| with open(path, 'r') as f: | |
| lines = f.readlines() | |
| load_data = [] | |
| for line in lines: | |
| load_data.append(line.strip().replace('\n', '')) | |
| return load_data | |
| class MKGformerModel(nn.Module): | |
| def __init__(self, text_config, vision_config): | |
| super().__init__() | |
| self.model = UnimoForMaskedLM(text_config, vision_config) | |
| def farword(self, batch): | |
| return self.model(**batch, return_dict=True) | |
| # tokenizer | |
| tokenizer = BertTokenizer.from_pretrained('bert-base-uncased') | |
| # entity and relation | |
| ent2text = load_dict_text('./dataset/MarKG/entity2text.txt') | |
| rel2text = load_dict_text('./dataset/MarKG/relation2text.txt') | |
| analogy_entities = load_text('./dataset/MARS/analogy_entities.txt') | |
| analogy_relations = load_text('./dataset/MARS/analogy_relations.txt') | |
| ent2description = load_dict_text('./dataset/MarKG/entity2textlong.txt') | |
| text2ent = {text: ent for ent, text in ent2text.items()} | |
| ent2token = {ent: f"[ENTITY_{i}]" for i, ent in enumerate(ent2description)} | |
| rel2token = {rel: f"[RELATION_{i}]" for i, rel in enumerate(rel2text)} | |
| analogy_ent2token = {ent : f"[ENTITY_{i}]" for i, ent in enumerate(ent2description) if ent in analogy_entities} | |
| analogy_rel2token = {rel : f"[RELATION_{i}]" for i, rel in enumerate(rel2text) if rel in analogy_relations} | |
| entity_list = list(ent2token.values()) | |
| relation_list = list(rel2token.values()) | |
| analogy_ent_list = list(analogy_ent2token.values()) | |
| analogy_rel_list = list(analogy_rel2token.values()) | |
| num_added_tokens = tokenizer.add_special_tokens({'additional_special_tokens': entity_list}) | |
| num_added_tokens = tokenizer.add_special_tokens({'additional_special_tokens': relation_list}) | |
| vocab = tokenizer.get_added_vocab() # dict: word: idx | |
| relation_id_st = vocab[relation_list[0]] | |
| relation_id_ed = vocab[relation_list[-1]] + 1 | |
| entity_id_st = vocab[entity_list[0]] | |
| entity_id_ed = vocab[entity_list[-1]] + 1 | |
| # analogy entities and relations | |
| analogy_entity_ids = [vocab[ent] for ent in analogy_ent_list] | |
| analogy_relation_ids = [vocab[rel] for rel in analogy_rel_list] | |
| num_added_tokens = tokenizer.add_special_tokens({'additional_special_tokens': ["[R]"]}) | |
| # model | |
| checkpoint_path = hf_hub_download(repo_id='flow3rdown/mkgformer_mart_ft', filename="mkgformer_mart_ft", repo_type='model') | |
| clip_config = CLIPConfig.from_pretrained('openai/clip-vit-base-patch32').vision_config | |
| clip_config.device = 'cpu' | |
| bert_config = BertConfig.from_pretrained('bert-base-uncased') | |
| mkgformer = MKGformerModel(clip_config, bert_config) | |
| mkgformer.model.resize_token_embeddings(len(tokenizer)) | |
| mkgformer.load_state_dict(torch.load(checkpoint_path, map_location='cpu')["state_dict"]) | |
| # processor | |
| 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, ?) | |
| ques_ent_text = ent2description[question_id] | |
| inputs = tokenizer( | |
| 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) | |
| 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 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] | |
| 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=question_img, return_tensors='pt')['pixel_values'].unsqueeze(1) | |
| pixel_values = torch.cat((pixel_values, torch.zeros_like(pixel_values)), dim=1) | |
| inputs['pixel_values'] = pixel_values | |
| 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): | |
| # (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(): | |
| gr.Markdown(""" $(I_h, I_t) : (T_q, ?)$ | |
| """) | |
| with gr.Row(): | |
| with gr.Column(): | |
| head_image = gr.Image(type='pil', label="Head Image") | |
| head_ent = gr.Textbox(lines=1, label="Head Entity") | |
| with gr.Column(): | |
| tail_image = gr.Image(type='pil', label="Tail Image") | |
| tail_ent = gr.Textbox(lines=1, label="Tail Entity") | |
| with gr.Column(): | |
| question_text = gr.Textbox(lines=1, label="Question Name") | |
| question_ent = gr.Textbox(lines=1, label="Question Entity") | |
| submit_btn = gr.Button("Submit") | |
| output_text = gr.Textbox(label="Output") | |
| submit_btn.click(fn=single_inference_iit, | |
| inputs=[head_image, head_ent, tail_image, tail_ent, question_text, question_ent], | |
| outputs=[output_text]) | |
| examples=[['examples/tree.jpg', 'Q10884', 'examples/forest.jpg', 'Q4421', "Anhui", 'Q40956']] | |
| ex = gr.Examples( | |
| examples=examples, | |
| fn=single_inference_iit, | |
| inputs=[head_image, head_ent, tail_image, tail_ent, question_text, question_ent], | |
| outputs=[output_text], | |
| cache_examples=False, | |
| run_on_click=False | |
| ) | |
| def single_tab_tti(): | |
| with gr.Column(): | |
| gr.Markdown(""" $(T_h, T_t) : (I_q, ?)$ | |
| """) | |
| with gr.Row(): | |
| with gr.Column(): | |
| head_text = gr.Textbox(lines=1, label="Head Name") | |
| head_ent = gr.Textbox(lines=1, label="Head Entity") | |
| with gr.Column(): | |
| tail_text = gr.Textbox(lines=1, label="Tail Name") | |
| tail_ent = gr.Textbox(lines=1, label="Tail Entity") | |
| with gr.Column(): | |
| question_image = gr.Image(type='pil', label="Question Image") | |
| question_ent = gr.Textbox(lines=1, label="Question Entity") | |
| submit_btn = gr.Button("Submit") | |
| output_text = gr.Textbox(label="Output") | |
| submit_btn.click(fn=single_inference_tti, | |
| inputs=[head_text, head_ent, tail_text, tail_ent, question_image, question_ent], | |
| outputs=[output_text]) | |
| examples=[['scrap', 'Q3217573', 'watch', 'Q178794', 'examples/root.jpg', 'Q111029']] | |
| ex = gr.Examples( | |
| examples=examples, | |
| fn=single_inference_iit, | |
| inputs=[head_text, head_ent, tail_text, tail_ent, question_image, question_ent], | |
| outputs=[output_text], | |
| cache_examples=False, | |
| run_on_click=False | |
| ) | |
| def blended_tab_iti(): | |
| with gr.Column(): | |
| gr.Markdown(""" $(I_h, T_t) : (I_q, ?)$ | |
| """) | |
| with gr.Row(): | |
| with gr.Column(): | |
| head_image = gr.Image(type='pil', label="Head Image") | |
| head_ent = gr.Textbox(lines=1, label="Head Entity") | |
| with gr.Column(): | |
| tail_txt = gr.Textbox(lines=1, label="Tail Name") | |
| tail_ent = gr.Textbox(lines=1, label="Tail Entity") | |
| with gr.Column(): | |
| question_image = gr.Image(type='pil', label="Question Image") | |
| question_ent = gr.Textbox(lines=1, label="Question Entity") | |
| submit_btn = gr.Button("Submit") | |
| output_text = gr.Textbox(label="Output") | |
| submit_btn.click(fn=blended_inference_iti, | |
| inputs=[head_image, head_ent, tail_txt, tail_ent, question_image, question_ent], | |
| outputs=[output_text]) | |
| examples=[['examples/watermelon.jpg', 'Q38645', 'fruit', 'Q3314483', 'examples/coffee.jpeg', 'Q8486']] | |
| ex = gr.Examples( | |
| examples=examples, | |
| fn=single_inference_iit, | |
| inputs=[head_image, head_ent, tail_txt, tail_ent, question_image, question_ent], | |
| outputs=[output_text], | |
| cache_examples=False, | |
| run_on_click=False | |
| ) | |
| TITLE = """MKG Analogy""" | |
| with gr.Blocks() as block: | |
| with gr.Column(elem_id="col-container"): | |
| gr.HTML(TITLE) | |
| with gr.Tab("Single Analogical Reasoning"): | |
| single_tab_iit() | |
| single_tab_tti() | |
| with gr.Tab("Blended Analogical Reasoning"): | |
| blended_tab_iti() | |
| # gr.HTML(ARTICLE) | |
| block.queue(max_size=64).launch(enable_queue=True) |