Update user agent and enhance predict function with additional parameters for improved flexibility

app.py

@@ -37,7 +37,7 @@ torch.set_default_device('cuda')
 
 logger = build_logger("gradio_web_server", "gradio_web_server.log")
 
-headers = {"User-Agent": "Vintern-Chat Client"}
+headers = {"User-Agent": "Vintern-1B-3.5-Demo Client"}
 
 no_change_btn = gr.Button()
 enable_btn = gr.Button(interactive=True)
@@ -181,21 +181,21 @@ model = AutoModel.from_pretrained(
 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, use_fast=False)
 
 @spaces.GPU
-def predict(message, image_path):
-        pixel_values = load_image(image_path, max_num=6).to(torch.bfloat16).cuda()
-        print(f"pixel_values: {pixel_values}")
-        generation_config = dict(max_new_tokens= 700, do_sample=False, num_beams = 3, repetition_penalty=2.5)
-        print(f"######################")
-        print(f"message: {message}")
+def predict(message, 
+            image_path, 
+            history,
+            max_input_tiles=6, 
+            temperature=1.0,
+            max_output_tokens=700,
+            top_p=0.7,
+            repetition_penalty=2.5):
+        pixel_values = load_image(image_path, max_num=max_input_tiles).to(torch.bfloat16).cuda()
+        generation_config = dict(temperature=temperature, max_new_tokens= max_output_tokens, top_p=top_p, do_sample=False, num_beams = 3, repetition_penalty=repetition_penalty)
         if pixel_values is not None:
             question = '<image>\n'+message
         else:
             question = message
-        print("Model: ", model)
-        print("Tokenizer: ", tokenizer)
-        print("Question: ", question)
-        response, conv_history = model.chat(tokenizer, pixel_values, question, generation_config, history=None, return_history=True)
-        print(f"AI response: {response}")
+        response, conv_history = model.chat(tokenizer, pixel_values, question, generation_config, history=history, return_history=True)
         return response, conv_history
 
 def http_bot(
@@ -220,7 +220,6 @@ def http_bot(
         return
 
     if model is None:
-        # state.messages[-1][-1] = server_error_msg
         state.update_message(Conversation.ASSISTANT, server_error_msg)
         yield (
             state,
@@ -246,26 +245,15 @@ def http_bot(
 
     try:
         # Stream output
-        # response = requests.post(worker_addr, json=pload, headers=headers, stream=True, timeout=300)
-        print(f"all_image_paths: {all_image_paths}")
         message = state.get_user_message(source=state.USER)
-        print(f"message: {message}")
-        response, conv_history = predict(message, all_image_paths[0])
+        logger.info(f"==== User message ====\n{message}")
+        logger.info(f"==== Image paths ====\n{all_image_paths}")
+        response, conv_history = predict(message, all_image_paths[0], max_input_tiles, temperature, max_new_tokens, top_p, repetition_penalty)
 
-        # streamer = TextIteratorStreamer(
-        #     tokenizer, skip_prompt=True, skip_special_tokens=True
-        # )
-        # generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024)
-
-        # thread = Thread(target=model.generate, kwargs=generation_kwargs)
-        # thread.start()
-        
         # response = "This is a test response"
         buffer = ""
         for new_text in response:
             buffer += new_text
-            # Remove <|im_end|> or similar tokens from the output
-            buffer = buffer.replace("<|im_end|>", "")
                 
             state.update_message(Conversation.ASSISTANT, buffer + state.streaming_placeholder, None)
             yield (
@@ -275,8 +263,7 @@ def http_bot(
             ) + (disable_btn,) * 5
 
     except Exception as e:
-        logger.error(f"Error in http_bot: {e}")
-        traceback.print_exc()
+        logger.error(f"Error in http_bot: {e} \n{traceback.format_exc()}")
         state.update_message(Conversation.ASSISTANT, server_error_msg, None)
         yield (
             state,
@@ -293,7 +280,7 @@ def http_bot(
 
     ai_response = state.return_last_message()
 
-    logger.info(f"==== response ====\n{ai_response}")
+    logger.info(f"==== AI response ====\n{ai_response}")
    
     state.end_of_current_turn()
 
@@ -321,9 +308,10 @@ def http_bot(
 title_html = """
 <div style="text-align: center;">
     <img src="https://lh3.googleusercontent.com/pw/AP1GczMmW-aFQ4dNaR_LCAllh4UZLLx9fTZ1ITHeGVMWx-1bwlIWz4VsWJSGb3_9C7CQfvboqJH41y2Sbc5ToC9ZmKeV4-buf_DEevIMU0HtaLWgHAPOqBiIbG6LaE8CvDqniLZzvB9UX8TR_-YgvYzPFt2z=w1472-h832-s-no-gm?authuser=0" style="height: 100; width: 100%;">
+    <p>🔥Vintern-1B-v3_5🔥</p>
     <p>Vintern-1B: An Efficient Multimodal Large Language Model for Vietnamese</p>
     <a href="https://huggingface.co/papers/2408.12480">[📖 Vintern Paper]</a>
-    <a href="https://huggingface.co/5CD-AI">[🤗 5CD-AI Huggingface]</a>
+    <a href="https://huggingface.co/5CD-AI">[🤗 Huggingface]</a>
 </div>
 """
 
@@ -411,7 +399,7 @@ def build_demo():
     )
 
     with gr.Blocks(
-        title="Vintern-Chat",
+        title="Vintern-1B-v3_5-Demo",
         theme=gr.themes.Default(),
         css=block_css,
     ) as demo:
@@ -424,7 +412,7 @@ def build_demo():
 
                 with gr.Accordion("Settings", open=False) as setting_row:
                     system_prompt = gr.Textbox(
-                        value="请尽可能详细地回答用户的问题。",
+                        value="Bạn là một trợ lý AI đa phương thức hữu ích, hãy trả lời câu hỏi người dùng một cách chi tiết.",
                         label="System Prompt",
                         interactive=True,
                     )
@@ -501,7 +489,7 @@ def build_demo():
             with gr.Column(scale=8):
                 chatbot = gr.Chatbot(
                     elem_id="chatbot",
-                    label="Vintern",
+                    label="Vintern-1B-v3_5-Demo",
                     height=580,
                     show_copy_button=True,
                     show_share_button=True,