some quick fix

app.py

@@ -22,6 +22,7 @@ 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."
 
 
@@ -29,6 +30,10 @@ 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 True
+    
     url = "https://api.openai.com/v1/moderations"
     headers = {
         "Content-Type": "application/json",
@@ -60,31 +65,32 @@ def regenerate(state: Conversation):
         prev_human_msg[1] = prev_human_msg[1][:2]
     state.skip_next = False
     
-    return (state, state.to_gradio_chatbot(), "", None, disable_btn, disable_btn)
+    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) + (disable_btn,) * 2 + knwl_none
+    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,) * 2
+        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,) * 2
+        return (state, state.to_gradio_chatbot(), moderation_msg, None) + (no_change_btn,) * 3
 
-    text = (text, image)
-    if len(state.get_images(return_pil=True)) > 0:
-        state = default_conversation.copy()
+    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,) * 2
+    return (state, state.to_gradio_chatbot(), "", None) + (disable_btn,) * 3
 
 
 def search(image, pos, topk, knwl_db, knwl_idx):
@@ -150,9 +156,10 @@ def retrieve_knowledge(image):
     return knwl_embd, knwl_text
 
 
-def generate(state, temperature, top_p, max_new_tokens, add_knwl, do_sampling, do_beam_search):
+@torch.inference_mode()
+def generate(state: Conversation, 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
+        yield (state, state.to_gradio_chatbot()) + (no_change_btn,) * 3 + knwl_unchange
         return
 
     if len(state.messages) == state.offset + 2:  # First round of conversation
@@ -177,11 +184,16 @@ def generate(state, temperature, top_p, max_new_tokens, add_knwl, do_sampling, d
             for pos in range(knwl_pos):
                 try:
                     txt = ""
-                    for k, v in knwl[query_type][str(pos)].items():
+                    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
-                    knwl_img[idx] = images[pos]
+
+                    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
@@ -189,13 +201,13 @@ def generate(state, temperature, top_p, max_new_tokens, add_knwl, do_sampling, d
     else:
         knwl_embd = None
         knwl_vis = knwl_none
+    yield (state, state.to_gradio_chatbot()) + (disable_btn,) * 3 + knwl_vis
 
     # generate output
-    prompt = state.get_prompt()
+    prompt = state.get_prompt().replace("USER: <image>\n", "")
     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,
@@ -203,6 +215,7 @@ def generate(state, temperature, top_p, max_new_tokens, add_knwl, do_sampling, d
             max_length=min(int(max_new_tokens), 1024),
             top_p=float(top_p),
             temperature=float(temperature),
+            length_penalty=0.0,
             auto_cast=True
         )[0]
         streamer = [new_text, ]
@@ -220,6 +233,7 @@ def generate(state, temperature, top_p, max_new_tokens, add_knwl, do_sampling, d
                 temperature=float(temperature),
                 streamer=streamer,
                 num_beams=1,
+                length_penalty=0.0,
                 auto_cast=True
             )
         )
@@ -229,10 +243,10 @@ def generate(state, temperature, top_p, max_new_tokens, add_knwl, do_sampling, d
     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
+        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,) * 2 + knwl_vis
+    yield (state, state.to_gradio_chatbot()) + (enable_btn,) * 3 + knwl_unchange
         
 
 title_markdown = ("""
@@ -268,15 +282,18 @@ def build_demo():
                     ["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.Row():
-                    submit_btn = gr.Button(value="📝 Submit")
-                    regenerate_btn = gr.Button(value="🔄 Regenerate", interactive=False)
-                    clear_btn = gr.Button(value="🗑️ Clear", interactive=False)
+                    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):
                     with gr.Row():
@@ -290,7 +307,7 @@ def build_demo():
             with gr.Column(scale=6):
                 chatbot = gr.Chatbot(elem_id="chatbot", label="GPT-K Chatbot", height=550)
             
-                gr.Markdown("Retrieved Knowledge")
+                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):
@@ -307,7 +324,7 @@ def build_demo():
         gr.Markdown(learn_more_markdown)
 
         # Register listeners
-        btn_list = [regenerate_btn, clear_btn]
+        btn_list = [submit_btn, regenerate_btn, clear_btn]
         regenerate_btn.click(
             regenerate, [state], [state, chatbot, textbox, imagebox] + btn_list
         ).then(

model/ckpt/gptk-vicuna7b.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:bb27e6bbdc6f93ac950d265287c8388824f106e17bcab5cf5254810ca9c6790f
-size 564340835
+oid sha256:502c7a85d5e0d17eb4e823ed8779565dbac832fa16fd659b69e39b9b024c9d26
+size 564340993

requirements.txt

@@ -1,9 +1,13 @@
-h5py>=3.8.0
-transformers==4.30.2
-faiss-gpu==1.7.2
-timm==0.4.12
-openai
 --extra-index-url https://download.pytorch.org/whl/cu113
 torch==1.11.0+cu113
 torchvision==0.12.0+cu113
 torchaudio==0.11.0
+
+transformers==4.30.2
+faiss-gpu==1.7.2
+timm==0.4.12
+openai
+open_clip_torch
+omegaconf
+h5py>=3.8.0
+spacy>=3.5.0