Update app.py

app.py

@@ -62,6 +62,7 @@ repo_name = os.environ["LOG_REPO"]
 
 external_log_dir = "./logs"
 LOGDIR = external_log_dir
+VOTEDIR = "./votes"
 
 
 def install_gradio_4_35_0():
@@ -87,6 +88,38 @@ def get_conv_log_filename():
     name = os.path.join(LOGDIR, f"{t.year}-{t.month:02d}-{t.day:02d}-user_conv.json")
     return name
 
+def get_conv_vote_filename():
+    t = datetime.datetime.now()
+    name = os.path.join(VOTEDIR, f"{t.year}-{t.month:02d}-{t.day:02d}-user_vote.json")
+    if not os.path.isfile(name):
+        os.makedirs(os.path.dirname(name), exist_ok=True)
+    return name
+
+def vote_last_response(state, vote_type, model_selector):
+    with open(get_conv_vote_filename(), "a") as fout:
+        data = {
+            "type": vote_type,
+            "model": model_selector,
+            "state": state,
+        }
+        fout.write(json.dumps(data) + "\n")
+    api.upload_file(
+        path_or_fileobj=get_conv_vote_filename(),
+        path_in_repo=get_conv_vote_filename().replace("./votes/", ""),
+        repo_id=repo_name,
+        repo_type="dataset")
+
+
+def upvote_last_response(state):
+    vote_last_response(state, "upvote", "Pangea-7b")
+    gr.Info("Thank you for your voting!")
+    return state
+
+def downvote_last_response(state):
+    vote_last_response(state, "downvote", "Pangea-7b")
+    gr.Info("Thank you for your voting!")
+    return state
+
 class InferenceDemo(object):
     def __init__(
         self, args, model_path, tokenizer, model, image_processor, context_len
@@ -125,6 +158,22 @@ class InferenceDemo(object):
         self.conversation = conv_templates[args.conv_mode].copy()
         self.num_frames = args.num_frames
 
+class ChatSessionManager:
+    def __init__(self):
+        self.chatbot_instance = None
+
+    def initialize_chatbot(self, args, model_path, tokenizer, model, image_processor, context_len):
+        self.chatbot_instance = InferenceDemo(args, model_path, tokenizer, model, image_processor, context_len)
+        print(f"Initialized Chatbot instance with ID: {id(self.chatbot_instance)}")
+
+    def reset_chatbot(self):
+        self.chatbot_instance = None
+
+    def get_chatbot(self, args, model_path, tokenizer, model, image_processor, context_len):
+        if self.chatbot_instance is None:
+            self.initialize_chatbot(args, model_path, tokenizer, model, image_processor, context_len)
+        return self.chatbot_instance
+        
 
 def is_valid_video_filename(name):
     video_extensions = ["avi", "mp4", "mov", "mkv", "flv", "wmv", "mjpeg"]
@@ -178,13 +227,6 @@ def load_image(image_file):
     return image
 
 
-def clear_history(history):
-
-    our_chatbot.conversation = conv_templates[our_chatbot.conv_mode].copy()
-
-    return None
-
-
 def clear_response(history):
     for index_conv in range(1, len(history)):
         # loop until get a text response from our model.
@@ -195,40 +237,69 @@ def clear_response(history):
     history = history[:-index_conv]
     return history, question
 
+chat_manager = ChatSessionManager()
+
+
+def clear_history(history):
+    chatbot_instance = chat_manager.get_chatbot(args, model_path, tokenizer, model, image_processor, context_len)
+    chatbot_instance.conversation = conv_templates[chatbot_instance.conv_mode].copy()
+    return None
 
-# def print_like_dislike(x: gr.LikeData):
-#     print(x.index, x.value, x.liked)
 
 
 def add_message(history, message):
-    # history=[]
-    global our_chatbot
-    if len(history) == 0:
-        our_chatbot = InferenceDemo(
-            args, model_path, tokenizer, model, image_processor, context_len
-        )
+    global chat_image_num
+    if not history:
+        history = []
+        our_chatbot = chat_manager.get_chatbot(args, model_path, tokenizer, model, image_processor, context_len)
+        chat_image_num = 0
+        
+    if len(message["files"]) <= 1:
+        for x in message["files"]:
+            history.append(((x,), None))
+            chat_image_num += 1
+        if chat_image_num > 1:
+            history = []
+            chat_manager.reset_chatbot()
+            our_chatbot = chat_manager.get_chatbot(args, model_path, tokenizer, model, image_processor, context_len)
+            chat_image_num = 0
+            for x in message["files"]:
+                history.append(((x,), None))
+                chat_image_num += 1
+                
+        if message["text"] is not None:
+            history.append((message["text"], None))
+        
+        print(f"### Chatbot instance ID: {id(our_chatbot)}")
+        return history, gr.MultimodalTextbox(value=None, interactive=False)
+    else:
+        for x in message["files"]:
+            history.append(((x,), None))
+        if message["text"] is not None:
+            history.append((message["text"], None))
 
-    for x in message["files"]:
-        history.append(((x,), None))
-    if message["text"] is not None:
-        history.append((message["text"], None))
-    return history, gr.MultimodalTextbox(value=None, interactive=False)
+        return history, gr.MultimodalTextbox(value=None, interactive=False)
 
 
 @spaces.GPU
 def bot(history, temperature, top_p, max_output_tokens):
-    print("### turn start history",history)
-    print("### turn start conv",our_chatbot.conversation)
+    our_chatbot = chat_manager.get_chatbot(args, model_path, tokenizer, model, image_processor, context_len)
+    print(f"### Chatbot instance ID: {id(our_chatbot)}")
     text = history[-1][0]
     images_this_term = []
     text_this_term = ""
-    # import pdb;pdb.set_trace()
+    
     num_new_images = 0
+    previous_image = False
     for i, message in enumerate(history[:-1]):
         if type(message[0]) is tuple:
+            if previous_image:
+                gr.Warning("Only one image can be uploaded in a conversation. Please reduce the number of images and start a new conversation.")
+                our_chatbot.conversation = conv_templates[our_chatbot.conv_mode].copy()
+                return None
+           
             images_this_term.append(message[0][0])
             if is_valid_video_filename(message[0][0]):
-                # 不接受视频
                 raise ValueError("Video is not supported")
                 num_new_images += our_chatbot.num_frames
             elif is_valid_image_filename(message[0][0]):
@@ -236,15 +307,10 @@ def bot(history, temperature, top_p, max_output_tokens):
                 num_new_images += 1
             else:
                 raise ValueError("Invalid image file")
+            previous_image = True
         else:
             num_new_images = 0
-
-    # for message in history[-i-1:]:
-    #     images_this_term.append(message[0][0])
-
-    assert len(images_this_term) > 0, "must have an image"
-    # image_files = (args.image_file).split(',')
-    # image = [load_image(f) for f in images_this_term if f]
+            previous_image = False
 
     all_image_hash = []
     all_image_path = []
@@ -288,9 +354,7 @@ def bot(history, temperature, top_p, max_output_tokens):
 
     image_tensor = torch.stack(image_tensor)
     image_token = DEFAULT_IMAGE_TOKEN * num_new_images
-    # if our_chatbot.model.config.mm_use_im_start_end:
-    #     inp = DEFAULT_IM_START_TOKEN + image_token + DEFAULT_IM_END_TOKEN + "\n" + inp
-    # else:
+    
     inp = text
     inp = image_token + "\n" + inp
     our_chatbot.conversation.append_message(our_chatbot.conversation.roles[0], inp)
@@ -298,13 +362,6 @@ def bot(history, temperature, top_p, max_output_tokens):
     our_chatbot.conversation.append_message(our_chatbot.conversation.roles[1], None)
     prompt = our_chatbot.conversation.get_prompt()
 
-    # input_ids = (
-    #     tokenizer_image_token(
-    #         prompt, our_chatbot.tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt"
-    #     )
-    #     .unsqueeze(0)
-    #     .to(our_chatbot.model.device)
-    # )
     input_ids = tokenizer_image_token(
             prompt, our_chatbot.tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt"
         ).unsqueeze(0).to(our_chatbot.model.device)
@@ -318,9 +375,7 @@ def bot(history, temperature, top_p, max_output_tokens):
     stopping_criteria = KeywordsStoppingCriteria(
         keywords, our_chatbot.tokenizer, input_ids
     )
-    # streamer = TextStreamer(
-    #     our_chatbot.tokenizer, skip_prompt=True, skip_special_tokens=True
-    # )
+    
     streamer = TextIteratorStreamer(
         our_chatbot.tokenizer, skip_prompt=True, skip_special_tokens=True
     )
@@ -328,27 +383,6 @@ def bot(history, temperature, top_p, max_output_tokens):
     print(input_ids.device)
     print(image_tensor.device)
     
-    # with torch.inference_mode():
-    #     output_ids = our_chatbot.model.generate(
-    #         input_ids,
-    #         images=image_tensor,
-    #         do_sample=True,
-    #         temperature=0.7,
-    #         top_p=1.0,
-    #         max_new_tokens=4096,
-    #         streamer=streamer,
-    #         use_cache=False,
-    #         stopping_criteria=[stopping_criteria],
-    #     )
-
-    # outputs = our_chatbot.tokenizer.decode(output_ids[0]).strip()
-    # if outputs.endswith(stop_str):
-    #     outputs = outputs[: -len(stop_str)]
-    # our_chatbot.conversation.messages[-1][-1] = outputs
-
-    # history[-1] = [text, outputs]
-    
-    # return history
     generate_kwargs = dict(
         inputs=input_ids,
         streamer=streamer,
@@ -367,13 +401,12 @@ def bot(history, temperature, top_p, max_output_tokens):
     outputs = []
     for stream_token in streamer:
         outputs.append(stream_token)
-        # print("### stream_token",stream_token)
-        # our_chatbot.conversation.messages[-1][-1] = "".join(outputs)
+        
         history[-1] = [text, "".join(outputs)]
         yield history
     our_chatbot.conversation.messages[-1][-1] = "".join(outputs)
-    print("### turn end history", history)
-    print("### turn end conv",our_chatbot.conversation)
+    # print("### turn end history", history)
+    # print("### turn end conv",our_chatbot.conversation)
         
     with open(get_conv_log_filename(), "a") as fout:
         data = {
@@ -637,17 +670,25 @@ with gr.Blocks(
         gr.Markdown(learn_more_markdown)
         gr.Markdown(bibtext)
 
-    chat_msg = chat_input.submit(
-        add_message, [chatbot, chat_input], [chatbot, chat_input]
-    )
-    bot_msg = chat_msg.then(bot, [chatbot,temperature, top_p, max_output_tokens], chatbot, api_name="bot_response")
-    bot_msg.then(lambda: gr.MultimodalTextbox(interactive=True), None, [chat_input])
+    chat_input.submit(
+    add_message, [chatbot, chat_input], [chatbot, chat_input]
+    ).then(bot, [chatbot, temperature, top_p, max_output_tokens], chatbot, api_name="bot_response").then(lambda: gr.MultimodalTextbox(interactive=True), None, [chat_input])
+  
 
     # chatbot.like(print_like_dislike, None, None)
     clear_btn.click(
         fn=clear_history, inputs=[chatbot], outputs=[chatbot], api_name="clear_all"
     )
 
+    upvote_btn.click(
+        fn=upvote_last_response, inputs=chatbot, outputs=chatbot, api_name="upvote_last_response"
+    )
+
+    
+    downvote_btn.click(
+        fn=downvote_last_response, inputs=chatbot, outputs=chatbot, api_name="upvote_last_response"
+    )
+
 
 demo.queue()
     
@@ -678,5 +719,5 @@ if __name__ == "__main__":
     model_name = get_model_name_from_path(args.model_path)
     tokenizer, model, image_processor, context_len = load_pretrained_model(args.model_path, args.model_base, model_name, args.load_8bit, args.load_4bit)
     model=model.to(torch.device('cuda'))
-    our_chatbot = None
+    chat_image_num = 0
     demo.launch()