Spaces:
Runtime error
Runtime error
| from pathlib import Path | |
| import os | |
| import time | |
| import gradio as gr | |
| import requests | |
| import numpy as np | |
| from pathlib import Path | |
| import torch | |
| import torch.nn.functional as F | |
| import open_clip | |
| import faiss | |
| from transformers import TextIteratorStreamer | |
| from threading import Thread | |
| from conversation import default_conversation, conv_templates, Conversation | |
| from knowledge import TextDB | |
| from knowledge.transforms import five_crop, nine_crop | |
| from knowledge.utils import refine_cosine | |
| from model import get_gptk_model, get_gptk_image_transform | |
| no_change_btn = gr.Button.update() | |
| enable_btn = gr.Button.update(interactive=True) | |
| disable_btn = gr.Button.update(interactive=False) | |
| knwl_none = (None, ) * 30 | |
| knwl_unchange = (gr.Image.update(), ) * 15 + (gr.Textbox.update(), ) * 15 | |
| moderation_msg = "YOUR INPUT VIOLATES OUR CONTENT MODERATION GUIDELINES. PLEASE TRY AGAIN." | |
| def violates_moderation(text): | |
| """ | |
| Check whether the text violates OpenAI moderation API. | |
| """ | |
| if "OPENAI_API_KEY" not in os.environ: | |
| print("OPENAI_API_KEY not found, skip content moderation check...") | |
| return False | |
| url = "https://api.openai.com/v1/moderations" | |
| headers = { | |
| "Content-Type": "application/json", | |
| "Authorization": "Bearer " + os.environ["OPENAI_API_KEY"] | |
| } | |
| text = text.replace("\n", "") | |
| data = "{" + '"input": ' + f'"{text}"' + "}" | |
| data = data.encode("utf-8") | |
| try: | |
| ret = requests.post(url, headers=headers, data=data, timeout=5) | |
| flagged = ret.json()["results"][0]["flagged"] | |
| except requests.exceptions.RequestException as e: | |
| flagged = False | |
| except KeyError as e: | |
| flagged = False | |
| return flagged | |
| def load_demo(): | |
| state = default_conversation.copy() | |
| return state | |
| def regenerate(state: Conversation): | |
| state.messages[-1][-1] = None | |
| prev_human_msg = state.messages[-2] | |
| if type(prev_human_msg[1]) in (tuple, list): | |
| prev_human_msg[1] = prev_human_msg[1][:2] | |
| state.skip_next = False | |
| return (state, state.to_gradio_chatbot(), "", None, disable_btn, disable_btn, disable_btn) | |
| def clear_history(): | |
| state = default_conversation.copy() | |
| return (state, state.to_gradio_chatbot(), "", None) + (enable_btn, disable_btn, disable_btn) + knwl_none | |
| def add_text(state: Conversation, text, image): | |
| if len(text) <= 0 and image is None: | |
| state.skip_next = True | |
| return (state, state.to_gradio_chatbot(), "", None) + (no_change_btn,) * 3 | |
| if violates_moderation(text): | |
| state.skip_next = True | |
| return (state, state.to_gradio_chatbot(), moderation_msg, None) + (no_change_btn,) * 3 | |
| if image is not None: | |
| text = (text, image) | |
| if len(state.get_images(return_pil=True)) > 0: | |
| state = default_conversation.copy() | |
| state.append_message(state.roles[0], text) | |
| state.append_message(state.roles[1], None) | |
| state.skip_next = False | |
| return (state, state.to_gradio_chatbot(), "", None) + (disable_btn,) * 3 | |
| def search(image, pos, topk, knwl_db, knwl_idx): | |
| with torch.cuda.amp.autocast(): | |
| image = query_trans(image).unsqueeze(0).to(device) | |
| query = F.normalize(query_enc(image), dim=-1) | |
| query = query.cpu().numpy() | |
| _, I = knwl_idx.search(query, 4*topk) | |
| score, I = refine_cosine(knwl_db.feature, query, I, device, topk) | |
| score, I = score.flatten(), I.flatten() | |
| embd, text = knwl_db[I] | |
| pos = np.full((topk, ), fill_value=pos) | |
| query = torch.FloatTensor(query).unsqueeze(0).to(device) | |
| embd = torch.FloatTensor(embd).unsqueeze(0).to(device) | |
| pos = torch.LongTensor(pos).unsqueeze(0).to(device) | |
| score = torch.FloatTensor(score).unsqueeze(0).to(device) | |
| return query, embd, pos, score, text | |
| def retrieve_knowledge(image): | |
| knwl_embd = {} | |
| knwl_text = {} | |
| for query_type, topk_q in topk.items(): | |
| if topk_q == 0: continue | |
| if query_type == "whole": | |
| images = [image, ] | |
| knwl_text[query_type] = {i: {} for i in range(1)} | |
| elif query_type == "five": | |
| images = five_crop(image) | |
| knwl_text[query_type] = {i: {} for i in range(5)} | |
| elif query_type == "nine": | |
| images = nine_crop(image) | |
| knwl_text[query_type] = {i: {} for i in range(9)} | |
| else: | |
| raise ValueError | |
| knwl_embd[query_type] = {} | |
| for knwl_type, (knwl_db_t, knwl_idx_t) in knwl_db.items(): | |
| query, embed, pos, score = [], [], [], [] | |
| for i, img in enumerate(images): | |
| query_i, embed_i, pos_i, score_i, text_i = search( | |
| img, i, topk_q, knwl_db_t, knwl_idx_t | |
| ) | |
| query.append(query_i) | |
| embed.append(embed_i) | |
| pos.append(pos_i) | |
| score.append(score_i) | |
| knwl_text[query_type][i][knwl_type] = text_i | |
| query = torch.cat(query, dim=1) | |
| embed = torch.cat(embed, dim=1) | |
| pos = torch.cat(pos, dim=1) | |
| score = torch.cat(score, dim=1) | |
| knwl_embd[query_type][knwl_type] = { | |
| "embed": embed, "query": query, "pos": pos, "score": score | |
| } | |
| return knwl_embd, knwl_text | |
| def generate(state: Conversation, temperature, top_p, max_new_tokens): | |
| if state.skip_next: # This generate call is skipped due to invalid inputs | |
| yield (state, state.to_gradio_chatbot()) + (no_change_btn,) * 3 + knwl_unchange | |
| return | |
| if len(state.messages) == state.offset + 2: # First round of conversation | |
| new_state = conv_templates["gptk"].copy() | |
| new_state.append_message(new_state.roles[0], state.messages[-2][1]) | |
| new_state.append_message(new_state.roles[1], None) | |
| state = new_state | |
| # retrieve and visualize knowledge | |
| image = state.get_images(return_pil=True)[0] | |
| knwl_embd, knwl = retrieve_knowledge(image) | |
| knwl_img, knwl_txt, idx = [None, ] * 15, ["", ] * 15, 0 | |
| for query_type, knwl_pos in (("whole", 1), ("five", 5), ("nine", 9)): | |
| if query_type == "whole": | |
| images = [image, ] | |
| elif query_type == "five": | |
| images = five_crop(image) | |
| elif query_type == "nine": | |
| images = nine_crop(image) | |
| for pos in range(knwl_pos): | |
| try: | |
| txt = "" | |
| for k, v in knwl[query_type][pos].items(): | |
| v = ", ".join([vi.replace("_", " ") for vi in v]) | |
| txt += f"**[{k.upper()}]:** {v}\n\n" | |
| knwl_txt[idx] += txt | |
| img = images[pos] | |
| img = query_trans.transforms[0](img) | |
| img = query_trans.transforms[1](img) | |
| img = query_trans.transforms[2](img) | |
| knwl_img[idx] = img | |
| except KeyError: | |
| pass | |
| idx += 1 | |
| knwl_vis = tuple(knwl_img + knwl_txt) | |
| yield (state, state.to_gradio_chatbot()) + (disable_btn,) * 3 + knwl_vis | |
| # generate output | |
| prompt = state.get_prompt().replace("USER: <image>\n", "") | |
| prompt = prompt.split("USER:")[-1].replace("ASSISTANT:", "") | |
| image_pt = gptk_trans(image).to(device).unsqueeze(0) | |
| samples = {"image": image_pt, "knowledge": knwl_embd, "prompt": prompt} | |
| streamer = TextIteratorStreamer( | |
| gptk_model.llm_tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=15 | |
| ) | |
| thread = Thread( | |
| target=gptk_model.generate, | |
| kwargs=dict( | |
| samples=samples, | |
| use_nucleus_sampling=(temperature > 0.001), | |
| max_length=min(int(max_new_tokens), 1024), | |
| top_p=float(top_p), | |
| temperature=float(temperature), | |
| streamer=streamer, | |
| num_beams=1, | |
| length_penalty=0.0, | |
| auto_cast=True | |
| ) | |
| ) | |
| thread.start() | |
| generated_text = "" | |
| for new_text in streamer: | |
| generated_text += new_text | |
| state.messages[-1][-1] = generated_text + "▌" | |
| yield (state, state.to_gradio_chatbot()) + (disable_btn,) * 3 + knwl_unchange | |
| time.sleep(0.03) | |
| state.messages[-1][-1] = state.messages[-1][-1][:-1] | |
| yield (state, state.to_gradio_chatbot()) + (enable_btn,) * 3 + knwl_unchange | |
| title_markdown = (""" | |
| # GPT-K: Knowledge Augmented Vision-and-Language Assistant | |
| """) | |
| tos_markdown = (""" | |
| ### Terms of use | |
| By using this service, users are required to agree to the following terms: | |
| The service is a research preview intended for non-commercial use only. It only provides limited safety measures and may generate offensive content. It must not be used for any illegal, harmful, violent, racist, or sexual purposes. The service may collect user dialogue data for future research. | |
| Please click the "Flag" button if you get any inappropriate answer! We will collect those to keep improving our moderator. | |
| For an optimal experience, please use desktop computers for this demo, as mobile devices may compromise its quality. | |
| """) | |
| learn_more_markdown = (""" | |
| ### License | |
| The service is a research preview intended for non-commercial use only, subject to the model [License](https://github.com/facebookresearch/llama/blob/main/MODEL_CARD.md) of LLaMA, [Terms of Use](https://openai.com/policies/terms-of-use) of the data generated by OpenAI, and [Privacy Practices](https://chrome.google.com/webstore/detail/sharegpt-share-your-chatg/daiacboceoaocpibfodeljbdfacokfjb) of ShareGPT. Please contact us if you find any potential violation. | |
| """) | |
| def build_demo(): | |
| textbox = gr.Textbox(show_label=False, placeholder="Enter text and press ENTER", container=False) | |
| imagebox = gr.Image(type="pil") | |
| with gr.Blocks(title="GPT-K", theme=gr.themes.Base()) as demo: | |
| state = gr.State() | |
| gr.Markdown(title_markdown) | |
| with gr.Row(): | |
| with gr.Column(scale=3): | |
| gr.Examples(examples=[ | |
| ["examples/mona_lisa.jpg", "Discuss the historical impact and the significance of this painting in the art world."], | |
| ["examples/mona_lisa_dog.jpg", "Describe this photo in detail."], | |
| ["examples/horseshoe_bend.jpg", "What are the possible reasons of the formation of this sight?"], | |
| ], inputs=[imagebox, textbox]) | |
| imagebox.render() | |
| with gr.Row(): | |
| with gr.Column(scale=8): | |
| textbox.render() | |
| with gr.Column(scale=1, min_width=60): | |
| submit_btn = gr.Button(value="Submit") | |
| with gr.Row(): | |
| regenerate_btn = gr.Button(value="🔄 Regenerate", interactive=False, scale=1) | |
| clear_btn = gr.Button(value="🗑️ Clear", interactive=False, scale=1) | |
| with gr.Accordion("Parameters", open=True): | |
| temperature = gr.Slider(minimum=0.0, maximum=1.0, value=0.0, step=0.1, interactive=True, label="Temperature",) | |
| top_p = gr.Slider(minimum=0.0, maximum=1.0, value=0.7, step=0.1, interactive=True, label="Top P",) | |
| max_output_tokens = gr.Slider(minimum=0, maximum=1024, value=512, step=64, interactive=True, label="Max output tokens",) | |
| with gr.Column(scale=6): | |
| chatbot = gr.Chatbot(elem_id="chatbot", label="GPT-K Chatbot", height=550) | |
| gr.Markdown("## Retrieved Knowledge") | |
| knwl_img, knwl_txt = [], [] | |
| for query_type, knwl_pos in (("whole", 1), ("five", 5), ("nine", 9)): | |
| with gr.Tab(query_type): | |
| for p in range(knwl_pos): | |
| with gr.Tab(str(p)): | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| knwl_img.append(gr.Image(type="pil", show_label=False, interactive=False)) | |
| with gr.Column(scale=7): | |
| knwl_txt.append(gr.Markdown()) | |
| knwl_vis = knwl_img + knwl_txt | |
| gr.Markdown(tos_markdown) | |
| gr.Markdown(learn_more_markdown) | |
| # Register listeners | |
| btn_list = [submit_btn, regenerate_btn, clear_btn] | |
| regenerate_btn.click( | |
| regenerate, [state], [state, chatbot, textbox, imagebox] + btn_list | |
| ).then( | |
| generate, | |
| [state, temperature, top_p, max_output_tokens], | |
| [state, chatbot] + btn_list + knwl_vis | |
| ) | |
| clear_btn.click( | |
| clear_history, None, [state, chatbot, textbox, imagebox] + btn_list + knwl_vis | |
| ) | |
| textbox.submit( | |
| add_text, [state, textbox, imagebox], [state, chatbot, textbox, imagebox] + btn_list | |
| ).then( | |
| generate, | |
| [state, temperature, top_p, max_output_tokens], | |
| [state, chatbot] + btn_list + knwl_vis | |
| ) | |
| submit_btn.click( | |
| add_text, [state, textbox, imagebox], [state, chatbot, textbox, imagebox] + btn_list | |
| ).then( | |
| generate, | |
| [state, temperature, top_p, max_output_tokens], | |
| [state, chatbot] + btn_list + knwl_vis | |
| ) | |
| demo.load(load_demo, None, [state]) | |
| return demo | |
| def build_knowledge(): | |
| def get_knwl(knowledge_db): | |
| knwl_db = TextDB(Path(knowledge_db)/"knowledge_db.hdf5") | |
| knwl_db.feature = knwl_db.feature | |
| knwl_idx = faiss.read_index(str(Path(knowledge_db)/"faiss.index")) | |
| knwl_idx.add(knwl_db.feature.astype(np.float32)) | |
| return knwl_db, knwl_idx | |
| knwl_db = { | |
| "obj": get_knwl('knowledge/(dataset-object)(clip-model-ViT-g-14)(dbscan)(eps-0.15)(ms-1)'), | |
| "act": get_knwl('knowledge/(dataset-action)(clip-model-ViT-g-14)(dbscan)(eps-0.15)(ms-1)'), | |
| "attr": get_knwl('knowledge/(dataset-attribute)(clip-model-ViT-g-14)(dbscan)(eps-0.15)(ms-1)'), | |
| } | |
| d_knwl = knwl_db["obj"][0].feature.shape[-1] | |
| return knwl_db, d_knwl | |
| def build_query_model(): | |
| query_enc, _, query_trans = open_clip.create_model_and_transforms( | |
| "ViT-g-14", pretrained="laion2b_s34b_b88k", precision='fp16' | |
| ) | |
| query_enc = query_enc.visual.to(device).eval() | |
| return query_enc, query_trans | |
| def build_gptk_model(): | |
| _, gptk_trans = get_gptk_image_transform() | |
| topk = {"whole": 60, "five": 24, "nine": 16} | |
| gptk_model = get_gptk_model(d_knwl=d_knwl, topk=topk) | |
| gptk_ckpt = "model/ckpt/gptk-vicuna7b.pt" | |
| gptk_ckpt = torch.load(gptk_ckpt, map_location="cpu") | |
| gptk_model.load_state_dict(gptk_ckpt, strict=False) | |
| gptk_model = gptk_model.to(device).eval() | |
| return gptk_model, gptk_trans, topk | |
| if torch.cuda.is_available(): | |
| device = torch.device("cuda") | |
| else: | |
| device = torch.device("cpu") | |
| knwl_db, d_knwl = build_knowledge() | |
| gptk_model, gptk_trans, topk = build_gptk_model() | |
| query_enc, query_trans = build_query_model() | |
| demo = build_demo() | |
| demo.queue().launch() |