Adds Tool Use Examples

app.py

@@ -1,8 +1,123 @@
 import gradio as gr
 from transformers import AutoTokenizer
+import json
+import random
+import math
 
 tokenizer = AutoTokenizer.from_pretrained("openai/gpt-oss-20b")
 
+# Mock tool functions
+def weather_tool(location, days=1):
+    """Get weather forecast for a location"""
+    weather_conditions = ["sunny", "cloudy", "rainy", "snowy", "partly cloudy"]
+    temps = random.randint(15, 35)
+    condition = random.choice(weather_conditions)
+    return {
+        "location": location,
+        "days": days,
+        "forecast": f"{condition}, {temps}°C"
+    }
+
+def calculator_tool(operation, a, b=None):
+    """Perform mathematical calculations"""
+    if operation == "add":
+        return a + b
+    elif operation == "subtract":
+        return a - b
+    elif operation == "multiply":
+        return a * b
+    elif operation == "divide":
+        return a / b if b != 0 else "Error: Division by zero"
+    elif operation == "sqrt":
+        return math.sqrt(a)
+    elif operation == "power":
+        return a ** b
+    else:
+        return "Error: Unknown operation"
+
+def search_tool(query, num_results=3):
+    """Mock web search results"""
+    mock_results = [
+        {"title": f"Result about {query} - Article 1", "url": f"https://example.com/{query.replace(' ', '-')}-1", "snippet": f"This is a comprehensive guide about {query} with detailed information..."},
+        {"title": f"{query} - Wikipedia", "url": f"https://en.wikipedia.org/wiki/{query.replace(' ', '_')}", "snippet": f"{query} is an important topic that covers various aspects..."},
+        {"title": f"Latest news on {query}", "url": f"https://news.example.com/{query.replace(' ', '-')}", "snippet": f"Recent developments and updates related to {query}..."},
+    ]
+    return mock_results[:num_results]
+
+def code_executor_tool(code):
+    """Execute simple Python code (safe expressions only)"""
+    try:
+        # Only allow simple mathematical expressions for safety
+        allowed_names = {"__builtins__": {"abs": abs, "max": max, "min": min, "sum": sum, "len": len}}
+        result = eval(code, {"__builtins__": {}}, allowed_names)
+        return f"Result: {result}"
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+# Tool definitions for function calling
+AVAILABLE_TOOLS = [
+    {
+        "type": "function",
+        "function": {
+            "name": "get_weather",
+            "description": "Get weather forecast for a specific location",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "location": {"type": "string", "description": "The city and state/country"},
+                    "days": {"type": "integer", "description": "Number of days for forecast (1-7)", "default": 1}
+                },
+                "required": ["location"]
+            }
+        }
+    },
+    {
+        "type": "function", 
+        "function": {
+            "name": "calculate",
+            "description": "Perform mathematical calculations",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "operation": {"type": "string", "enum": ["add", "subtract", "multiply", "divide", "sqrt", "power"]},
+                    "a": {"type": "number", "description": "First number"},
+                    "b": {"type": "number", "description": "Second number (not needed for sqrt)"}
+                },
+                "required": ["operation", "a"]
+            }
+        }
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "web_search", 
+            "description": "Search the web for information",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "query": {"type": "string", "description": "Search query"},
+                    "num_results": {"type": "integer", "description": "Number of results (1-10)", "default": 3}
+                },
+                "required": ["query"]
+            }
+        }
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "execute_code",
+            "description": "Execute simple Python code expressions",
+            "parameters": {
+                "type": "object", 
+                "properties": {
+                    "code": {"type": "string", "description": "Python code expression to execute"}
+                },
+                "required": ["code"]
+            }
+        }
+    }
+]
+
 
 def tokenize_dialogue(dialogue_data):
     """
@@ -44,6 +159,161 @@ def tokenize_dialogue(dialogue_data):
     
     return gr.HighlightedText(value=decoded_text, color_map=color_map), len(token_ids)
 
+def tokenize_tool_conversation(messages_with_tools):
+    """
+    Tokenize a conversation that includes tool calls and responses
+    """
+    if tokenizer is None:
+        raise ValueError("Tokenizer not loaded. Please check your installation.")
+    
+    # Preprocess messages to handle None content
+    processed_messages = []
+    for message in messages_with_tools:
+        processed_message = message.copy()
+        if processed_message.get("content") is None:
+            processed_message["content"] = ""
+        processed_messages.append(processed_message)
+    
+    formatted_input = tokenizer.apply_chat_template(
+        processed_messages,
+        add_generation_prompt=False,
+        return_tensors="np"
+    )
+    
+    token_ids = formatted_input[0].tolist()
+    decoded_text = []
+    colors = ["#FF6B6B", "#4ECDC4", "#45B7D1", "#96CEB4", "#FFEAA7", "#DDA0DD", "#98FB98", "#F0E68C"]
+    color_map = {}
+    
+    for i, token_id in enumerate(token_ids):
+        color = colors[i % len(colors)]
+        if token_id not in color_map:
+            color_map[str(token_id)] = color
+        decoded_text.append((tokenizer.decode([token_id]), str(token_id)))
+    
+    return gr.HighlightedText(value=decoded_text, color_map=color_map), len(token_ids)
+
+def execute_tool_call(tool_name, arguments):
+    """Execute a tool call and return the result"""
+    try:
+        if tool_name == "get_weather":
+            return weather_tool(**arguments)
+        elif tool_name == "calculate":
+            return calculator_tool(**arguments)
+        elif tool_name == "web_search":
+            return search_tool(**arguments)
+        elif tool_name == "execute_code":
+            return code_executor_tool(**arguments)
+        else:
+            return {"error": f"Unknown tool: {tool_name}"}
+    except Exception as e:
+        return {"error": str(e)}
+
+def create_tool_conversation_examples():
+    """Create example conversations with tool use"""
+    examples = {
+        "Weather Query": [
+            {"role": "system", "content": "You are a helpful assistant with access to weather information."},
+            {"role": "user", "content": "What's the weather like in Tokyo today?"},
+            {
+                "role": "assistant", 
+                "content": "",
+                "tool_calls": [
+                    {
+                        "id": "call_1",
+                        "type": "function",
+                        "function": {
+                            "name": "get_weather",
+                            "arguments": json.dumps({"location": "Tokyo, Japan", "days": 1})
+                        }
+                    }
+                ]
+            },
+            {
+                "role": "tool",
+                "content": json.dumps(weather_tool("Tokyo, Japan", 1)),
+                "tool_call_id": "call_1"
+            },
+            {"role": "assistant", "content": "The weather in Tokyo today is sunny with a temperature of 25°C. It looks like a great day to be outside!"}
+        ],
+        
+        "Math Calculation": [
+            {"role": "system", "content": "You are a helpful assistant that can perform calculations."},
+            {"role": "user", "content": "What's 15% tip on a $87.50 bill?"},
+            {
+                "role": "assistant",
+                "content": "",
+                "tool_calls": [
+                    {
+                        "id": "call_2",
+                        "type": "function", 
+                        "function": {
+                            "name": "calculate",
+                            "arguments": json.dumps({"operation": "multiply", "a": 87.50, "b": 0.15})
+                        }
+                    }
+                ]
+            },
+            {
+                "role": "tool",
+                "content": json.dumps({"result": calculator_tool("multiply", 87.50, 0.15)}),
+                "tool_call_id": "call_2"
+            },
+            {"role": "assistant", "content": "A 15% tip on an $87.50 bill would be $13.13. So your total would be $100.63."}
+        ],
+
+        "Web Search": [
+            {"role": "system", "content": "You are a helpful assistant that can search for information."},
+            {"role": "user", "content": "Find me information about machine learning trends in 2024"},
+            {
+                "role": "assistant",
+                "content": "",
+                "tool_calls": [
+                    {
+                        "id": "call_3",
+                        "type": "function",
+                        "function": {
+                            "name": "web_search", 
+                            "arguments": json.dumps({"query": "machine learning trends 2024", "num_results": 3})
+                        }
+                    }
+                ]
+            },
+            {
+                "role": "tool",
+                "content": json.dumps(search_tool("machine learning trends 2024", 3)),
+                "tool_call_id": "call_3"
+            },
+            {"role": "assistant", "content": "I found several resources about machine learning trends in 2024. Based on the search results, key trends include advances in large language models, improved efficiency in AI training, and greater focus on responsible AI development."}
+        ],
+
+        "Code Execution": [
+            {"role": "system", "content": "You are a helpful assistant that can execute Python code."},
+            {"role": "user", "content": "Calculate the sum of numbers from 1 to 100"},
+            {
+                "role": "assistant",
+                "content": "",
+                "tool_calls": [
+                    {
+                        "id": "call_4", 
+                        "type": "function",
+                        "function": {
+                            "name": "execute_code",
+                            "arguments": json.dumps({"code": "sum(range(1, 101))"})
+                        }
+                    }
+                ]
+            },
+            {
+                "role": "tool",
+                "content": json.dumps({"result": code_executor_tool("sum(range(1, 101))")}),
+                "tool_call_id": "call_4"
+            },
+            {"role": "assistant", "content": "The sum of numbers from 1 to 100 is 5,050."}
+        ]
+    }
+    return examples
+
 def create_sample_dialogue():
     """
     Create a sample dialogue for demonstration
@@ -57,57 +327,102 @@ def create_sample_dialogue():
 
 with gr.Blocks(title="GPT-OSS Tokenizer Explorer") as demo:
     gr.Markdown("# GPT-OSS Tokenizer Explorer")
-    gr.Markdown("Enter a dialogue and see how the GPT-OSS tokenizer processes it. Use the format `speaker: message` in the dialogue component.")
+    gr.Markdown("Explore how the GPT-OSS tokenizer processes regular conversations and tool-calling scenarios.")
     
-    with gr.Row():
-        with gr.Column(scale=1):
-            gr.Markdown("### Input Dialogue")
-            
-            dialogue_input = gr.Dialogue(
-                speakers=["system", "user", "assistant"],
-                label="Enter your dialogue",
-                placeholder="Type 'system:', 'user:', or 'assistant:' followed by your message",
-                show_submit_button=True,
-                show_copy_button=True,
-                type="dialogue",
-                ui_mode="dialogue-only",
-            )
+    with gr.Tabs():
+        with gr.TabItem("Regular Dialogue"):
+            gr.Markdown("Enter a dialogue and see how the GPT-OSS tokenizer processes it.")
             
             with gr.Row():
-                sample_btn = gr.Button("Load Sample", variant="secondary")
-                clear_btn = gr.Button("Clear", variant="secondary")
+                with gr.Column(scale=1):
+                    gr.Markdown("### Input Dialogue")
+                    
+                    dialogue_input = gr.Dialogue(
+                        speakers=["system", "user", "assistant"],
+                        label="Enter your dialogue",
+                        placeholder="Type 'system:', 'user:', or 'assistant:' followed by your message",
+                        show_submit_button=True,
+                        show_copy_button=True,
+                        type="dialogue",
+                        ui_mode="dialogue-only",
+                    )
+                    
+                    with gr.Row():
+                        sample_btn = gr.Button("Load Sample", variant="secondary")
+                        clear_btn = gr.Button("Clear", variant="secondary")
+                
+                with gr.Column(scale=1):
+                    gr.Markdown("### Tokenization Results")
+                    
+                    highlighted_output = gr.HighlightedText(
+                        label="Tokenized Output",
+                        show_inline_category=False
+                    )
+                    
+                    token_count = gr.Label(
+                        value="Total Tokens: 0",
+                        label="Token Count"
+                    )
         
-        with gr.Column(scale=1):
-            gr.Markdown("### Tokenization Results")
-            
-            highlighted_output = gr.HighlightedText(
-                label="Tokenized Output",
-                show_inline_category=False
-            )
+        with gr.TabItem("Tool Use Examples"):
+            gr.Markdown("See how the GPT-OSS tokenizer handles function calling and tool use conversations.")
             
-            token_count = gr.Label(
-                value="Total Tokens: 0",
-                label="Token Count"
-            )
+            with gr.Row():
+                with gr.Column(scale=1):
+                    gr.Markdown("### Tool Use Scenarios")
+                    
+                    example_dropdown = gr.Dropdown(
+                        choices=["Weather Query", "Math Calculation", "Web Search", "Code Execution"],
+                        value="Weather Query",
+                        label="Select Example Scenario"
+                    )
+                    
+                    load_example_btn = gr.Button("Load Example", variant="primary")
+                    
+                    gr.Markdown("### Available Tools")
+                    tools_display = gr.JSON(
+                        value=AVAILABLE_TOOLS,
+                        label="Tool Definitions"
+                    )
+                
+                with gr.Column(scale=1):
+                    gr.Markdown("### Tool Conversation Tokenization")
+                    
+                    tool_highlighted_output = gr.HighlightedText(
+                        label="Tokenized Tool Conversation",
+                        show_inline_category=False
+                    )
+                    
+                    tool_token_count = gr.Label(
+                        value="Total Tokens: 0",
+                        label="Token Count"
+                    )
+                    
+                    gr.Markdown("### Conversation Preview")
+                    conversation_display = gr.JSON(
+                        label="Conversation Structure",
+                        value=[]
+                    )
     
     with gr.Accordion("How to use", open=False):
         gr.Markdown("""
-        ### Instructions:
+        ### Regular Dialogue Tab:
         1. **Enter dialogue**: Use the dialogue component to enter conversations
         2. **Speaker format**: Type `system:`, `user:`, or `assistant:` followed by your message
         3. **Submit**: Click 'Tokenize Dialogue' to process the conversation
         4. **View results**: See the tokenization details in the output area
         
-        ### Example:
-        ```
-        system: You are a helpful assistant.
-        user: Hello! How are you today?
-        assistant: I'm doing well, thank you for asking!
-        ```
+        ### Tool Use Examples Tab:
+        1. **Select scenario**: Choose from weather query, math calculation, web search, or code execution
+        2. **Load example**: Click 'Load Example' to see a tool-calling conversation
+        3. **Compare tokenization**: See how tool calls differ from regular messages
+        4. **Explore tools**: View available tool definitions and their parameters
         
         ### What you'll see:
         - **Total tokens**: Number of tokens in the conversation
-        - **Tokenized output**: How the tokenizer formats the conversation
+        - **Tokenized output**: How the tokenizer formats conversations and tool calls
+        - **Tool definitions**: JSON schema for available functions
+        - **Conversation structure**: The complete message flow including tool calls and responses
         """)
     
     def process_dialogue(dialogue):
@@ -121,6 +436,20 @@ with gr.Blocks(title="GPT-OSS Tokenizer Explorer") as demo:
     def clear_dialogue():
         return None, [], "Total Tokens: 0"
     
+    def load_tool_example(example_name):
+        """Load a tool use example and tokenize it"""
+        examples = create_tool_conversation_examples()
+        if example_name not in examples:
+            return gr.HighlightedText(value=[]), "Total Tokens: 0", []
+        
+        conversation = examples[example_name]
+        try:
+            result_text, token_count_val = tokenize_tool_conversation(conversation)
+            return result_text, f"Total Tokens: {token_count_val}", conversation
+        except Exception as e:
+            error_msg = f"Error tokenizing conversation: {str(e)}"
+            return gr.HighlightedText(value=[(error_msg, "error")]), "Total Tokens: 0", conversation
+    
     sample_btn.click(
         fn=create_sample_dialogue,
         outputs=[dialogue_input]
@@ -136,6 +465,13 @@ with gr.Blocks(title="GPT-OSS Tokenizer Explorer") as demo:
         inputs=[dialogue_input],
         outputs=[highlighted_output, token_count]
     )
+    
+    # Tool use event handlers
+    load_example_btn.click(
+        fn=load_tool_example,
+        inputs=[example_dropdown],
+        outputs=[tool_highlighted_output, tool_token_count, conversation_display]
+    )
 
 if __name__ == "__main__":
     demo.launch()