Update src/app.py

src/app.py

@@ -1,131 +1,105 @@
 """Template Demo for IBM Granite Hugging Face spaces."""
 
-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
-
-# Vision model imports
-from transformers import LlavaNextProcessor, LlavaNextForConditionalGeneration
-import random
-
-from themes.research_monochrome import theme
-
-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"""
-TITLE = "IBM Granite 3.1 8b Instruct & Vision Preview"
-DESCRIPTION = """
-<p>Granite 3.1 8b instruct 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. Upload an image to use the vision model. Keep in mind that AI can occasionally make mistakes.
-<span class="gr_docs_link">
-<a href="https://www.ibm.com/granite/docs/">View Granite Instruct Documentation <i class="fa fa-external-link"></i></a>
-</span>
-<span class="gr_docs_link">
-<a href="https://www.ibm.com/granite/vision/docs/">View Granite Vision 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_TEMPERATURE = 0.2
-VISION_TOP_P = 0.95
-VISION_TOP_K = 50
-VISION_MAX_TOKENS = 128
-
-
-if not torch.cuda.is_available():
-    print("This demo may not work on CPU.")
-
-# Text model loading
-text_model = AutoModelForCausalLM.from_pretrained(
-    "ibm-granite/granite-3.1-8b-instruct", torch_dtype=torch.float16, device_map="auto"
-)
-text_tokenizer = AutoTokenizer.from_pretrained("ibm-granite/granite-3.1-8b-instruct")
-text_tokenizer.use_default_system_prompt = False
-
-# Vision model loading
-vision_model_path = "ibm-granite/granite-vision-3.1-2b-preview"
-vision_processor = LlavaNextProcessor.from_pretrained(vision_model_path, use_fast=True)
-
-# Option 1: Use the default settings (like the original demo)
-vision_model = LlavaNextForConditionalGeneration.from_pretrained(
-    vision_model_path,
-    torch_dtype="auto",
-    device_map="auto"
-)
-
-# --- OR ---
-
-# Option 2: Use torch.float16 but ensure the custom model code is trusted
-# vision_model = LlavaNextForConditionalGeneration.from_pretrained(
-#     vision_model_path,
-#     torch_dtype=torch.float16,
-#     device_map="auto",
-#     trust_remote_code=True
-# )
-
-
-@spaces.GPU
-def generate(
-    message: str,
-    chat_history: list[dict],
-    temperature: float = TEMPERATURE,
-    repetition_penalty: float = REPETITION_PENALTY,
-    top_p: float = TOP_P,
-    top_k: float = TOP_K,
-    max_new_tokens: int = MAX_NEW_TOKENS,
-) -> Iterator[str]:
-    """Generate function for text chat demo."""
-    # Build messages
-    conversation = []
-    conversation.append({"role": "system", "content": SYS_PROMPT})
-    conversation += chat_history
-    conversation.append({"role": "user", "content": message})
-
-    # Convert messages to prompt format
-    input_ids = text_tokenizer.apply_chat_template(
-        conversation,
-        return_tensors="pt",
-        add_generation_prompt=True,
-        truncation=True,
-        max_length=MAX_INPUT_TOKEN_LENGTH - max_new_tokens,
-    )
-
-    input_ids = input_ids.to(text_model.device)
-    streamer = TextIteratorStreamer(text_tokenizer, skip_prompt=True, skip_special_tokens=True)
-    generate_kwargs = dict(
-        {"input_ids": input_ids},
-        streamer=streamer,
-        max_new_tokens=max_new_tokens,
-        do_sample=True,
-        top_p=top_p,
-        top_k=top_k,
-        temperature=temperature,
-        num_beams=1,
-        repetition_penalty=repetition_penalty,
-    )
-
-    t = Thread(target=text_model.generate, kwargs=generate_kwargs)
-    t.start()
-
-    outputs = []
-    for text in streamer:
-        outputs.append(text)
-        yield "".join(outputs)
-
+ 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 themes.research_monochrome import theme
+
+ # Vision imports
+ import random
+ from transformers import LlavaNextProcessor, LlavaNextForConditionalGeneration
+
+ 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"""
+ TITLE = "IBM Granite 3.1 8b Instruct & Vision Preview"
+ DESCRIPTION = """
+ <p>Granite 3.1 8b instruct is an open-source LLM supporting a 128k context window. Start with one of the sample prompts
+ or upload an image and ask a question. 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
+
+ if not torch.cuda.is_available():
+     print("This demo may not work on CPU.")
+
+ # Text Model and Tokenizer
+ text_model = AutoModelForCausalLM.from_pretrained(
+     "ibm-granite/granite-3.1-8b-instruct", torch_dtype=torch.float16, device_map="auto"
+ )
+ text_tokenizer = AutoTokenizer.from_pretrained("ibm-granite/granite-3.1-8b-instruct")
+ text_tokenizer.use_default_system_prompt = False
+
+ # Vision Model and Processor
+ vision_model_path = "ibm-granite/granite-vision-3.1-2b-preview"
+ vision_processor = LlavaNextProcessor.from_pretrained(vision_model_path, use_fast=True)
+ vision_model = LlavaNextForConditionalGeneration.from_pretrained(vision_model_path, torch_dtype="auto", device_map="auto")
+
+
+ @spaces.GPU
+ def generate(
+     message: str,
+     chat_history: list[dict],
+     temperature: float = TEMPERATURE,
+     repetition_penalty: float = REPETITION_PENALTY,
+     top_p: float = TOP_P,
+     top_k: float = TOP_K,
+     max_new_tokens: int = MAX_NEW_TOKENS,
+ ) -> Iterator[str]:
+     """Generate function for text chat demo."""
+     # Build messages
+     conversation = []
+     conversation.append({"role": "system", "content": SYS_PROMPT})
+     conversation += chat_history
+     conversation.append({"role": "user", "content": message})
+
+     # Convert messages to prompt format
+     input_ids = text_tokenizer.apply_chat_template(
+         conversation,
+         return_tensors="pt",
+         add_generation_prompt=True,
+         truncation=True,
+         max_length=MAX_INPUT_TOKEN_LENGTH - max_new_tokens,
+     )
+
+     input_ids = input_ids.to(text_model.device)
+     streamer = TextIteratorStreamer(text_tokenizer, skip_prompt=True, skip_special_tokens=True)
+     generate_kwargs = dict(
+         {"input_ids": input_ids},
+         streamer=streamer,
+         max_new_tokens=max_new_tokens,
+         do_sample=True,
+         top_p=top_p,
+         top_k=top_k,
+         temperature=temperature,
+         num_beams=1,
+         repetition_penalty=repetition_penalty,
+     )
+
+     t = Thread(target=text_model.generate, kwargs=generate_kwargs)
+     t.start()
+
+     outputs = []
+     for text in streamer:
+         outputs.append(text)
+         yield "".join(outputs)
 
 def get_text_from_content(content):
     texts = []
@@ -137,7 +111,7 @@ def get_text_from_content(content):
     return " ".join(texts)
 
 @spaces.GPU
-def chat_inference(image, text, conversation, temperature=VISION_TEMPERATURE, top_p=VISION_TOP_P, top_k=VISION_TOP_K, max_tokens=VISION_MAX_TOKENS):
+def chat_inference(image, text, temperature, top_p, top_k, max_tokens, conversation):
     if conversation is None:
         conversation = []
 
@@ -194,156 +168,126 @@ def conversation_display(conversation):
     return chat_history
 
 def clear_chat():
-    return [], [], "", None, [] # Cleared state for both text and vision
+    return [], [], "", None
 
 css_file_path = Path(Path(__file__).parent / "app.css")
 head_file_path = Path(Path(__file__).parent / "app_head.html")
 
-# Advanced settings (displayed in Accordion) - Text Model
-text_temperature_slider = gr.Slider(
-    minimum=0, maximum=1.0, value=TEMPERATURE, step=0.1, label="Text Temperature", elem_classes=["gr_accordion_element"]
+# Advanced settings (displayed in Accordion) - Common settings for both models
+temperature_slider = gr.Slider(
+    minimum=0, maximum=1.0, value=TEMPERATURE, step=0.1, label="Temperature", elem_classes=["gr_accordion_element"]
 )
-text_top_p_slider = gr.Slider(
-    minimum=0, maximum=1.0, value=TOP_P, step=0.05, label="Text Top P", elem_classes=["gr_accordion_element"]
+top_p_slider = gr.Slider(
+    minimum=0, maximum=1.0, value=TOP_P, step=0.05, label="Top P", elem_classes=["gr_accordion_element"]
 )
-text_top_k_slider = gr.Slider(
-    minimum=0, maximum=100, value=TOP_K, step=1, label="Text Top K", elem_classes=["gr_accordion_element"]
+top_k_slider = gr.Slider(
+    minimum=0, maximum=100, value=TOP_K, step=1, label="Top K", elem_classes=["gr_accordion_element"]
 )
-text_repetition_penalty_slider = gr.Slider(
+
+# Advanced settings specific to Text model
+repetition_penalty_slider = gr.Slider(
     minimum=0,
     maximum=2.0,
     value=REPETITION_PENALTY,
     step=0.05,
-    label="Text Repetition Penalty",
+    label="Repetition Penalty (Text Model)",
     elem_classes=["gr_accordion_element"],
 )
-text_max_new_tokens_slider = gr.Slider(
+max_new_tokens_slider = gr.Slider(
     minimum=1,
     maximum=2000,
     value=MAX_NEW_TOKENS,
     step=1,
-    label="Text Max New Tokens",
+    label="Max New Tokens (Text Model)",
     elem_classes=["gr_accordion_element"],
 )
-text_chat_interface_accordion = gr.Accordion(label="Text Model Advanced Settings", open=False)
 
-# Advanced settings (displayed in Accordion) - Vision Model
-vision_temperature_slider = gr.Slider(
-    minimum=0.0, maximum=2.0, value=VISION_TEMPERATURE, step=0.01, label="Vision Temperature", elem_classes=["gr_accordion_element"]
-)
-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"]
+# Advanced settings specific to Vision model
+max_tokens_slider_vision = gr.Slider(
+    minimum=10,
+    maximum=300,
+    value=128,
+    step=1,
+    label="Max Tokens (Vision Model)",
+    elem_classes=["gr_accordion_element"],
 )
-vision_chat_interface_accordion = gr.Accordion(label="Vision Model Advanced Settings", open=False)
 
+chat_interface_accordion = gr.Accordion(label="Advanced Settings", open=False)
 
 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", elem_id="chatbot", height=500, type='messages')
-    text_input = gr.Textbox(lines=2, placeholder="Enter your message here", label="Message")
-    image_input = gr.Image(type="pil", label="Upload Image (optional)")
-
-    with text_chat_interface_accordion:
-        text_temperature_slider
-        text_top_p_slider
-        text_top_k_slider
-        text_repetition_penalty_slider
-        text_max_new_tokens_slider
-
-    with vision_chat_interface_accordion:
-        vision_temperature_slider
-        vision_top_p_slider
-        vision_top_k_slider
-        vision_max_tokens_slider
-
-
-    clear_button = gr.Button("Clear Chat")
-    send_button = gr.Button("Send Message") # Changed from "Chat" to "Send Message" for clarity
-
-    text_state = gr.State([]) # State for text chatbot history
-    vision_state = gr.State([]) # State for vision chatbot history
-    chatbot_type_state = gr.State("text") # State to track which chatbot is in use
-
-    def send_message(image_input, text_input, chatbot_type_state, text_state, vision_state,
-                     text_temperature, text_repetition_penalty, text_top_p, text_top_k, text_max_new_tokens,
-                     vision_temperature, vision_top_p, vision_top_k, vision_max_tokens):
+    state = gr.State([]) # State for vision chat history
+    chat_history_state = gr.State([]) # State for text chat history
+
+    with gr.Row():
+        with gr.Column(scale=2):
+            image_input = gr.Image(type="pil", label="Upload Image (optional)")
+            with gr.Accordion(label="Vision Model Settings", open=False):
+                max_tokens_input_vision = max_tokens_slider_vision
+            with gr.Accordion(label="Text Model Settings", open=False):
+                repetition_penalty_input = repetition_penalty_slider
+                max_new_tokens_input = max_new_tokens_slider
+            with chat_interface_accordion: # Common Settings
+                temperature_input = temperature_slider
+                top_p_input = top_p_slider
+                top_k_input = top_k_slider
+
+        with gr.Column(scale=3):
+            chatbot = gr.Chatbot(label="Chat History", elem_id="chatbot", type='messages')
+            text_input = gr.Textbox(lines=2, placeholder="Enter your message here", label="Message")
+            with gr.Row():
+                send_button = gr.Button("Chat")
+                clear_button = gr.Button("Clear Chat")
+
+    def process_chat(image_input, text_input, temperature_input, top_p_input, top_k_input, repetition_penalty_input, max_new_tokens_input, max_tokens_input_vision, state, chat_history_state):
         if image_input:
-            chatbot_type_state = "vision"
-            history = vision_state
-            gen_kwargs_vision = {
-                "temperature": vision_temperature,
-                "top_p": vision_top_p,
-                "top_k": vision_top_k,
-                "max_tokens": vision_max_tokens,
-                "conversation": history
-            }
-            chat_output, updated_vision_state = chat_inference(image=image_input, text=text_input, **gen_kwargs_vision)
-            return chat_output, updated_vision_state, chatbot_type_state, gr.ChatInterface.update(visible=False), gr.Chatbot.update(visible=True) # Hide text interface, show vision chatbot
-
+            # Use Vision model
+            return chat_inference(image_input, text_input, temperature_input, top_p_input, top_k_input, max_tokens_input_vision, state)
         else:
-            chatbot_type_state = "text"
-            history = text_state
-            gen_kwargs_text = {
-                "temperature": text_temperature,
-                "repetition_penalty": text_repetition_penalty,
-                "top_p": text_top_p,
-                "top_k": text_top_k,
-                "max_new_tokens": text_max_new_tokens,
-                "message": text_input,
-                "chat_history": history
-            }
+            # Use Text model
+            return generate(text_input, chat_history_state, temperature_input, repetition_penalty_input, top_p_input, top_k_input, max_new_tokens_input), None # Return None for state as text model doesn't use it
 
-            chat_output_iterator = generate(**gen_kwargs_text)
-            output_text = ""
-            for text_chunk in chat_output_iterator:
-                output_text = text_chunk
-
-            updated_text_state = history + [{"role": "user", "content": text_input}, {"role": "assistant", "content": output_text}]
-            text_chatbot_history = updated_text_state # format for chatbot display
-            formatted_history = []
-            for message in text_chatbot_history:
-                formatted_history.append((message["content"] if message["role"] == "user" else None, message["content"] if message["role"] == "assistant" else None))
+    def process_chat_wrapper(image_input_val, text_input_val, temperature_input_val, top_p_input_val, top_k_input_val, repetition_penalty_input_val, max_new_tokens_input_val, max_tokens_input_vision_val, state_val, chat_history_state_val):
+        if image_input_val:
+            chatbot_output, updated_state = process_chat(image_input_val, text_input_val, temperature_input_val, top_p_input_val, top_k_input_val, repetition_penalty_input_val, max_new_tokens_input_val, max_tokens_input_vision_val, state_val, chat_history_state_val)
+            return chatbot_output, updated_state, chat_history_state_val # Return vision state and keep text state unchanged
+        else:
+            chatbot_output_generator, _ = process_chat(image_input_val, text_input_val, temperature_input_val, top_p_input_val, top_k_input_val, repetition_penalty_input_val, max_new_tokens_input_val, max_tokens_input_vision_val, state_val, chat_history_state_val)
+            updated_chat_history = []
+            full_response = ""
+            for response_chunk in chatbot_output_generator:
+                full_response = response_chunk
+            if chat_history_state_val is None:
+                updated_chat_history = []
+            else:
+                updated_chat_history = chat_history_state_val
 
+            updated_chat_history.append({"role": "user", "content": text_input_val})
+            updated_chat_history.append({"role": "assistant", "content": full_response})
 
-            return formatted_history, updated_text_state, chatbot_type_state, gr.ChatInterface.update(visible=True), gr.Chatbot.update(visible=False) # Show text interface, hide vision chatbot
+            return updated_chat_history, state_val, updated_chat_history # Return text chat history, keep vision state unchanged, return updated text history for chatbot display
 
 
     send_button.click(
-        send_message,
-        inputs=[image_input, text_input, chatbot_type_state, text_state, vision_state,
-                text_temperature_slider, text_repetition_penalty_slider, text_top_p_slider, text_top_k_slider, text_max_new_tokens_slider,
-                vision_temperature_slider, vision_top_p_slider, vision_top_k_slider, vision_max_tokens_slider],
-        outputs=[chatbot, vision_state, chatbot_type_state, gr.ChatInterface(), gr.Chatbot()] # Dummy ChatInterface output, real Chatbot output
+        process_chat_wrapper,
+        inputs=[image_input, text_input, temperature_input, top_p_input, top_k_input, repetition_penalty_input, max_new_tokens_input, max_tokens_input_vision, state, chat_history_state],
+        outputs=[chatbot, state, chat_history_state] # Keep both states as output
     )
 
     clear_button.click(
         clear_chat,
         inputs=None,
-        outputs=[chatbot, vision_state, text_input, image_input, text_state] # Added text_state to clear
+        outputs=[chatbot, state, text_input, image_input] # clear_chat clears vision state and input. Need to clear text state also.
     )
 
-
-    image_examples_dir = Path(__file__).parent / "themes" # Create a folder 'themes' in the same directory as app.py and put images there
-    image_examples = [
-        str(image_examples_dir / "cat.jpg"), # Replace "cat.jpg" with actual image file names in 'image_examples'
-        str(image_examples_dir / "dog.jpg"), # Replace "dog.jpg" with actual image file names in 'image_examples'
-        str(image_examples_dir / "horse.jpg") # Replace "horse.jpg" with actual image file names in 'image_examples'
-    ]
-
     gr.Examples(
         examples=[
-            ["Explain the concept of quantum computing to someone with no background in physics or computer science.", None],
-            ["What is OpenShift?", None],
-            ["What's the importance of low latency inference?", None],
-            ["Help me boost productivity habits.", None],
+            ["Explain the concept of quantum computing to someone with no background in physics or computer science."],
+            ["What is OpenShift?"],
+            ["What's the importance of low latency inference?"],
+            ["Help me boost productivity habits."],
             [
                 """Explain the following code in a concise manner:
 
@@ -384,7 +328,7 @@ class Pair {
         this.y = y;
     }
 }
-```""", None
+```"""
             ],
             [
                 """Generate a Java code block from the following explanation:
@@ -395,27 +339,13 @@ The findPairs method takes two arguments: an array of integers and a difference
 
 The Pair class is a simple data structure that stores two integers.
 
-The main method creates an array of integers, initializes the difference value, and calls the findPairs method to find all pairs in the array. Finally, the code iterates over the list of pairs and prints each pair to the console.""" , None # noqa: E501
+The main method creates an array of integers, initializes the difference value, and calls the findPairs method to find all pairs in the array. Finally, the code iterates over the list of pairs and prints each pair to the console."""  # noqa: E501
             ],
-            ["What is in this image?", image_examples[0]], # Vision example using local image
-            ["Describe this image in detail", image_examples[1]], # Vision example using local image
-            ["Identify the object in the image", image_examples[2]], # Vision example using local image
+             ["https://raw.githubusercontent.com/gradio-app/gradio/main/test/test_files/bus.png", "What is this?"] # Vision example
         ],
-        inputs=[text_input, image_input],
-        example_labels=[
-            "Explain quantum computing",
-            "What is OpenShift?",
-            "Importance of low latency inference",
-            "Boosting productivity habits",
-            "Explain and document your code",
-            "Generate Java Code",
-            "Vision Example 1: What is in this image?",
-            "Vision Example 2: Describe this image",
-            "Vision Example 3: Identify object",
-        ],
-        cache_examples=False,
+        inputs=[text_input, text_input, text_input, text_input, text_input, text_input, image_input, image_input] , # Duplicated text_input to match example count, last two are image_input for vision example
+        examples_per_page=7
     )
 
-
 if __name__ == "__main__":
     demo.queue().launch()