app.py

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):
    """
    Tokenize the dialogue using the GPT-OSS tokenizer
    """
    if tokenizer is None:
        raise ValueError("Tokenizer not loaded. Please check your installation.")
    
    messages = []
    for message in dialogue_data:
        role = message.get("speaker", "user")
        content = message.get("text", "")
        
        if role == "system":
            messages.append({"role": "system", "content": content})
        elif role == "user":
            messages.append({"role": "user", "content": content})
        elif role == "assistant":
            messages.append({"role": "assistant", "content": content})
    
    formatted_input = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="np"
    )
    
    token_ids = formatted_input[0].tolist()
    decoded_text = []
    colors = ["#FF6B6B", "#4ECDC4", "#45B7D1", "#96CEB4", "#FFEAA7"]
    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)))
    
    print("decoded_text", decoded_text)
    
    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
    """
    return [
        {"speaker": "system", "text": "You are a helpful assistant."},
        {"speaker": "user", "text": "Hello! How are you today?"},
        {"speaker": "assistant", "text": "I'm doing well, thank you for asking! How can I help you today?"},
        {"speaker": "user", "text": "Can you explain what MXFP4 quantization is?"}
    ]

with gr.Blocks(title="GPT-OSS Tokenizer Explorer") as demo:
    gr.Markdown("# GPT-OSS Tokenizer Explorer")
    gr.Markdown("Explore how the GPT-OSS tokenizer processes regular conversations and tool-calling scenarios.")
    
    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():
                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.TabItem("Tool Use Examples"):
            gr.Markdown("See how the GPT-OSS tokenizer handles function calling and tool use conversations.")
            
            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("""
        ### 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
        
        ### 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 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):
        if not dialogue:
            return "Please enter some dialogue first.", {}, "Total Tokens: 0"
        
        result_text, token_count_val = tokenize_dialogue(dialogue)
        
        return result_text, f"Total Tokens: {token_count_val}"
    
    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]
    )
    
    clear_btn.click(
        fn=clear_dialogue,
        outputs=[dialogue_input, highlighted_output, token_count]
    )
    
    dialogue_input.submit(
        fn=process_dialogue,
        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()