Update src/app.py

src/app.py

@@ -1,51 +1,29 @@
 """Developed by Ruslan Magana Vsevolodovna"""
-
 from collections.abc import Iterator
 from datetime import datetime
 from pathlib import Path
 from threading import Thread
-
 import gradio as gr
 import spaces
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 from transformers import LlavaNextProcessor, LlavaNextForConditionalGeneration
 import random
-
 from themes.research_monochrome import theme
 
 # =============================================================================
 # Constants & Prompts
 # =============================================================================
-today_date = datetime.today().strftime("%B %-d, %Y") 
-SYS_PROMPT = """
-Today's Date: {today_date}.
-You are a helpful AI assistant.
-Respond in the following format:
-<reasoning>
-...
-</reasoning>
-<answer>
-...
-</answer>
-"""
-
+today_date = datetime.today().strftime("%B %-d, %Y")  # noqa: DTZ002
+SYS_PROMPT = f"""Knowledge Cutoff Date: April 2024.Today's Date: {today_date}.You are Granite, developed by IBM. You are a helpful AI assistant. Respond in the following format:<reasoning>Step-by-step reasoning to arrive at the answer.</reasoning><answer>The final answer to the user's query.</answer> If reasoning is not applicable, you can directly provide the <answer>."""
 TITLE = "IBM Granite 3.1 8b Reasoning & Vision Preview"
-DESCRIPTION = """
-<p>Granite 3.1 8b Reasoning is an open‐source LLM supporting a 128k context window and Granite Vision 3.1 2B Preview for vision‐language capabilities. Start with one of the sample prompts
-or enter your own. Keep in mind that AI can occasionally make mistakes.
-<span class="gr_docs_link">
-<a href="https://www.ibm.com/granite/docs/">View Documentation <i class="fa fa-external-link"></i></a>
-</span>
-</p>
-"""
+DESCRIPTION = """<p>Granite 3.1 8b Reasoning is an open‐source LLM supporting a 128k context window and Granite Vision 3.1 2B Preview for vision‐language capabilities. Start with one of the sample promptsor enter your own. Keep in mind that AI can occasionally make mistakes.<span class="gr_docs_link"><a href="https://www.ibm.com/granite/docs/">View Documentation <i class="fa fa-external-link"></i></a></span></p>"""
 MAX_INPUT_TOKEN_LENGTH = 128_000
 MAX_NEW_TOKENS = 1024
 TEMPERATURE = 0.7
 TOP_P = 0.85
 TOP_K = 50
 REPETITION_PENALTY = 1.05
-
 # Vision defaults (advanced settings)
 VISION_TEMPERATURE = 0.2
 VISION_TOP_P = 0.95
@@ -54,18 +32,13 @@ VISION_MAX_TOKENS = 128
 
 if not torch.cuda.is_available():
     print("This demo may not work on CPU.")
-
 # =============================================================================
 # Text Model Loading
 # =============================================================================
-
 #Standard Model
 #granite_text_model="ibm-granite/granite-3.1-8b-instruct"
-
 #With Reasoning
 granite_text_model="ruslanmv/granite-3.1-8b-Reasoning"
-
-
 text_model = AutoModelForCausalLM.from_pretrained(
     granite_text_model,
     torch_dtype=torch.float16,
@@ -73,7 +46,6 @@ text_model = AutoModelForCausalLM.from_pretrained(
 )
 tokenizer = AutoTokenizer.from_pretrained(granite_text_model)
 tokenizer.use_default_system_prompt = False
-
 # =============================================================================
 # Vision Model Loading
 # =============================================================================
@@ -83,9 +55,8 @@ vision_model = LlavaNextForConditionalGeneration.from_pretrained(
     vision_model_path,
     torch_dtype=torch.float16,
     device_map="auto",
-    trust_remote_code=True  # Ensure the custom code is used so that weight shapes match.
+    trust_remote_code=True  # Ensure the custom code is used so that weight shapes match.)
 )
-
 # =============================================================================
 # Text Generation Function (for text-only chat)
 # =============================================================================
@@ -99,7 +70,7 @@ def generate(
     top_k: float = TOP_K,
     max_new_tokens: int = MAX_NEW_TOKENS,
 ) -> Iterator[str]:
-    """Generate function for text chat demo."""
+    """Generate function for text chat demo with chain of thought display."""
     conversation = []
     conversation.append({"role": "system", "content": SYS_PROMPT})
     conversation.extend(chat_history)
@@ -126,10 +97,45 @@ def generate(
     )
     t = Thread(target=text_model.generate, kwargs=generate_kwargs)
     t.start()
+
     outputs = []
+    reasoning_started = False
+    answer_started = False
+    collected_reasoning = ""
+    collected_answer = ""
+
     for text in streamer:
         outputs.append(text)
-        yield "".join(outputs)
+        current_output = "".join(outputs)
+
+        if "<reasoning>" in current_output and not reasoning_started:
+            reasoning_started = True
+            reasoning_start_index = current_output.find("<reasoning>") + len("<reasoning>")
+            collected_reasoning = current_output[reasoning_start_index:]
+            yield "[Reasoning]: "  # Indicate start of reasoning in chatbot
+            outputs = [collected_reasoning] # Reset outputs to only include reasoning part
+
+        elif reasoning_started and "<answer>" in current_output and not answer_started:
+            answer_started = True
+            reasoning_end_index = current_output.find("<answer>")
+            collected_reasoning = current_output[len("<reasoning>"):reasoning_end_index] # Correctly extract reasoning part
+
+            answer_start_index = current_output.find("<answer>") + len("<answer>")
+            collected_answer = current_output[answer_start_index:]
+            yield "\n[Answer]: " # Indicate start of answer in chatbot
+            outputs = [collected_answer] # Reset outputs to only include answer part
+            yield collected_answer # Yield initial part of answer
+
+        elif reasoning_started and not answer_started:
+             collected_reasoning = text # Accumulate reasoning tokens
+             yield text # Stream reasoning tokens
+
+        elif answer_started:
+            collected_answer += text # Accumulate answer tokens
+            yield text # Stream answer tokens
+        else:
+            yield text # In case no tags are found, stream as before
+
 
 # =============================================================================
 # Vision Chat Inference Function (for image+text chat)
@@ -172,7 +178,30 @@ def chat_inference(image, text, conversation, temperature=VISION_TEMPERATURE, to
     }
     output = vision_model.generate(**inputs, **generation_kwargs)
     assistant_response = vision_processor.decode(output[0], skip_special_tokens=True)
-    conversation.append({"role": "assistant", "content": [{"type": "text", "text": assistant_response.strip()}]})
+
+    reasoning = ""
+    answer = ""
+    if "<reasoning>" in assistant_response and "<answer>" in assistant_response:
+        reasoning_start = assistant_response.find("<reasoning>") + len("<reasoning>")
+        reasoning_end = assistant_response.find("</reasoning>")
+        reasoning = assistant_response[reasoning_start:reasoning_end].strip()
+
+        answer_start = assistant_response.find("<answer>") + len("<answer>")
+        answer_end = assistant_response.find("</answer>")
+
+        if answer_end != -1: # Handle cases where answer end tag is present
+            answer = assistant_response[answer_start:answer_end].strip()
+        else: # Fallback if answer end tag is missing (less robust)
+            answer = assistant_response[answer_start:].strip()
+
+
+    formatted_response_content = []
+    if reasoning:
+        formatted_response_content.append({"type": "text", "text": f"[Reasoning]: {reasoning}"})
+    formatted_response_content.append({"type": "text", "text": f"[Answer]: {answer}"})
+
+
+    conversation.append({"role": "assistant", "content": formatted_response_content})
     return display_vision_conversation(conversation), conversation
 
 # =============================================================================
@@ -206,7 +235,12 @@ def display_vision_conversation(conversation):
             assistant_msg = ""
             if i + 1 < len(conversation) and conversation[i+1]["role"] == "assistant":
                 # Extract assistant text; remove any special tokens if present.
-                assistant_msg = conversation[i+1]["content"][0]["text"].split("<|assistant|>")[-1].strip()
+                assistant_content = conversation[i+1]["content"]
+                assistant_text_parts = []
+                for item in assistant_content:
+                     if item["type"] == "text":
+                        assistant_text_parts.append(item["text"])
+                assistant_msg = "\n".join(assistant_text_parts).strip()
                 i += 2
             else:
                 i += 1
@@ -214,7 +248,6 @@ def display_vision_conversation(conversation):
         else:
             i += 1
     return chat_history
-
 # =============================================================================
 # Unified Send-Message Function
 # =============================================================================
@@ -251,29 +284,28 @@ def send_message(image, text,
             top_k=text_top_k,
             max_new_tokens=text_max_new_tokens
         ):
-            output_text = chunk
+            output_text += chunk # Accumulate for display function to process correctly.
+
         conv.append({"role": "user", "content": text})
-        conv.append({"role": "assistant", "content": output_text})
+        conv.append({"role": "assistant", "content": output_text}) # Store full output with tags
         text_state = conv
-        chat_history = display_text_conversation(text_state)
+        chat_history = display_text_conversation(text_state) # Display function handles tag parsing now.
         return chat_history, text_state, vision_state
 
 def clear_chat():
     # Clear the conversation and input fields.
     return [], [], [], None  # (chat_history, text_state, vision_state, cleared text and image inputs)
-
 # =============================================================================
 # UI Layout with Gradio
 # =============================================================================
 css_file_path = Path(Path(__file__).parent / "app.css")
 head_file_path = Path(Path(__file__).parent / "app_head.html")
-
 with gr.Blocks(fill_height=True, css_paths=css_file_path, head_paths=head_file_path, theme=theme, title=TITLE) as demo:
     gr.HTML(f"<h1>{TITLE}</h1>", elem_classes=["gr_title"])
     gr.HTML(DESCRIPTION)
-    
+
     chatbot = gr.Chatbot(label="Chat History", height=500)
-    
+
     with gr.Row():
         with gr.Column(scale=2):
             image_input = gr.Image(type="pil", label="Upload Image (optional)")
@@ -290,14 +322,13 @@ with gr.Blocks(fill_height=True, css_paths=css_file_path, head_paths=head_file_p
                 vision_top_p_slider = gr.Slider(minimum=0.0, maximum=1.0, value=VISION_TOP_P, step=0.01, label="Vision Top p", elem_classes=["gr_accordion_element"])
                 vision_top_k_slider = gr.Slider(minimum=0, maximum=100, value=VISION_TOP_K, step=1, label="Vision Top k", elem_classes=["gr_accordion_element"])
                 vision_max_tokens_slider = gr.Slider(minimum=10, maximum=300, value=VISION_MAX_TOKENS, step=1, label="Vision Max Tokens", elem_classes=["gr_accordion_element"])
-    
-    send_button = gr.Button("Send Message")
+        send_button = gr.Button("Send Message")
     clear_button = gr.Button("Clear Chat")
-    
+
     # Conversation state variables for each branch.
     text_state = gr.State([])
     vision_state = gr.State([])
-    
+
     send_button.click(
         send_message,
         inputs=[
@@ -308,13 +339,13 @@ with gr.Blocks(fill_height=True, css_paths=css_file_path, head_paths=head_file_p
         ],
         outputs=[chatbot, text_state, vision_state]
     )
-    
+
     clear_button.click(
         clear_chat,
         inputs=None,
         outputs=[chatbot, text_state, vision_state, text_input, image_input]
     )
-    
+
     gr.Examples(
         examples=[
             ["https://raw.githubusercontent.com/gradio-app/gradio/main/test/test_files/cheetah1.jpg", "What is in this image?"],
@@ -339,4 +370,4 @@ with gr.Blocks(fill_height=True, css_paths=css_file_path, head_paths=head_file_p
     )
 
 if __name__ == "__main__":
-    demo.queue().launch()
+    demo.queue().launch()