Update src/app.py

src/app.py

@@ -1,106 +1,122 @@
 """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
-
- 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)
+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")  # 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. 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
+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"
+)
+tokenizer = AutoTokenizer.from_pretrained("ibm-granite/granite-3.1-8b-instruct")
+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)
+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.
+)
 
+# =============================================================================
+# Text Generation Function (for text-only chat)
+# =============================================================================
+@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."""
+    conversation = []
+    conversation.append({"role": "system", "content": SYS_PROMPT})
+    conversation.extend(chat_history)
+    conversation.append({"role": "user", "content": message})
+    input_ids = 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(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)
+
+# =============================================================================
+# Vision Chat Inference Function (for image+text chat)
+# =============================================================================
 def get_text_from_content(content):
     texts = []
     for item in content:
@@ -111,23 +127,17 @@ def get_text_from_content(content):
     return " ".join(texts)
 
 @spaces.GPU
-def chat_inference(image, text, temperature, top_p, top_k, max_tokens, conversation):
+def chat_inference(image, text, conversation, temperature=VISION_TEMPERATURE, top_p=VISION_TOP_P, top_k=VISION_TOP_K, max_tokens=VISION_MAX_TOKENS):
     if conversation is None:
         conversation = []
-
     user_content = []
     if image is not None:
         user_content.append({"type": "image", "image": image})
     if text and text.strip():
         user_content.append({"type": "text", "text": text.strip()})
     if not user_content:
-        return conversation_display(conversation), conversation
-
-    conversation.append({
-        "role": "user",
-        "content": user_content
-    })
-
+        return display_vision_conversation(conversation), conversation
+    conversation.append({"role": "user", "content": user_content})
     inputs = vision_processor.apply_chat_template(
         conversation,
         add_generation_prompt=True,
@@ -135,9 +145,7 @@ def chat_inference(image, text, temperature, top_p, top_k, max_tokens, conversat
         return_dict=True,
         return_tensors="pt"
     ).to("cuda")
-
     torch.manual_seed(random.randint(0, 10000))
-
     generation_kwargs = {
         "max_new_tokens": max_tokens,
         "temperature": temperature,
@@ -145,207 +153,165 @@ def chat_inference(image, text, temperature, top_p, top_k, max_tokens, conversat
         "top_k": top_k,
         "do_sample": True,
     }
-
     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()}]})
+    return display_vision_conversation(conversation), conversation
+
+# =============================================================================
+# Helper Functions to Format Conversation for Display
+# =============================================================================
+def display_text_conversation(conversation):
+    """Convert a text conversation (list of dicts) into a list of (user, assistant) tuples."""
+    chat_history = []
+    i = 0
+    while i < len(conversation):
+        if conversation[i]["role"] == "user":
+            user_msg = conversation[i]["content"]
+            assistant_msg = ""
+            if i + 1 < len(conversation) and conversation[i+1]["role"] == "assistant":
+                assistant_msg = conversation[i+1]["content"]
+                i += 2
+            else:
+                i += 1
+            chat_history.append((user_msg, assistant_msg))
+        else:
+            i += 1
+    return chat_history
 
-    conversation.append({
-        "role": "assistant",
-        "content": [{"type": "text", "text": assistant_response.strip()}]
-    })
-
-    return conversation_display(conversation), conversation
-
-def conversation_display(conversation):
+def display_vision_conversation(conversation):
+    """Convert a vision conversation (with mixed content types) into a list of (user, assistant) tuples."""
     chat_history = []
-    for msg in conversation:
-        if msg["role"] == "user":
-            user_text = get_text_from_content(msg["content"])
-            chat_history.append({"role": "user", "content": user_text})
-        elif msg["role"] == "assistant":
-            assistant_text = msg["content"][0]["text"].split("<|assistant|>")[-1].strip()
-            chat_history.append({"role": "assistant", "content": assistant_text})
+    i = 0
+    while i < len(conversation):
+        if conversation[i]["role"] == "user":
+            user_msg = get_text_from_content(conversation[i]["content"])
+            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()
+                i += 2
+            else:
+                i += 1
+            chat_history.append((user_msg, assistant_msg))
+        else:
+            i += 1
     return chat_history
 
+# =============================================================================
+# Unified Send-Message Function
+# =============================================================================
+def send_message(image, text,
+                 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,
+                 text_state, vision_state):
+    """
+    If an image is uploaded, use the vision model; otherwise, use the text model.
+    Returns updated conversation (as a list of tuples) and state for each branch.
+    """
+    if image is not None:
+        # Vision branch
+        conv = vision_state if vision_state is not None else []
+        chat_history, updated_conv = chat_inference(
+            image, text, conv,
+            temperature=vision_temperature,
+            top_p=vision_top_p,
+            top_k=vision_top_k,
+            max_tokens=vision_max_tokens
+        )
+        vision_state = updated_conv
+        # In vision mode, the conversation display is produced from the vision branch.
+        return chat_history, text_state, vision_state
+    else:
+        # Text branch
+        conv = text_state if text_state is not None else []
+        output_text = ""
+        for chunk in generate(
+            text, conv,
+            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
+        ):
+            output_text = chunk
+        conv.append({"role": "user", "content": text})
+        conv.append({"role": "assistant", "content": output_text})
+        text_state = conv
+        chat_history = display_text_conversation(text_state)
+        return chat_history, text_state, vision_state
+
 def clear_chat():
-    return [], [], "", None
+    # 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")
 
-# 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"]
-)
-top_p_slider = gr.Slider(
-    minimum=0, maximum=1.0, value=TOP_P, step=0.05, label="Top P", 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"]
-)
-
-# Advanced settings specific to Text model
-repetition_penalty_slider = gr.Slider(
-    minimum=0,
-    maximum=2.0,
-    value=REPETITION_PENALTY,
-    step=0.05,
-    label="Repetition Penalty (Text Model)",
-    elem_classes=["gr_accordion_element"],
-)
-max_new_tokens_slider = gr.Slider(
-    minimum=1,
-    maximum=2000,
-    value=MAX_NEW_TOKENS,
-    step=1,
-    label="Max New Tokens (Text Model)",
-    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"],
-)
-
-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)
-
-    state = gr.State([]) # State for vision chat history
-    chat_history_state = gr.State([]) # State for text chat history
-
+    
+    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)")
-            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:
-            # Use Vision model
-            return chat_inference(image_input, text_input, temperature_input, top_p_input, top_k_input, max_tokens_input_vision, state)
-        else:
-            # 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
-
-    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 updated_chat_history, state_val, updated_chat_history # Return text chat history, keep vision state unchanged, return updated text history for chatbot display
-
-
+        with gr.Column(scale=1):
+            with gr.Accordion("Text Advanced Settings", open=False):
+                text_temperature_slider = gr.Slider(minimum=0, maximum=1.0, value=TEMPERATURE, step=0.1, label="Temperature", elem_classes=["gr_accordion_element"])
+                repetition_penalty_slider = gr.Slider(minimum=0, maximum=2.0, value=REPETITION_PENALTY, step=0.05, label="Repetition Penalty", 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"])
+                top_k_slider = gr.Slider(minimum=0, maximum=100, value=TOP_K, step=1, label="Top K", elem_classes=["gr_accordion_element"])
+                max_new_tokens_slider = gr.Slider(minimum=1, maximum=2000, value=MAX_NEW_TOKENS, step=1, label="Max New Tokens", elem_classes=["gr_accordion_element"])
+            with gr.Accordion("Vision Advanced Settings", open=False):
+                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"])
+    
+    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(
-        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
+        send_message,
+        inputs=[
+            image_input, text_input,
+            text_temperature_slider, repetition_penalty_slider, top_p_slider, top_k_slider, max_new_tokens_slider,
+            vision_temperature_slider, vision_top_p_slider, vision_top_k_slider, vision_max_tokens_slider,
+            text_state, vision_state
+        ],
+        outputs=[chatbot, text_state, vision_state]
     )
-
+    
     clear_button.click(
         clear_chat,
         inputs=None,
-        outputs=[chatbot, state, text_input, image_input] # clear_chat clears vision state and input. Need to clear text state also.
+        outputs=[chatbot, text_state, vision_state, text_input, image_input]
     )
-
+    
     gr.Examples(
         examples=[
-            ["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:
-
-```java
-import java.util.ArrayList;
-import java.util.List;
-
-public class Main {
-
-    public static void main(String[] args) {
-        int[] arr = {1, 5, 3, 4, 2};
-        int diff = 3;
-        List<Pair> pairs = findPairs(arr, diff);
-        for (Pair pair : pairs) {
-            System.out.println(pair.x + " " + pair.y);
-        }
-    }
-
-    public static List<Pair> findPairs(int[] arr, int diff) {
-        List<Pair> pairs = new ArrayList<>();
-        for (int i = 0; i < arr.length; i++) {
-            for (int j = i + 1; j < arr.length; j++) {
-                if (Math.abs(arr[i] - arr[j]) < diff) {
-                    pairs.add(new Pair(arr[i], arr[j]));
-                }
-            }
-        }
-
-        return pairs;
-    }
-}
-
-class Pair {
-    int x;
-    int y;
-    public Pair(int x, int y) {
-        this.x = x;
-        this.y = y;
-    }
-}
-```"""
-            ],
-            [
-                """Generate a Java code block from the following explanation:
-
-The code in the Main class finds all pairs in an array whose absolute difference is less than a given value.
-
-The findPairs method takes two arguments: an array of integers and a difference value. It iterates over the array and compares each element to every other element in the array. If the absolute difference between the two elements is less than the difference value, a new Pair object is created and added to a list.
-
-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."""  # noqa: E501
-            ],
-             ["https://raw.githubusercontent.com/gradio-app/gradio/main/test/test_files/bus.png", "What is this?"] # Vision example
+            ["https://raw.githubusercontent.com/gradio-app/gradio/main/test/test_files/bus.png", "What is in this image?"],
+            ["Explain quantum computing to a beginner.", None],
+            ["What is OpenShift?", None]
+        ],
+        inputs=[image_input, text_input],
+        example_labels=[
+            "Vision Example: What is in this image?",
+            "Explain quantum computing",
+            "What is OpenShift?"
         ],
-        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
+        cache_examples=False,
     )
 
 if __name__ == "__main__":
-    demo.queue().launch()
+    demo.queue().launch()