implement gpt-k demo

.gitattributes

@@ -33,3 +33,8 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+.bin filter=lfs diff=lfs merge=lfs -text
+.pt filter=lfs diff=lfs merge=lfs -text
+*.hdf5 filter=lfs diff=lfs merge=lfs -text
+*.index filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text

.vscode/settings.json

@@ -0,0 +1,6 @@
+{
+    "[python]": {
+        "editor.defaultFormatter": "ms-python.autopep8"
+    },
+    "python.formatting.provider": "none"
+}

README.md

@@ -1,5 +1,5 @@
 ---
-title: K GPT
+title: GPT-K
 emoji: 🚀
 colorFrom: green
 colorTo: red

app.py

@@ -0,0 +1,387 @@
+from pathlib import Path
+import datetime
+import json
+import os
+import time
+import gradio as gr
+import requests
+import numpy as np
+
+import torch
+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
+moderation_msg = "YOUR INPUT VIOLATES OUR CONTENT MODERATION GUIDELINES. PLEASE TRY AGAIN."
+
+
+def violates_moderation(text):
+    """
+    Check whether the text violates OpenAI moderation API.
+    """
+    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)
+
+
+def clear_history():
+    state = default_conversation.copy()
+    return (state, state.to_gradio_chatbot(), "", None) + (disable_btn,) * 2 + 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,) * 2
+    
+    if violates_moderation(text):
+        state.skip_next = True
+        return (state, state.to_gradio_chatbot(), moderation_msg, None) + (no_change_btn,) * 2
+
+    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,) * 2
+
+
+def search(image, pos, topk, knwl_db, knwl_idx):
+    with torch.cuda.amp.autocast():
+        image = query_trans(image).unsqueeze(0).to(device)
+        query = query_enc.encode_image(image, normalize=True)
+    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, temperature, top_p, max_new_tokens, add_knwl, do_sampling, do_beam_search):
+    if state.skip_next:  # This generate call is skipped due to invalid inputs
+        yield (state, state.to_gradio_chatbot()) + (no_change_btn,) * 2 + knwl_none
+        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]
+    if bool(add_knwl):
+        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][str(pos)].items():
+                        v = ", ".join([vi.replace("_", " ") for vi in v])
+                        txt += f"**[{k.upper()}]:** {v}\n\n"
+                    knwl_txt[idx] += txt
+                    knwl_img[idx] = images[pos]
+                except KeyError:
+                    pass
+                idx += 1
+        knwl_vis = tuple(knwl_img + knwl_txt)
+    else:
+        knwl_embd = None
+        knwl_vis = knwl_none
+
+    # generate output
+    prompt = state.get_prompt()
+    prompt = prompt.split("USER:")[-1].replace("ASSISTANT:", "")
+    image_pt = image_trans(image).to(device).unsqueeze(0)
+    samples = {"image": image_pt, "knowledge": knwl_embd, "prompt": prompt}
+    
+    if bool(do_beam_search):
+        new_text = gptk_model.generate(
+            samples=samples,
+            use_nucleus_sampling=bool(do_sampling),
+            max_length=min(int(max_new_tokens), 1024),
+            top_p=float(top_p),
+            temperature=float(temperature),
+            auto_cast=True
+        )[0]
+        streamer = [new_text, ]
+    else:
+        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=bool(do_sampling),
+                max_length=min(int(max_new_tokens), 1024),
+                top_p=float(top_p),
+                temperature=float(temperature),
+                streamer=streamer,
+                num_beams=1,
+                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,) * 2 + knwl_vis
+        time.sleep(0.03)
+    state.messages[-1][-1] = state.messages[-1][-1][:-1]
+    yield (state, state.to_gradio_chatbot()) + (enable_btn,) * 2 + knwl_vis
+        
+
+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")
+    state = gr.State()
+
+    with gr.Blocks(title="GPT-K", theme=gr.themes.Base()) as demo:
+        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/diamond_head.jpg", "What is the name of this famous sight in the photo?"],
+                    ["examples/horseshoe_bend.jpg", "What are the possible reasons of the formation of this sight?"],
+                    ["examples/titanic.jpg", "What happen in the scene in this movie?"],
+                ], inputs=[imagebox, textbox])
+
+                imagebox.render()
+                textbox.render()
+                with gr.Column():
+                    submit_btn = gr.Button(value="📝 Submit")
+                    regenerate_btn = gr.Button(value="🔄  Regenerate", interactive=False)
+                    clear_btn = gr.Button(value="🗑️  Clear", interactive=False)
+
+                with gr.Accordion("Parameters", open=True):
+                    with gr.Row():
+                        add_knwl = gr.Checkbox(value=True, interactive=True, label="Knowledge")
+                        do_sampling = gr.Checkbox(value=False, interactive=True, label="Sampling")
+                        do_beam_search = gr.Checkbox(value=False, interactive=True, label="Beam search")
+                    temperature = gr.Slider(minimum=0.1, maximum=2.0, value=1.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="LLaVA 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 = [regenerate_btn, clear_btn]
+        regenerate_btn.click(
+            regenerate, [state], [state, chatbot, textbox, imagebox] + btn_list
+        ).then(
+            generate,
+            [state, temperature, top_p, max_output_tokens, add_knwl, do_sampling, do_beam_search],
+            [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, add_knwl, do_sampling, do_beam_search],
+            [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, add_knwl, do_sampling, do_beam_search],
+            [state, chatbot] + btn_list + knwl_vis
+        )
+
+        demo.load(load_demo, None, [state, ])
+    
+    return demo
+
+
+def build_model():
+    if torch.cuda.is_available():
+        device = torch.device("cuda")
+    else:
+        device = torch.device("cpu")
+
+    query_enc, _, query_trans = open_clip.create_model_and_transforms(
+        "ViT-g-14", pretrained="laion2b_s34b_b88k", precision='fp16'
+    )
+    query_enc = query_enc.to(device).eval()
+
+    def get_knwl(knowledge_db):
+        knwl_db = TextDB(Path(knowledge_db)/"knowledge_db.hdf5")
+        knwl_idx = faiss.read_index(str(Path(knowledge_db)/"faiss.index"))
+        knwl_idx.add(knwl_db.feature)
+
+        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]
+
+    _, image_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/mp_rank_00_model_states.pt"
+    gptk_ckpt = torch.load(gptk_ckpt, map_location="cpu")
+    gptk_ckpt = {
+        ".".join(k.split(".")[2:]): v
+        for k, v in gptk_ckpt["module"].items()
+    }
+    gptk_model.load_state_dict(gptk_ckpt)
+    gptk_model = gptk_model.to(device).eval()
+
+    return knwl_db, query_enc, query_trans, gptk_model, image_trans, topk, device
+
+
+knwl_db, query_enc, query_trans, gptk_model, image_trans, topk, device = build_model()
+demo = build_demo()
+demo.queue().launch()

conversation.py

@@ -0,0 +1,364 @@
+import dataclasses
+from enum import auto, Enum
+from typing import List, Tuple
+
+
+class SeparatorStyle(Enum):
+    """Different separator style."""
+    SINGLE = auto()
+    TWO = auto()
+    MPT = auto()
+    PLAIN = auto()
+    LLAMA_2 = auto()
+
+
+@dataclasses.dataclass
+class Conversation:
+    """A class that keeps all conversation history."""
+    system: str
+    roles: List[str]
+    messages: List[List[str]]
+    offset: int
+    sep_style: SeparatorStyle = SeparatorStyle.SINGLE
+    sep: str = "###"
+    sep2: str = None
+    version: str = "Unknown"
+
+    skip_next: bool = False
+
+    def get_prompt(self):
+        messages = self.messages
+        if len(messages) > 0 and type(messages[0][1]) is tuple:
+            messages = self.messages.copy()
+            init_role, init_msg = messages[0].copy()
+            init_msg = init_msg[0].replace("<image>", "").strip()
+            if 'mmtag' in self.version:
+                messages[0] = (init_role, init_msg)
+                messages.insert(0, (self.roles[0], "<Image><image></Image>"))
+                messages.insert(1, (self.roles[1], "Received."))
+            else:
+                messages[0] = (init_role, "<image>\n" + init_msg)
+
+        if self.sep_style == SeparatorStyle.SINGLE:
+            ret = self.system + self.sep
+            for role, message in messages:
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + ": " + message + self.sep
+                else:
+                    ret += role + ":"
+        elif self.sep_style == SeparatorStyle.TWO:
+            seps = [self.sep, self.sep2]
+            ret = self.system + seps[0]
+            for i, (role, message) in enumerate(messages):
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + ": " + message + seps[i % 2]
+                else:
+                    ret += role + ":"
+        elif self.sep_style == SeparatorStyle.MPT:
+            ret = self.system + self.sep
+            for role, message in messages:
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + message + self.sep
+                else:
+                    ret += role
+        elif self.sep_style == SeparatorStyle.LLAMA_2:
+            wrap_sys = lambda msg: f"<<SYS>>\n{msg}\n<</SYS>>\n\n"
+            wrap_inst = lambda msg: f"[INST] {msg} [/INST]"
+            ret = ""
+
+            for i, (role, message) in enumerate(messages):
+                if i == 0:
+                    assert message, "first message should not be none"
+                    assert role == self.roles[0], "first message should come from user"
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    if i == 0: message = wrap_sys(self.system) + message
+                    if i % 2 == 0:
+                        message = wrap_inst(message)
+                        ret += self.sep + message
+                    else:
+                        ret += " " + message + " " + self.sep2
+                else:
+                    ret += ""
+            ret = ret.lstrip(self.sep)
+        elif self.sep_style == SeparatorStyle.PLAIN:
+            seps = [self.sep, self.sep2]
+            ret = self.system
+            for i, (role, message) in enumerate(messages):
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += message + seps[i % 2]
+                else:
+                    ret += ""
+        else:
+            raise ValueError(f"Invalid style: {self.sep_style}")
+
+        return ret
+
+    def append_message(self, role, message):
+        self.messages.append([role, message])
+
+    def get_images(self, return_pil=False):
+        images = []
+        for i, (role, msg) in enumerate(self.messages[self.offset:]):
+            if i % 2 == 0:
+                if type(msg) is tuple:
+                    image = msg[1].convert('RGB')
+                    if return_pil:
+                        images.append(image)
+                    else:
+                        import base64
+                        from io import BytesIO
+                        buffered = BytesIO()
+                        image.save(buffered, format="PNG")
+                        img_b64_str = base64.b64encode(buffered.getvalue()).decode()
+                        images.append(img_b64_str)
+        return images
+
+    def to_gradio_chatbot(self):
+        ret = []
+        for i, (role, msg) in enumerate(self.messages[self.offset:]):
+            if i % 2 == 0:
+                if type(msg) is tuple:
+                    import base64
+                    from io import BytesIO
+                    msg, image = msg
+                    max_hw, min_hw = max(image.size), min(image.size)
+                    aspect_ratio = max_hw / min_hw
+                    max_len, min_len = 800, 400
+                    shortest_edge = int(min(max_len / aspect_ratio, min_len, min_hw))
+                    longest_edge = int(shortest_edge * aspect_ratio)
+                    W, H = image.size
+                    if H > W:
+                        H, W = longest_edge, shortest_edge
+                    else:
+                        H, W = shortest_edge, longest_edge
+                    image = image.resize((W, H))
+                    buffered = BytesIO()
+                    image.save(buffered, format="JPEG")
+                    img_b64_str = base64.b64encode(buffered.getvalue()).decode()
+                    img_str = f'<img src="data:image/png;base64,{img_b64_str}" alt="user upload image" />'
+                    ret.append([img_str, None])
+                    msg = msg.replace('<image>', '').strip()
+                    if len(msg) > 0:
+                        ret.append([msg, None])
+                else:
+                    ret.append([msg, None])
+            else:
+                ret[-1][-1] = msg
+        return ret
+
+    def copy(self):
+        return Conversation(
+            system=self.system,
+            roles=self.roles,
+            messages=[[x, y] for x, y in self.messages],
+            offset=self.offset,
+            sep_style=self.sep_style,
+            sep=self.sep,
+            sep2=self.sep2,
+            version=self.version)
+
+    def dict(self):
+        if len(self.get_images()) > 0:
+            return {
+                "system": self.system,
+                "roles": self.roles,
+                "messages": [[x, y[0] if type(y) is tuple else y] for x, y in self.messages],
+                "offset": self.offset,
+                "sep": self.sep,
+                "sep2": self.sep2,
+            }
+        return {
+            "system": self.system,
+            "roles": self.roles,
+            "messages": self.messages,
+            "offset": self.offset,
+            "sep": self.sep,
+            "sep2": self.sep2,
+        }
+
+
+conv_gptk = Conversation(
+    system="",
+    roles=("USER", "ASSISTANT"),
+    version="v1",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.SINGLE,
+    sep=""
+)
+
+
+conv_vicuna_v0 = Conversation(
+    system="A chat between a curious human and an artificial intelligence assistant. "
+           "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+    roles=("Human", "Assistant"),
+    messages=(
+        ("Human", "What are the key differences between renewable and non-renewable energy sources?"),
+        ("Assistant",
+            "Renewable energy sources are those that can be replenished naturally in a relatively "
+            "short amount of time, such as solar, wind, hydro, geothermal, and biomass. "
+            "Non-renewable energy sources, on the other hand, are finite and will eventually be "
+            "depleted, such as coal, oil, and natural gas. Here are some key differences between "
+            "renewable and non-renewable energy sources:\n"
+            "1. Availability: Renewable energy sources are virtually inexhaustible, while non-renewable "
+            "energy sources are finite and will eventually run out.\n"
+            "2. Environmental impact: Renewable energy sources have a much lower environmental impact "
+            "than non-renewable sources, which can lead to air and water pollution, greenhouse gas emissions, "
+            "and other negative effects.\n"
+            "3. Cost: Renewable energy sources can be more expensive to initially set up, but they typically "
+            "have lower operational costs than non-renewable sources.\n"
+            "4. Reliability: Renewable energy sources are often more reliable and can be used in more remote "
+            "locations than non-renewable sources.\n"
+            "5. Flexibility: Renewable energy sources are often more flexible and can be adapted to different "
+            "situations and needs, while non-renewable sources are more rigid and inflexible.\n"
+            "6. Sustainability: Renewable energy sources are more sustainable over the long term, while "
+            "non-renewable sources are not, and their depletion can lead to economic and social instability.\n")
+    ),
+    offset=2,
+    sep_style=SeparatorStyle.SINGLE,
+    sep="###",
+)
+
+conv_vicuna_v1 = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+    "The assistant gives helpful, detailed, and polite answers to the user's questions.",
+    roles=("USER", "ASSISTANT"),
+    version="v1",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="</s>",
+)
+
+conv_llama_2 = Conversation(
+    system="""You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe.  Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.
+
+If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.""",
+    roles=("USER", "ASSISTANT"),
+    version="llama_v2",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.LLAMA_2,
+    sep="<s>",
+    sep2="</s>",
+)
+
+conv_llava_llama_2 = Conversation(
+    system="You are a helpful language and vision assistant. "
+           "You are able to understand the visual content that the user provides, "
+           "and assist the user with a variety of tasks using natural language.",
+    roles=("USER", "ASSISTANT"),
+    version="llama_v2",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.LLAMA_2,
+    sep="<s>",
+    sep2="</s>",
+)
+
+conv_mpt = Conversation(
+    system="""<|im_start|>system
+A conversation between a user and an LLM-based AI assistant. The assistant gives helpful and honest answers.""",
+    roles=("<|im_start|>user\n", "<|im_start|>assistant\n"),
+    version="mpt",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.MPT,
+    sep="<|im_end|>",
+)
+
+conv_llava_plain = Conversation(
+    system="",
+    roles=("", ""),
+    messages=(
+    ),
+    offset=0,
+    sep_style=SeparatorStyle.PLAIN,
+    sep="\n",
+)
+
+conv_llava_v0 = Conversation(
+    system="A chat between a curious human and an artificial intelligence assistant. "
+           "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+    roles=("Human", "Assistant"),
+    messages=(
+        ("Human", "Hi!"),
+        ("Assistant", "Hi there! How can I help you today?")
+    ),
+    offset=2,
+    sep_style=SeparatorStyle.SINGLE,
+    sep="###",
+)
+
+conv_llava_v0_mmtag = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+           "The assistant is able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language."
+           "The visual content will be provided with the following format: <Image>visual content</Image>.",
+    roles=("Human", "Assistant"),
+    messages=(
+    ),
+    offset=0,
+    sep_style=SeparatorStyle.SINGLE,
+    sep="###",
+    version="v0_mmtag",
+)
+
+conv_llava_v1 = Conversation(
+    system="A chat between a curious human and an artificial intelligence assistant. "
+           "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+    roles=("USER", "ASSISTANT"),
+    version="v1",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="</s>",
+)
+
+conv_llava_v1_mmtag = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+           "The assistant is able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language."
+           "The visual content will be provided with the following format: <Image>visual content</Image>.",
+    roles=("USER", "ASSISTANT"),
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="</s>",
+    version="v1_mmtag",
+)
+
+default_conversation = conv_vicuna_v0
+conv_templates = {
+    "default": conv_vicuna_v0,
+    "v0": conv_vicuna_v0,
+    "v1": conv_vicuna_v1,
+    "vicuna_v1": conv_vicuna_v1,
+    "llama_2": conv_llama_2,
+    "gptk": conv_gptk,
+
+    "plain": conv_llava_plain,
+    "v0_plain": conv_llava_plain,
+    "llava_v0": conv_llava_v0,
+    "v0_mmtag": conv_llava_v0_mmtag,
+    "llava_v1": conv_llava_v1,
+    "v1_mmtag": conv_llava_v1_mmtag,
+    "llava_llama_2": conv_llava_llama_2,
+
+    "mpt": conv_mpt,
+}
+
+
+if __name__ == "__main__":
+    print(default_conversation.get_prompt())

knowledge/(dataset-action)(clip-model-ViT-g-14)(dbscan)(eps-0.15)(ms-1)/faiss.index

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fb05eb3ab8b8e775c1e10ab21a4f8d409b77a47ffacbc606050c2055bd78549a
+size 45

knowledge/(dataset-action)(clip-model-ViT-g-14)(dbscan)(eps-0.15)(ms-1)/info.txt

@@ -0,0 +1 @@
+n_samples = 148,620; n_clusters = 43,296; noise_ratio = 0.000%

knowledge/(dataset-action)(clip-model-ViT-g-14)(dbscan)(eps-0.15)(ms-1)/knowledge_db.hdf5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6281557260322cacbfbe58d710e3dd537e823d6d6565da7c9fea27e30ced5e31
+size 166074480

knowledge/(dataset-action)(clip-model-ViT-g-14)(dbscan)(eps-0.15)(ms-1)/labels.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:881bf21972ffb9a9155d185282530a75a4ca4ffdb75c8a05d38dda901c0f366c
+size 1189088

knowledge/(dataset-attribute)(clip-model-ViT-g-14)(dbscan)(eps-0.15)(ms-1)/faiss.index

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1efe5c6accd575c85403aaeccaf24c6fb1cfff05bd6a0f1ecdbdbc0ce0a5befa
+size 9093259

knowledge/(dataset-attribute)(clip-model-ViT-g-14)(dbscan)(eps-0.15)(ms-1)/info.txt

@@ -0,0 +1 @@
+n_samples = 191,836; n_clusters = 77,073; noise_ratio = 0.000%

knowledge/(dataset-attribute)(clip-model-ViT-g-14)(dbscan)(eps-0.15)(ms-1)/knowledge_db.hdf5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:51449f86c49a1651debdaf7ec1b4c1020db911785bd5f51e0766a4bfefe1897f
+size 295832959

knowledge/(dataset-attribute)(clip-model-ViT-g-14)(dbscan)(eps-0.15)(ms-1)/labels.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6cf9a479e50595e52e593f961d4f3dcc822d9c0caf097fed3498a64c175f7e2c
+size 1534816

knowledge/(dataset-object)(clip-model-ViT-g-14)(dbscan)(eps-0.15)(ms-1)/faiss.index

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ece63b94bf3672252b77fbbf47a3070a378280ef3eafb682f99340fc74e1d096
+size 18702475

knowledge/(dataset-object)(clip-model-ViT-g-14)(dbscan)(eps-0.15)(ms-1)/info.txt

@@ -0,0 +1 @@
+n_samples = 770,808; n_clusters = 325,813; noise_ratio = 0.000%

knowledge/(dataset-object)(clip-model-ViT-g-14)(dbscan)(eps-0.15)(ms-1)/knowledge_db.hdf5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9dcca3e4560c724f42128b8d476dd28ad0305ad66125213050c7fec7715d6a8b
+size 1251033850

knowledge/(dataset-object)(clip-model-ViT-g-14)(dbscan)(eps-0.15)(ms-1)/labels.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c79a747b0551e46056391ad988317604dd29a8905acb3167127550dcc6b90890
+size 6166592

knowledge/__init__.py

@@ -0,0 +1,2 @@
+from .text_db import TextDB
+from .retrieve import *

knowledge/cluster.py

@@ -0,0 +1,178 @@
+import argparse
+from pathlib import Path
+import numpy as np
+from tqdm import tqdm
+import h5py
+import time
+
+import faiss
+import torch
+from pytorch_lightning import seed_everything
+
+import sys
+sys.path.append('.')
+from knowledge.text_db import TextDB
+from knowledge.utils import nn_search, build_faiss_index, refine_cosine
+
+
+UNSEEN = -2
+NOISE = -1
+
+
+def dbscan(X, faiss_index, device, eps=0.1, min_points=1, k=2048, bs=512):
+    neighbors = []
+    N = (len(X) - 1) // bs + 1
+    for i in tqdm(range(N), dynamic_ncols=True, desc="Find nearest neighbors", mininterval=1.0):
+        Xi = X[i*bs: (i+1)*bs]
+        _, I = faiss_index.search(Xi, k*2)
+        S, I = refine_cosine(X, Xi, I, device, k)
+
+        for sim, idx in zip(S, I):
+            dist = 1. - sim
+            neighbors.append(idx[dist < eps])
+        
+    cluster_id = 0
+    n_points = len(X)
+    labels = np.array([
+        NOISE if len(neighbors[i]) < min_points else UNSEEN
+        for i in range(n_points)
+    ])
+
+    with tqdm(total=n_points, dynamic_ncols=True, desc="DBSCAN clustering", mininterval=1.0) as pbar:
+        for i in range(n_points):
+            if labels[i] == UNSEEN:
+                seeds = np.array([i, ])
+                labels[seeds] = cluster_id
+
+                while len(seeds) > 0:
+                    neighbor_seeds = set()
+                    for s in seeds:
+                        n = neighbors[s]
+                        if len(n) > 0:
+                            l = np.array(list(set(labels[n])))
+                            l = l[np.logical_and(l >= 0, l != cluster_id)]
+                            for li in l:
+                                labels[labels == li] = cluster_id
+
+                            n = n[labels[n] == UNSEEN]
+                            neighbor_seeds.update(n)
+
+                    seeds = np.array(list(neighbor_seeds))
+                    if len(seeds) > 0:
+                        assert np.all(labels[seeds] == UNSEEN)
+                        labels[seeds] = cluster_id
+
+                cluster_id += 1
+            
+            pbar.set_postfix(num_clusters=cluster_id)
+            pbar.update()
+    
+    label_set = np.sort(list(set(labels)))
+    label_set = label_set[label_set >= 0]
+    labels_mapping = {l1: l2 for l2, l1 in enumerate(label_set)}
+    labels_mapping[-1] = -1
+    labels = np.array([labels_mapping[l] for l in labels])
+            
+    return labels
+
+
+def extract_clusters(feat, text, labels, faiss_index, device, k=128, bs=8192):
+    clusters = {}
+    for i, l in enumerate(tqdm(labels, dynamic_ncols=True, desc="Label each samples", mininterval=1.0)):
+        if l >= 0:
+            try:
+                clusters[l]["feat"] += feat[i].astype(np.float64)
+                clusters[l]["N"] += 1
+            except KeyError:
+                clusters[l] = {"feat": feat[i].astype(np.float64), "N": 1}
+    
+    cc = []
+    for l in tqdm(list(clusters.keys()), dynamic_ncols=True, desc="Compute cluster centers", mininterval=1.0):
+        c = clusters[l]["feat"]/clusters[l]["N"]
+        cc.append(c.astype(np.float32))
+    cc = np.stack(cc)
+    cc /= np.linalg.norm(cc, keepdims=True, axis=-1)
+
+    idx = []
+    N = (len(cc) - 1) // bs + 1
+    for i in tqdm(range(N), dynamic_ncols=True, desc="Find nearest neighbors", mininterval=1.0):
+        cc_i = cc[i*bs: (i+1)*bs]
+        _, I = faiss_index.search(cc_i, k)
+        _, I = refine_cosine(feat, cc_i, I, device, 1)
+        idx.append(I[:, 0])
+    idx = np.unique(np.concatenate(idx))
+    text = [text[i] for i in idx]
+    feat = np.stack([feat[i] for i in idx])
+
+    return feat, text
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Cluster knowledge database using DBSCAN")
+    parser.add_argument("--knowledge_db", type=str, required=True)
+    parser.add_argument("--seed", type=int, default=12345)
+    parser.add_argument("--eps", type=float, default=0.1)
+    parser.add_argument("--ms", type=int, default=1)
+    parser.add_argument("--ratio", type=float, default=None)
+    parser.add_argument("--device", type=int, default=None)
+    args = parser.parse_args()
+
+    # parse exp name
+    args.knowledge_db = Path(args.knowledge_db)
+    exp_name = args.knowledge_db.parent.name
+    exp_name += f"(dbscan)(eps-{args.eps})(ms-{args.ms})"
+    save_root = args.knowledge_db.parent.parent/exp_name
+    setattr(args, "save_root", save_root)
+    args.save_root.mkdir(parents=True, exist_ok=True)
+
+    args.device = torch.device("cuda", args.device) \
+        if args.device is not None else torch.device("cpu")
+
+    seed_everything(args.seed, workers=True)
+    print(args)
+    
+    # load feature, text, and faiss index from knowledge db
+    knowledge_db = TextDB(args.knowledge_db)
+    feat = knowledge_db.feature.astype(np.float32)
+    text = knowledge_db.text
+    if args.ratio is not None:
+        N = int(len(feat) * args.ratio)
+        feat, text = feat[:N], text[:N]
+    faiss_index = faiss.read_index(str(args.knowledge_db.parent/"faiss.index"))
+    print("Add data to faiss index...", end="\r")
+    ts = time.time()
+    faiss_index.add(feat)
+    print(f"Add data to faiss index...done in {time.time() - ts:.2f} secs")
+    
+    # DBSCAN clustering
+    labels_file = args.save_root/"labels.npy"
+    if labels_file.exists():
+        labels = np.load(labels_file)
+    else:
+        labels = dbscan(feat, faiss_index, args.device, args.eps, args.ms)
+        with open(labels_file, 'wb') as f:
+            np.save(f, labels)
+
+    # extract clusters
+    feat, text = extract_clusters(feat, text, labels, faiss_index, args.device)
+    with h5py.File(args.save_root/f"knowledge_db.hdf5", "w") as f:
+        bs = 65536
+        N = (len(feat) - 1) // bs + 1
+        for i in tqdm(range(N), dynamic_ncols=True, desc="Saving clustered DB", mininterval=1.0):
+            g = f.create_group(str(i))
+            g.create_dataset("feature", data=feat[i*bs: (i+1)*bs], compression="gzip")
+            g.create_dataset("text", data=text[i*bs: (i+1)*bs], compression="gzip")
+    
+    # build faiss index for the clustered DB
+    index = build_faiss_index(feat, gpus=[args.device.index, ])
+    faiss.write_index(index, str(args.save_root/"faiss.index"))
+
+    # some stats
+    noise_ratio = np.sum(labels == -1) / len(labels)
+    n_clusters, n_samples = len(text), len(labels)
+    msg = f"n_samples = {n_samples:,}; n_clusters = {n_clusters:,}; noise_ratio = {noise_ratio*100:.3f}%\n"
+    with open(save_root/"info.txt", "w") as f:
+        f.write(msg)
+    print(msg)
+
+

knowledge/retrieve.py

@@ -0,0 +1,327 @@
+import argparse
+from pathlib import Path
+import h5py
+import time
+import shutil
+import numpy as np
+import subprocess
+import time
+from tqdm import tqdm
+
+import faiss
+import open_clip
+import torch
+import torch.distributed as dist
+from torch.utils.data import DataLoader
+from pytorch_lightning import callbacks
+from pytorch_lightning import Trainer, LightningModule, seed_everything
+
+import sys
+sys.path.append('.')
+from dataset import coco, cc, llava
+from knowledge.utils import refine_cosine
+from knowledge import text_db
+from knowledge import TextDB
+from train.utils import ExpName
+
+
+class ImageCropsIdx:
+    def __init__(self, knowledge_idx, topk_w, topk_f, topk_n):
+        topk = {"whole": topk_w, "five": topk_f, "nine": topk_n}
+        self.topk = {k: v for k, v in topk.items() if v > 0}
+
+        self.knowledge_idx, self.fdim, self.file_hash = self.load(knowledge_idx, self.topk)
+
+    def load(self, knowledge_idx, topk):
+        with h5py.File(knowledge_idx, "r") as f:
+            fdim = f.attrs["fdim"]
+            file_hash = f.attrs["file_hash"]
+
+            knowledge_idx_ = {}
+            for i in tqdm(range(len(f)), desc="Load sentence idx", dynamic_ncols=True, mininterval=1.0):
+                knowledge_idx_[str(i)] = {"image_ids": f[f"{i}/image_ids"][:]}
+                for k, v in topk.items():
+                    knowledge_idx_[str(i)][k] = {
+                        "index": f[f"{i}/{k}/index"][:, :, :v],
+                        "score": f[f"{i}/{k}/score"][:, :, :v],
+                        "query": f[f"{i}/{k}/query"][:]
+                    }
+        
+        knowledge_idx = {}
+        for i in knowledge_idx_.keys():
+            for j, id in enumerate(knowledge_idx_[i]["image_ids"]):
+                knowledge_idx[id] = {}
+                for k in topk.keys():
+                    knowledge_idx[id][k] = {
+                        "index": knowledge_idx_[i][k]["index"][j],
+                        "score": knowledge_idx_[i][k]["score"][j],
+                        "query": knowledge_idx_[i][k]["query"][j],
+                    }
+
+        return knowledge_idx, fdim, file_hash
+    
+    def __getitem__(self, image_id):
+        return self.knowledge_idx[image_id]
+        
+
+class KnowAugImageCrops:
+    def __init__(self, knowledge_db: TextDB, knowledge_idx: ImageCropsIdx, return_txt=False):
+        self.knowledge_db = knowledge_db
+        self.knowledge_idx = knowledge_idx
+        assert knowledge_db.file_hash == knowledge_idx.file_hash
+
+        self.ncrop = {"whole": 1, "five": 5, "nine": 9}
+        self.topk = knowledge_idx.topk
+        self.fdim = knowledge_idx.fdim
+
+        self.return_txt = return_txt
+    
+    def __call__(self, image_id):
+        ret = {}
+        for k in self.topk.keys():
+            ki = self.knowledge_idx[image_id][k]["index"].flatten()
+            ke, kt = self.knowledge_db[ki]
+            kq = self.knowledge_idx[image_id][k]["query"]
+            kp = np.tile(np.arange(self.ncrop[k])[:, None], (1, self.topk[k])).flatten()
+            ks = self.knowledge_idx[image_id][k]["score"].flatten()
+
+            ke = torch.FloatTensor(ke)
+            kq = torch.FloatTensor(kq)
+            kp = torch.LongTensor(kp)
+            ks = torch.FloatTensor(ks)
+
+            ret[k] = {"embed": ke, "query": kq, "pos": kp, "score": ks}
+            if self.return_txt:
+                ret[k]["text"] = kt
+        
+        return ret
+
+
+class KnowAugImageCropsCombined:
+    def __init__(
+        self, 
+        knwl_aug_obj: KnowAugImageCrops,
+        knwl_aug_attr: KnowAugImageCrops,
+        knwl_aug_act: KnowAugImageCrops
+    ):
+        self.knwl_aug_obj = knwl_aug_obj
+        self.knwl_aug_act = knwl_aug_act
+        self.knwl_aug_attr = knwl_aug_attr
+        self.fdim = knwl_aug_obj.fdim
+
+    def __call__(self, image_id):
+        knwl_obj = self.knwl_aug_obj(image_id)
+        knwl_attr = self.knwl_aug_attr(image_id)
+        knwl_act = self.knwl_aug_act(image_id)
+
+        ret = {}
+        for k in knwl_obj.keys():
+            ret[k] = {
+                "obj": knwl_obj[k],
+                "attr": knwl_attr[k],
+                "act": knwl_act[k]
+            }
+
+        return ret
+        
+
+class ImageCropsIdxBuilder(LightningModule):
+    def __init__(self, args, model: open_clip.model.CLIP):
+        super().__init__()
+
+        self.args = args
+        self.save_root = args.save_root
+        self.k = args.k
+        self.model = model
+        
+    def on_validation_epoch_start(self):
+        if self.global_rank == 0:
+            knowledge_db = TextDB(self.args.knowledge_db)
+            self.feature = knowledge_db.feature
+            self.text = knowledge_db.text
+
+            self.faiss_index = faiss.read_index(
+                str(Path(self.args.knowledge_db).parent/"faiss.index")
+            )
+            print("\nAdd data to faiss index...", end="\r")
+            ts = time.time()
+            self.faiss_index.add(self.feature)
+            print(f"Add data to faiss index...done in {time.time() - ts:.2f} secs")
+
+            with h5py.File(self.save_root/"knowledge_idx.hdf5", "a") as f:
+                f.attrs["fdim"] = self.feature.shape[-1]
+                f.attrs["file_hash"] = knowledge_db.file_hash
+        
+        self.trainer.strategy.barrier()
+    
+    def all_gather_object(self, data):
+        if self.trainer.world_size > 1:
+            gathered = [None for _ in range(self.trainer.world_size)]
+            dist.all_gather_object(gathered, data)
+            data = gathered
+        else:
+            data = [data, ]
+        
+        return data
+    
+    def broadcast_object(self, data, src_rank=0):
+        if self.trainer.world_size > 1:
+            if self.global_rank == src_rank:
+                data_list = [data, ] * self.trainer.world_size
+            else:
+                data_list = [None, ] * self.trainer.world_size
+            
+            dist.broadcast_object_list(data_list, src=src_rank)
+            return data_list[0]
+        else:
+            return data
+    
+    def search(self, images, topk):
+        query = self.model.encode_image(images, normalize=True)
+        query = query.cpu().numpy()
+        query = self.all_gather_object(query)
+        query = np.concatenate(query)
+
+        if self.global_rank == 0:
+            _, I = self.faiss_index.search(query, 4*topk)
+            S, I = refine_cosine(self.feature, query, I, self.device, topk)
+        else:
+            S = I = None
+
+        return S, I, query
+    
+    def validation_step(self, batch, batch_idx):
+        orig_imgs, five_imgs, nine_imgs, ids = batch
+
+        ids = ids.cpu().numpy()
+        ids = np.concatenate(self.all_gather_object(ids))
+
+        S_w, I_w, Q_w = self.search(orig_imgs, topk=self.k)
+        
+        S_f, I_f, Q_f = [], [], []
+        for i in range(five_imgs.shape[1]):
+            Si, Ii, Qi = self.search(five_imgs[:, i], topk=self.k)
+            S_f.append(Si)
+            I_f.append(Ii)
+            Q_f.append(Qi)
+        
+        S_n, I_n, Q_n = [], [], []
+        for i in range(nine_imgs.shape[1]):
+            Si, Ii, Qi = self.search(nine_imgs[:, i], topk=self.k)
+            S_n.append(Si)
+            I_n.append(Ii)
+            Q_n.append(Qi)
+
+        if self.global_rank == 0:
+            S_w, I_w, Q_w = np.expand_dims(S_w, axis=1), np.expand_dims(I_w, axis=1), np.expand_dims(Q_w, axis=1)
+            S_f, I_f, Q_f = np.stack(S_f, axis=1), np.stack(I_f, axis=1), np.stack(Q_f, axis=1)
+            S_n, I_n, Q_n = np.stack(S_n, axis=1), np.stack(I_n, axis=1), np.stack(Q_n, axis=1)
+
+            with h5py.File(self.save_root/"knowledge_idx.hdf5", "a") as f:
+                g = f.create_group(str(batch_idx))
+
+                g.create_dataset("image_ids", data=ids.astype(np.int32), compression="gzip")
+                
+                gw = g.create_group("whole")
+                gw.create_dataset("index", data=I_w.astype(np.int32), compression="gzip")
+                gw.create_dataset("score", data=S_w.astype(np.float32), compression="gzip")
+                gw.create_dataset("query", data=Q_w.astype(np.float32), compression="gzip")
+
+                gf = g.create_group("five")
+                gf.create_dataset("index", data=I_f.astype(np.int32), compression="gzip")
+                gf.create_dataset("score", data=S_f.astype(np.float32), compression="gzip")
+                gf.create_dataset("query", data=Q_f.astype(np.float32), compression="gzip")
+
+                gn = g.create_group("nine")
+                gn.create_dataset("index", data=I_n.astype(np.int32), compression="gzip")
+                gn.create_dataset("score", data=S_n.astype(np.float32), compression="gzip")
+                gn.create_dataset("query", data=Q_n.astype(np.float32), compression="gzip")
+
+    def on_validation_epoch_end(self):
+        if self.args.azcopy and self.global_rank == 0:
+            with open("azcopy/sas_output", "r") as f:
+                sas = f.readline()
+            sas_base, sas_key = sas.split("?")
+            sas = f"{sas_base}/knowledge_idx?{sas_key}"
+
+            cmd = ["azcopy/azcopy", "copy", str(self.args.save_root), sas, "--recursive=true"]
+            print(f"start copying data with command {cmd}")
+            ts = time.time()
+            subprocess.run(cmd)
+            print(f"done copying data in {time.time() - ts:.2f} secs")
+
+
+def main(args):
+    model, _, trans_img = open_clip.create_model_and_transforms(
+        args.clip_model, pretrained=text_db.CLIP_MODELS[args.clip_model]
+    )
+
+    print("load query dataset...")
+    if "coco" in args.query:
+        dset = coco.COCOImageCrops(Path(f"data/{args.query}"), trans=trans_img)
+        collate_crops = coco.collate_coco_crops
+    elif args.query == "cc3m":
+        dset = cc.CC3MImageCrops(Path("data/cc3m_instruct"), trans=trans_img)
+        collate_crops = cc.collate_cc_crops
+    elif args.query == "llava":
+        dset = llava.LLaVAImageCrops(Path("data/llava_bench"), trans=trans_img)
+        collate_crops = llava.collate_llava_crops
+    else:
+        raise ValueError
+    loader = DataLoader(
+        dset, batch_size=args.bs, shuffle=False, num_workers=args.num_workers,
+        drop_last=False, collate_fn=collate_crops
+    )
+
+    print("build model and trainer...")
+    pl_model = ImageCropsIdxBuilder(args, model)
+    model_summary = callbacks.RichModelSummary()
+    progress_bar = callbacks.TQDMProgressBar(args.refresh_rate)
+    trainer_callbacks = [model_summary, progress_bar]
+    trainer = Trainer(
+        sync_batchnorm=True,
+        precision=16,
+        accelerator='gpu',
+        devices=args.devices,
+        strategy="ddp",
+        default_root_dir=args.save_root,
+        callbacks=trainer_callbacks,
+        limit_val_batches=args.limit_val_batches
+    )
+
+    print("retrieve knowledge...")
+    trainer.validate(pl_model, dataloaders=loader)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Knowledge retrieval using image crops')
+    parser = Trainer.add_argparse_args(parser)
+    parser.add_argument('--query', type=str, choices=["coco14", "coco17", "cc3m", "llava"], required=True)
+    parser.add_argument('--knowledge_db', type=str, required=True)
+    parser.add_argument('--k', type=int, default=128)
+    parser.add_argument("--bs", type=int, default=128)
+    parser.add_argument("--num_workers", type=int, default=7)
+    parser.add_argument("--seed", type=int, default=12345)
+    parser.add_argument("--refresh_rate", type=int, default=1)
+    parser.add_argument("--azcopy", action="store_true")
+    args = parser.parse_args()
+
+    # parse exp_name
+    exp_name = ExpName(f"(query-{args.query})")
+    exp_name += Path(args.knowledge_db).parent.name
+    if args.azcopy:
+        setattr(args, "save_root", Path("azcopy")/str(exp_name))
+    else:
+        setattr(args, "save_root", Path("output")/"knowledge_idx"/str(exp_name))
+    shutil.rmtree(args.save_root, ignore_errors=True)
+    args.save_root.mkdir(parents=True, exist_ok=True)
+
+    # parse model
+    model = exp_name.get("clip-model")[1:-1]
+    model = model[len("clip-model-"):]
+    assert model in text_db.CLIP_MODELS.keys()
+    setattr(args, "clip_model", model)
+
+    print(args)
+    seed_everything(args.seed, workers=True)
+    main(args)

knowledge/text_db.py

@@ -0,0 +1,197 @@
+import argparse
+import itertools
+from pathlib import Path
+import shutil
+import h5py
+import time
+import subprocess
+from tqdm import tqdm
+import numpy as np
+import codecs
+
+import open_clip
+import faiss
+import torch
+import torch.distributed as dist
+from torch.utils.data import DataLoader
+from pytorch_lightning import callbacks
+from pytorch_lightning import Trainer, LightningModule, seed_everything
+
+import sys
+sys.path.append("./")
+from dataset import cc, words
+from knowledge.utils import file_hash, build_faiss_index
+
+
+class TextDB:
+    def __init__(self, text_db):
+        self.feature, self.text = self.load(text_db)
+        self.file_hash = file_hash(text_db)
+
+    def load(self, text_db):
+        with h5py.File(text_db, 'r') as f:
+            db_size = 0
+            for i in range(len(f)):
+                db_size += len(f[f"{i}/feature"])
+            _, d = f[f"0/feature"].shape
+
+        with h5py.File(text_db, 'r') as f:
+            feature = np.zeros((db_size, d), dtype=np.float32)
+            text = []
+            N = 0
+            for i in tqdm(range(len(f)), desc="Load text DB", dynamic_ncols=True, mininterval=1.0):
+                fi = f[f"{i}/feature"][:]
+                feature[N:N+len(fi)] = fi
+                N += len(fi)
+
+                text.extend(f[f"{i}/text"][:])
+        text = [codecs.decode(t) for t in text]
+
+        return feature, text
+    
+    def __getitem__(self, idx):
+        f = self.feature[idx]
+
+        try:
+            t = [self.text[i] for i in idx]
+        except TypeError:
+            t = self.text[idx]
+        
+        return f, t
+
+
+class TextDBBuilder(LightningModule):
+    def __init__(self, args, model: open_clip.model.CLIP):
+        super().__init__()
+        self.args = args
+        self.model = model
+    
+    def validation_step(self, batch, batch_idx):
+        token, text = batch
+        feat = self.model.encode_text(token, normalize=True)
+
+        if self.trainer.world_size > 1:
+            text_gathered = [None for _ in range(self.trainer.world_size)]
+            dist.all_gather_object(text_gathered, text)
+            text = list(itertools.chain.from_iterable(text_gathered))
+
+            feat_gathered = [None for _ in range(self.trainer.world_size)]
+            dist.all_gather_object(feat_gathered, feat)
+            feat = torch.cat([x.to(self.device) for x in feat_gathered])
+            feat = feat.cpu().numpy()
+
+        if self.global_rank == 0:
+            with h5py.File(self.args.save_root/"knowledge_db.hdf5", "a") as f:
+                g = f.create_group(str(batch_idx))
+                g.create_dataset("feature", data=feat, compression="gzip")
+                g.create_dataset("text", data=text, compression="gzip")
+        
+    def validation_epoch_end(self, outputs):
+        if self.global_rank == 0:
+            knowledge_db = TextDB(self.args.save_root/"knowledge_db.hdf5")
+            feat = knowledge_db.feature
+
+            if self.args.devices == "-1":
+                num_devices = torch.cuda.device_count()
+                devices = list(range(num_devices))
+            else:
+                devices = [int(x) for x in args.devices.split(",") if x]
+            print(f"CUDA devices: {devices}")
+
+            index = build_faiss_index(feat, gpus=devices)
+            faiss.write_index(index, str(self.args.save_root/"faiss.index"))
+        self.trainer.strategy.barrier()
+
+        if self.args.azcopy and self.global_rank == 0:
+            with open("azcopy/sas_output", "r") as f:
+                sas = f.readline()
+            sas_base, sas_key = sas.split("?")
+            sas = f"{sas_base}/knowledge_db?{sas_key}"
+
+            cmd = ["azcopy/azcopy", "copy", str(self.args.save_root), sas, "--recursive=true"]
+            print(f"start copying data with command {cmd}")
+            ts = time.time()
+            subprocess.run(cmd)
+            print(f"done copying data in {time.time() - ts:.2f} secs")
+        self.trainer.strategy.barrier()
+
+
+DATASETS = {
+    "object": words.ObjsDataset,
+    "attribute": words.AttrsDataset,
+    "action": words.ActsDataset,
+    "cc3m": cc.CC3MTextDataset,
+    "cc12m": cc.CC12MTextDataset
+}
+
+
+def main(args):
+    model, _, _ = open_clip.create_model_and_transforms(
+        args.clip_model, pretrained=CLIP_MODELS[args.clip_model]
+    )
+    trans_txt = open_clip.get_tokenizer(args.clip_model)
+
+    print("load dataset...")
+    dset = DATASETS[args.dataset](Path(args.data_root), trans_txt)
+    loader = DataLoader(
+        dset, batch_size=args.bs, shuffle=False, num_workers=args.num_workers,
+        drop_last=False, collate_fn=cc.collate_cc_txt
+    )
+
+    print("build model and trainer...")
+    pl_model = TextDBBuilder(args, model)
+    model_summary = callbacks.RichModelSummary()
+    progress_bar = callbacks.TQDMProgressBar(args.refresh_rate)
+    trainer_callbacks = [model_summary, progress_bar]
+    trainer = Trainer(
+        sync_batchnorm=True,
+        precision=16,
+        accelerator='gpu',
+        devices=args.devices,
+        strategy="ddp",
+        default_root_dir=args.save_root,
+        callbacks=trainer_callbacks,
+        limit_val_batches=args.limit_val_batches
+    )
+
+    print("compute textual features...")
+    trainer.validate(pl_model, dataloaders=loader)
+
+
+CLIP_MODELS = {
+    'ViT-B-32': 'openai',
+    'ViT-B-16': 'openai',
+    'ViT-L-14': 'openai',
+    'ViT-g-14': 'laion2b_s34b_b88k',
+    'ViT-bigG-14': 'laion2b_s39b_b160k',
+    'convnext_xxlarge': 'laion2b_s34b_b82k_augreg_soup',
+}
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Build knowledge database of words")
+    parser = Trainer.add_argparse_args(parser)
+    parser.add_argument(
+        "--dataset", type=str, required=True, choices=["object", "attribute", "action", "cc3m", "cc12m"]
+    )
+    parser.add_argument("--data_root", type=str, default="data/conceptnet/conceptnet-assertions-5.7.0.csv")
+    parser.add_argument("--clip_model", type=str, default="ViT-g-14", choices=CLIP_MODELS.keys())
+    parser.add_argument("--bs", type=int, default=2**10)
+    parser.add_argument("--num_workers", type=int, default=7)
+    parser.add_argument("--seed", type=int, default=12345)
+    parser.add_argument("--refresh_rate", type=int, default=1)
+    parser.add_argument("--azcopy", action="store_true")
+    args = parser.parse_args()
+
+    # feature dir
+    exp_name = f"(dataset-{args.dataset})(clip-model-{args.clip_model})"
+    if args.azcopy:
+        setattr(args, "save_root", Path("azcopy")/"knowledge_db"/exp_name)
+    else:
+        setattr(args, "save_root", Path("output")/"knowledge_db"/exp_name)
+    shutil.rmtree(args.save_root, ignore_errors=True)
+    args.save_root.mkdir(parents=True, exist_ok=True)
+    
+    print(args)
+    seed_everything(args.seed, workers=True)
+    main(args)

knowledge/transforms.py

@@ -0,0 +1,52 @@
+import itertools
+from torchvision.transforms import functional as F
+import re
+
+
+def five_crop(image, ratio=0.6):
+    w, h = image.size
+    hw = (h*ratio, w*ratio)
+
+    return F.five_crop(image, hw)
+
+def nine_crop(image, ratio=0.4):
+    w, h = image.size
+
+    t = (0, int((0.5-ratio/2)*h), int((1.0 - ratio)*h))
+    b = (int(ratio*h), int((0.5+ratio/2)*h), h)
+    l = (0, int((0.5-ratio/2)*w), int((1.0 - ratio)*w))
+    r = (int(ratio*w), int((0.5+ratio/2)*w), w)
+    h, w = list(zip(t, b)), list(zip(l, r))
+
+    images = []
+    for s in itertools.product(h, w):
+        h, w = s
+        top, left = h[0], w[0]
+        height, width = h[1]-h[0], w[1]-w[0]
+        images.append(F.crop(image, top, left, height, width))
+
+    return images
+
+
+def pre_caption(caption, max_words=None):
+    # Ref: https://github.com/salesforce/LAVIS/blob/main/lavis/processors/blip_processors.py#L49-L68
+    caption = re.sub(
+        r"([.!\"()*#:;~])",
+        " ",
+        caption.lower(),
+    )
+    caption = re.sub(
+        r"\s{2,}",
+        " ",
+        caption,
+    )
+    caption = caption.rstrip("\n")
+    caption = caption.strip(" ")
+
+    # truncate caption
+    caption_words = caption.split(" ")
+    if max_words is not None and len(caption_words) > max_words:
+        caption = " ".join(caption_words[: max_words])
+
+    return caption
+    

knowledge/utils.py

@@ -0,0 +1,127 @@
+from tqdm import tqdm
+import numpy as np
+import time
+import math
+import bisect
+import hashlib
+import faiss
+from faiss import StandardGpuResources, index_cpu_to_gpu_multiple_py
+import torch
+
+
+def file_hash(file):
+    # Ref: https://stackoverflow.com/a/59056837
+    with open(file, "rb") as f:
+        hash_fn = hashlib.blake2b()
+        chunk = f.read(8192)
+        while chunk:
+            hash_fn.update(chunk)
+            chunk = f.read(8192)
+
+    return hash_fn.hexdigest()
+
+
+def build_faiss_index(x, gpus=None):
+    # Ref: https://github.com/facebookresearch/faiss/wiki/Guidelines-to-choose-an-index
+    # Ref: https://gist.github.com/mdouze/46d6bbbaabca0b9778fca37ed2bcccf6
+
+    N, dim = x.shape
+    secs = [2**i for i in range(1, 15)]
+    d = secs[bisect.bisect_right(secs, dim) - 1] // 2
+    m = d // 4
+
+    if N <= 60000:
+        index_factory = "Flat"
+    elif N <= 2555904:
+        index_factory = f"IVF{int(8*math.sqrt(N))},Flat"
+    elif N <= 10223616:
+        index_factory = f"OPQ{m}_{d},IVF65536_HNSW32,PQ{m}x4fsr"
+    elif N <= 1e8:
+        index_factory = f"OPQ{m}_{d},IVF262144_HNSW32,PQ{m}x4fsr"
+    else:
+        index_factory = f"OPQ{m}_{d},IVF1048576_HNSW32,PQ{m}x4fsr"
+    print(f"train {index_factory} index on {N:,} x {dim} data")
+
+    index = faiss.index_factory(dim, index_factory)
+    if gpus is not None and N > 60000:
+        index_ivf = faiss.extract_index_ivf(index)
+        res = []
+        for _ in gpus:
+            r = StandardGpuResources()
+            r.noTempMemory()
+            res.append(r)
+        clustering_index = index_cpu_to_gpu_multiple_py(
+            res, faiss.IndexFlatL2(index_ivf.d), None, gpus
+        )
+        index_ivf.clustering_index = clustering_index
+
+    print("train index...", end="\r")
+    ts = time.time()
+    # commented out for index_factory = "Flat"
+    # assert not index.is_trained
+    index.train(x)
+    assert index.is_trained
+    print(f"train index...done in {time.time() - ts:.2f} secs")
+
+    index.nprobe = 64
+    index.quantizer_efSearch = 32
+
+    return index
+
+
+def nn_search(query, index, topk, bs=256, desc=None, disable_tqdm=True):
+    idx, dist = [], []
+    N = (len(query) - 1) // bs + 1
+    for i in tqdm(range(N), dynamic_ncols=True, desc=desc, disable=disable_tqdm):
+        D, I = index.search(query[i*bs: (i+1)*bs], topk)
+        idx.append(I)
+        dist.append(D)
+    idx = np.concatenate(idx)
+    dist = np.concatenate(dist)
+
+    return idx, dist
+
+
+def radius_search(query, index, r, bs=256, desc=None, disable_tqdm=True):
+    idx, dist = [], []
+    N = (len(query) - 1) // bs + 1
+    for i in tqdm(range(N), dynamic_ncols=True, desc=desc, disable=disable_tqdm):
+        L, D, I = index.range_search(query[i*bs: (i+1)*bs], r)
+        idx.extend([I[L[j]:L[j+1]] for j in range(len(L)-1)])
+        dist.extend([D[L[j]:L[j+1]] for j in range(len(L)-1)])
+
+    return idx, dist
+
+
+@torch.no_grad()
+def refine_cosine(Xa, Xq, I, device, k=None):
+    if k is not None:
+        assert k <= I.shape[1]
+    else:
+        k = I.shape[1]
+
+    Xi = torch.tensor(Xq, device=device).unsqueeze(1)  # bs x 1 x d
+    Xj = torch.tensor(Xa[I.flatten()], device=device)  # K * bs x d
+    Xj = Xj.reshape(*I.shape, Xq.shape[-1])  # bs x K x d
+
+    sim = torch.sum(Xi * Xj, dim=-1)  # bs x K
+    sort_idx = torch.argsort(sim, dim=1, descending=True).cpu().numpy()
+    I_refined, S_refined = [], []
+    for idx_i, sim_i, sort_i in zip(I, sim.cpu().numpy(), sort_idx):
+        I_refined.append(idx_i[sort_i][:k])
+        S_refined.append(sim_i[sort_i][:k])
+    I_refined = np.stack(I_refined)
+    S_refined = np.stack(S_refined)
+    
+    return S_refined, I_refined
+
+
+def test_nn_search():
+    key = np.random.random((3000000, 512)).astype(np.float32)
+    key /= np.linalg.norm(key, keepdims=True, axis=1)
+    index = build_faiss_index(key, -1)
+
+    query = np.random.random((100000, 512)).astype(np.float32)
+    query /= np.linalg.norm(query, keepdims=True, axis=1)
+    idx_r = nn_search(query, index, r=0.5)
+    idx_k = nn_search(query, index, topk=10)

model/.gitattributes

@@ -0,0 +1,2 @@
+*.hdf5 filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text

model/__init__.py

@@ -0,0 +1 @@
+from .gptk import get_gptk_model, get_gptk_image_transform

model/ckpt/mp_rank_00_model_states.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fab39af071b1e303f5976936a8662f75eb04952e03fa71bcb93291948892d2fd
+size 31462530292

model/eva_vit.py

@@ -0,0 +1,434 @@
+# Based on EVA, BEIT, timm and DeiT code bases
+# https://github.com/baaivision/EVA
+# https://github.com/rwightman/pytorch-image-models/tree/master/timm
+# https://github.com/microsoft/unilm/tree/master/beit
+# https://github.com/facebookresearch/deit/
+# https://github.com/facebookresearch/dino
+# --------------------------------------------------------'
+import math
+from functools import partial
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint as checkpoint
+from timm.models.layers import drop_path, to_2tuple, trunc_normal_
+
+import sys
+sys.path.append("./")
+from model.utils import download_cached_file
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
+    """
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+    
+    def extra_repr(self) -> str:
+        return 'p={}'.format(self.drop_prob)
+
+
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        # x = self.drop(x)
+        # commit this for the orignal BERT implement 
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(
+            self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0.,
+            proj_drop=0., window_size=None, attn_head_dim=None):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        if attn_head_dim is not None:
+            head_dim = attn_head_dim
+        all_head_dim = head_dim * self.num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, all_head_dim * 3, bias=False)
+        if qkv_bias:
+            self.q_bias = nn.Parameter(torch.zeros(all_head_dim))
+            self.v_bias = nn.Parameter(torch.zeros(all_head_dim))
+        else:
+            self.q_bias = None
+            self.v_bias = None
+
+        if window_size:
+            self.window_size = window_size
+            self.num_relative_distance = (2 * window_size[0] - 1) * (2 * window_size[1] - 1) + 3
+            self.relative_position_bias_table = nn.Parameter(
+                torch.zeros(self.num_relative_distance, num_heads))  # 2*Wh-1 * 2*Ww-1, nH
+            # cls to token & token 2 cls & cls to cls
+
+            # get pair-wise relative position index for each token inside the window
+            coords_h = torch.arange(window_size[0])
+            coords_w = torch.arange(window_size[1])
+            coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
+            coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
+            relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
+            relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
+            relative_coords[:, :, 0] += window_size[0] - 1  # shift to start from 0
+            relative_coords[:, :, 1] += window_size[1] - 1
+            relative_coords[:, :, 0] *= 2 * window_size[1] - 1
+            relative_position_index = \
+                torch.zeros(size=(window_size[0] * window_size[1] + 1, ) * 2, dtype=relative_coords.dtype)
+            relative_position_index[1:, 1:] = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
+            relative_position_index[0, 0:] = self.num_relative_distance - 3
+            relative_position_index[0:, 0] = self.num_relative_distance - 2
+            relative_position_index[0, 0] = self.num_relative_distance - 1
+
+            self.register_buffer("relative_position_index", relative_position_index)
+        else:
+            self.window_size = None
+            self.relative_position_bias_table = None
+            self.relative_position_index = None
+
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(all_head_dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x, rel_pos_bias=None):
+        B, N, C = x.shape
+        qkv_bias = None
+        if self.q_bias is not None:
+            qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias))
+        # qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
+        qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]   # make torchscript happy (cannot use tensor as tuple)
+
+        q = q * self.scale
+        attn = (q @ k.transpose(-2, -1))
+
+        if self.relative_position_bias_table is not None:
+            relative_position_bias = \
+                self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
+                    self.window_size[0] * self.window_size[1] + 1,
+                    self.window_size[0] * self.window_size[1] + 1, -1)  # Wh*Ww,Wh*Ww,nH
+            relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
+            attn = attn + relative_position_bias.unsqueeze(0)
+
+        if rel_pos_bias is not None:
+            attn = attn + rel_pos_bias
+        
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, -1)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class Block(nn.Module):
+
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., init_values=None, act_layer=nn.GELU, norm_layer=nn.LayerNorm,
+                 window_size=None, attn_head_dim=None):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
+            attn_drop=attn_drop, proj_drop=drop, window_size=window_size, attn_head_dim=attn_head_dim)
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+        if init_values is not None and init_values > 0:
+            self.gamma_1 = nn.Parameter(init_values * torch.ones((dim)),requires_grad=True)
+            self.gamma_2 = nn.Parameter(init_values * torch.ones((dim)),requires_grad=True)
+        else:
+            self.gamma_1, self.gamma_2 = None, None
+
+    def forward(self, x, rel_pos_bias=None):
+        if self.gamma_1 is None:
+            x = x + self.drop_path(self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias))
+            x = x + self.drop_path(self.mlp(self.norm2(x)))
+        else:
+            x = x + self.drop_path(self.gamma_1 * self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias))
+            x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))
+        return x
+
+
+class PatchEmbed(nn.Module):
+    """ Image to Patch Embedding
+    """
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
+        self.patch_shape = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+
+    def forward(self, x, **kwargs):
+        B, C, H, W = x.shape
+        # FIXME look at relaxing size constraints
+        assert H == self.img_size[0] and W == self.img_size[1], \
+            f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+        x = self.proj(x).flatten(2).transpose(1, 2)
+        return x
+
+
+class RelativePositionBias(nn.Module):
+
+    def __init__(self, window_size, num_heads):
+        super().__init__()
+        self.window_size = window_size
+        self.num_relative_distance = (2 * window_size[0] - 1) * (2 * window_size[1] - 1) + 3
+        self.relative_position_bias_table = nn.Parameter(
+            torch.zeros(self.num_relative_distance, num_heads))  # 2*Wh-1 * 2*Ww-1, nH
+        # cls to token & token 2 cls & cls to cls
+
+        # get pair-wise relative position index for each token inside the window
+        coords_h = torch.arange(window_size[0])
+        coords_w = torch.arange(window_size[1])
+        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
+        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
+        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
+        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
+        relative_coords[:, :, 0] += window_size[0] - 1  # shift to start from 0
+        relative_coords[:, :, 1] += window_size[1] - 1
+        relative_coords[:, :, 0] *= 2 * window_size[1] - 1
+        relative_position_index = \
+            torch.zeros(size=(window_size[0] * window_size[1] + 1,) * 2, dtype=relative_coords.dtype)
+        relative_position_index[1:, 1:] = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
+        relative_position_index[0, 0:] = self.num_relative_distance - 3
+        relative_position_index[0:, 0] = self.num_relative_distance - 2
+        relative_position_index[0, 0] = self.num_relative_distance - 1
+
+        self.register_buffer("relative_position_index", relative_position_index)
+
+        # trunc_normal_(self.relative_position_bias_table, std=.02)
+
+    def forward(self):
+        relative_position_bias = \
+            self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
+                self.window_size[0] * self.window_size[1] + 1,
+                self.window_size[0] * self.window_size[1] + 1, -1)  # Wh*Ww,Wh*Ww,nH
+        return relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
+
+
+class VisionTransformer(nn.Module):
+    """ Vision Transformer with support for patch or hybrid CNN input stage
+    """
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12,
+                 num_heads=12, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0.,
+                 drop_path_rate=0., norm_layer=nn.LayerNorm, init_values=None,
+                 use_abs_pos_emb=True, use_rel_pos_bias=False, use_shared_rel_pos_bias=False,
+                 use_mean_pooling=True, init_scale=0.001, use_checkpoint=False):
+        super().__init__()
+        self.image_size = img_size
+        self.num_classes = num_classes
+        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+
+        self.patch_embed = PatchEmbed(
+            img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
+        num_patches = self.patch_embed.num_patches
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        if use_abs_pos_emb:
+            self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
+        else:
+            self.pos_embed = None
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        if use_shared_rel_pos_bias:
+            self.rel_pos_bias = RelativePositionBias(window_size=self.patch_embed.patch_shape, num_heads=num_heads)
+        else:
+            self.rel_pos_bias = None
+        self.use_checkpoint = use_checkpoint
+        
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+        self.use_rel_pos_bias = use_rel_pos_bias
+        self.blocks = nn.ModuleList([
+            Block(
+                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
+                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer,
+                init_values=init_values, window_size=self.patch_embed.patch_shape if use_rel_pos_bias else None)
+            for i in range(depth)])
+#         self.norm = nn.Identity() if use_mean_pooling else norm_layer(embed_dim)
+#         self.fc_norm = norm_layer(embed_dim) if use_mean_pooling else None
+#         self.head = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+
+        if self.pos_embed is not None:
+            trunc_normal_(self.pos_embed, std=.02)
+        trunc_normal_(self.cls_token, std=.02)
+        # trunc_normal_(self.mask_token, std=.02)
+#         if isinstance(self.head, nn.Linear):
+#             trunc_normal_(self.head.weight, std=.02)
+        self.apply(self._init_weights)
+        self.fix_init_weight()
+#         if isinstance(self.head, nn.Linear):
+#             self.head.weight.data.mul_(init_scale)
+#             self.head.bias.data.mul_(init_scale)
+
+    def fix_init_weight(self):
+        def rescale(param, layer_id):
+            param.div_(math.sqrt(2.0 * layer_id))
+
+        for layer_id, layer in enumerate(self.blocks):
+            rescale(layer.attn.proj.weight.data, layer_id + 1)
+            rescale(layer.mlp.fc2.weight.data, layer_id + 1)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def get_classifier(self):
+        return self.head
+
+    def reset_classifier(self, num_classes, global_pool=''):
+        self.num_classes = num_classes
+        self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+
+    def forward_features(self, x):
+        x = self.patch_embed(x)
+        batch_size, seq_len, _ = x.size()
+
+        cls_tokens = self.cls_token.expand(batch_size, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
+        x = torch.cat((cls_tokens, x), dim=1)
+        if self.pos_embed is not None:
+            x = x + self.pos_embed
+        x = self.pos_drop(x)
+
+        rel_pos_bias = self.rel_pos_bias() if self.rel_pos_bias is not None else None
+        for blk in self.blocks:
+            if self.use_checkpoint:
+                x = checkpoint.checkpoint(blk, x, rel_pos_bias)
+            else:
+                x = blk(x, rel_pos_bias)
+        return x
+#         x = self.norm(x)
+
+#         if self.fc_norm is not None:
+#             t = x[:, 1:, :]
+#             return self.fc_norm(t.mean(1))
+#         else:
+#             return x[:, 0]
+
+    def forward(self, x):
+        x = self.forward_features(x)
+#         x = self.head(x)
+        return x
+
+    def get_intermediate_layers(self, x):
+        x = self.patch_embed(x)
+        batch_size, seq_len, _ = x.size()
+
+        cls_tokens = self.cls_token.expand(batch_size, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
+        x = torch.cat((cls_tokens, x), dim=1)
+        if self.pos_embed is not None:
+            x = x + self.pos_embed
+        x = self.pos_drop(x)
+
+        features = []
+        rel_pos_bias = self.rel_pos_bias() if self.rel_pos_bias is not None else None
+        for blk in self.blocks:
+            x = blk(x, rel_pos_bias)
+            features.append(x)
+
+        return features
+    
+    
+def interpolate_pos_embed(model, checkpoint_model):
+    if 'pos_embed' in checkpoint_model:
+        pos_embed_checkpoint = checkpoint_model['pos_embed'].float()
+        embedding_size = pos_embed_checkpoint.shape[-1]
+        num_patches = model.patch_embed.num_patches
+        num_extra_tokens = model.pos_embed.shape[-2] - num_patches
+        # height (== width) for the checkpoint position embedding
+        orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
+        # height (== width) for the new position embedding
+        new_size = int(num_patches ** 0.5)
+        # class_token and dist_token are kept unchanged
+        if orig_size != new_size:
+            print("Position interpolate from %dx%d to %dx%d" % (orig_size, orig_size, new_size, new_size))
+            extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
+            # only the position tokens are interpolated
+            pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
+            pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
+            pos_tokens = torch.nn.functional.interpolate(
+                pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)
+            pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
+            new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
+            checkpoint_model['pos_embed'] = new_pos_embed
+            
+            
+def convert_weights_to_fp16(model: nn.Module):
+    """Convert applicable model parameters to fp16"""
+
+    def _convert_weights_to_fp16(l):
+        if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
+            l.weight.data = l.weight.data.half()
+            if l.bias is not None:
+                l.bias.data = l.bias.data.half()
+
+#         if isinstance(l, (nn.MultiheadAttention, Attention)):
+#             for attr in [*[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]], "in_proj_bias", "bias_k", "bias_v"]:
+#                 tensor = getattr(l, attr)
+#                 if tensor is not None:
+#                     tensor.data = tensor.data.half()
+
+    model.apply(_convert_weights_to_fp16)
+    
+    
+def create_eva_vit_g(img_size=224,drop_path_rate=0.4,use_checkpoint=False,precision="fp16"):
+    model = VisionTransformer(
+        img_size=img_size,
+        patch_size=14,
+        use_mean_pooling=False,
+        embed_dim=1408,
+        depth=39,
+        num_heads=1408//88,
+        mlp_ratio=4.3637,
+        qkv_bias=True,
+        drop_path_rate=drop_path_rate,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        use_checkpoint=use_checkpoint,
+    )  
+    url = "https://storage.googleapis.com/sfr-vision-language-research/LAVIS/models/BLIP2/eva_vit_g.pth"
+    cached_file = download_cached_file(
+        url, check_hash=False, progress=True
+    )
+    state_dict = torch.load(cached_file, map_location="cpu")    
+    interpolate_pos_embed(model,state_dict)
+    
+    incompatible_keys = model.load_state_dict(state_dict, strict=False)
+#     print(incompatible_keys)
+    
+    if precision == "fp16":
+#         model.to("cuda") 
+        convert_weights_to_fp16(model)
+    return model

model/gptk-7b.yaml

@@ -0,0 +1,25 @@
+ # Copyright (c) 2022, salesforce.com, inc.
+ # All rights reserved.
+ # SPDX-License-Identifier: BSD-3-Clause
+ # For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
+
+arch: instruct_vicuna7b
+pretrained: "https://storage.googleapis.com/sfr-vision-language-research/LAVIS/models/InstructBLIP/instruct_blip_vicuna7b_trimmed.pth"
+
+# vit encoder
+image_size: 224
+drop_path_rate: 0
+use_grad_checkpoint: False
+vit_precision: "fp16"
+freeze_vit: True
+
+# Q-Former
+num_query_token: 32
+
+# path to Vicuna checkpoint
+llm_model: "model/llm/vicuna-7b-v1.1"
+# llm_model: "lmsys/vicuna-7b-v1.3"
+# llm_model: "lmsys/vicuna-7b-v1.5"
+
+# generation configs
+prompt: ""