Update app.py

app.py

@@ -136,64 +136,85 @@ model = AutoModel.from_pretrained(
 tokenizer = AutoTokenizer.from_pretrained("5CD-AI/Vintern-3B-R-beta", trust_remote_code=True, use_fast=False)
     
 @spaces.GPU
-def chat(message, history):
-    print("history",history)
-    print("message",message)
-
-    if len(history) != 0 and len(message["files"]) != 0:
-        return """Chúng tôi hiện chỉ hổ trợ 1 ảnh ở đầu ngữ cảnh! Vui lòng tạo mới cuộc trò chuyện.
-We currently only support one image at the start of the context! Please start a new conversation."""
+def chat(message, history, think_mode):
+    if not think_mode:
+        print("history",history)
+        print("message",message)
     
-    if len(history) == 0 and len(message["files"]) != 0:
-        if "path" in message["files"][0]:
-            test_image = message["files"][0]["path"]
+        if len(history) != 0 and len(message["files"]) != 0:
+            return """Chúng tôi hiện chỉ hổ trợ 1 ảnh ở đầu ngữ cảnh! Vui lòng tạo mới cuộc trò chuyện.
+    We currently only support one image at the start of the context! Please start a new conversation."""
+        
+        if len(history) == 0 and len(message["files"]) != 0:
+            if "path" in message["files"][0]:
+                test_image = message["files"][0]["path"]
+            else:
+                test_image = message["files"][0]
+            pixel_values = load_image(test_image, max_num=6).to(torch.bfloat16).cuda()
+        elif len(history) == 0 and len(message["files"]) == 0:
+            pixel_values = None     
+        elif history[0][0][0] is not None and os.path.isfile(history[0][0][0]):
+            test_image = history[0][0][0]
+            pixel_values = load_image(test_image, max_num=6).to(torch.bfloat16).cuda()
         else:
-            test_image = message["files"][0]
-        pixel_values = load_image(test_image, max_num=6).to(torch.bfloat16).cuda()
-    elif len(history) == 0 and len(message["files"]) == 0:
-        pixel_values = None     
-    elif history[0][0][0] is not None and os.path.isfile(history[0][0][0]):
-        test_image = history[0][0][0]
-        pixel_values = load_image(test_image, max_num=6).to(torch.bfloat16).cuda()
-    else:
-        pixel_values = None 
+            pixel_values = None 
+            
         
-    
-    generation_config = dict(max_new_tokens= 700, do_sample=False, num_beams = 3, repetition_penalty=2.5)
-    
-    if len(history) == 0:
-        if pixel_values is not None:
-            question = '<image>\n'+message["text"]
+        generation_config = dict(max_new_tokens= 700, do_sample=False, num_beams = 3, repetition_penalty=2.5)
+        
+        if len(history) == 0:
+            if pixel_values is not None:
+                question = '<image>\n'+message["text"]
+            else:
+                question = message["text"]
+            response, conv_history = model.chat(tokenizer, pixel_values, question, generation_config, history=None, return_history=True)
         else:
+            conv_history = []
+            if history[0][0][0] is not None and os.path.isfile(history[0][0][0]):
+                start_index = 1
+            else:
+                start_index = 0
+            
+            for i, chat_pair in enumerate(history[start_index:]):
+                if i == 0 and start_index == 1:
+                     conv_history.append(tuple(['<image>\n'+chat_pair[0],chat_pair[1]]))
+                else:
+                    conv_history.append(tuple(chat_pair))
+    
+                
+            print("conv_history",conv_history)
             question = message["text"]
-        response, conv_history = model.chat(tokenizer, pixel_values, question, generation_config, history=None, return_history=True)
+            response, conv_history = model.chat(tokenizer, pixel_values, question, generation_config, history=conv_history, return_history=True)
+            
+        print(f'User: {question}\nAssistant: {response}')
+    
+        # return response
+        buffer = ""
+        for new_text in response:
+          buffer += new_text
+          generated_text_without_prompt = buffer[:]
+          time.sleep(0.02)
+          yield generated_text_without_prompt
     else:
-        conv_history = []
-        if history[0][0][0] is not None and os.path.isfile(history[0][0][0]):
-            start_index = 1
-        else:
-            start_index = 0
-        
-        for i, chat_pair in enumerate(history[start_index:]):
-            if i == 0 and start_index == 1:
-                 conv_history.append(tuple(['<image>\n'+chat_pair[0],chat_pair[1]]))
-            else:
-                conv_history.append(tuple(chat_pair))
+        buffer = ""
+        thinking = """🔄 Model is analyzing the data...\n
+I am processing your request carefully. First, I need to understand the question clearly.  
+Then, I retrieve relevant information and analyze different possibilities.  
+Finally, I generate a structured response that best fits your input.  
+\nThis process ensures that I provide the most accurate and meaningful answer possible.
+"""
 
-            
-        print("conv_history",conv_history)
-        question = message["text"]
-        response, conv_history = model.chat(tokenizer, pixel_values, question, generation_config, history=conv_history, return_history=True)
-        
-    print(f'User: {question}\nAssistant: {response}')
+        accumulated_text = "💡 **Thinking process:** *(Click to expand)*\n\n"
+        accumulated_text += "<pre><code>\n"
+
+        temp_text = ""
+        for char in thinking:
+            temp_text += char
+            yield accumulated_text + temp_text + "\n</code></pre>\n\n---\n"
+            time.sleep(0.02)
+
+        accumulated_text += temp_text + "\n</code></pre>\n\n---\n"
 
-    # return response
-    buffer = ""
-    for new_text in response:
-      buffer += new_text
-      generated_text_without_prompt = buffer[:]
-      time.sleep(0.02)
-      yield generated_text_without_prompt
 
 CSS ="""
 #component-10 {
@@ -250,7 +271,16 @@ button.svelte-1lcyrx4[aria-label="user's message: a file of type image/jpeg, "]
     margin-bottom: var(--spacing-lg);
 }
 
-#think-button { position: absolute; bottom: 10px; left: 10px; font-size: 12px; padding: 5px 10px; }
+#think-button { 
+    position: absolute; 
+    bottom: 20px; /* Dời xuống 10px nữa */
+    left: 10px; 
+    font-size: 12px; 
+    padding: 5px 10px; 
+    border-radius: 10px; /* Bo tròn góc */
+    background-color: #f4f4f4; /* Màu nền nhẹ */
+    border: 1px solid #ccc; /* Viền nhẹ */
+}
 """
 
 js = """
@@ -290,7 +320,7 @@ demo = gr.Blocks(css=CSS,js=js, theme='NoCrypt/miku')
 # Hàm bật/tắt Think Mode
 def toggle_think_mode(current_state):
     new_state = not current_state
-    button_label = "💡Think✅" if new_state else "💡Think"
+    button_label = "🧠Think💡" if new_state else "🧠Think"
     return new_state, button_label
 
     
@@ -315,7 +345,7 @@ with demo:
     )
     
     with gr.Row(variant="compact"):
-        think_button = gr.Button("💡Think", elem_id="think-button", variant="secondary")  
+        think_button = gr.Button("🧠Think", elem_id="think-button", variant="secondary")  
 
     # Khi nhấn nút, trạng thái think_mode thay đổi + đổi nhãn nút
     think_button.click(toggle_think_mode, inputs=[think_mode], outputs=[think_mode, think_button])