File size: 2,748 Bytes
be82a8a
fab8c45
 
 
be82a8a
fab8c45
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
be82a8a
fab8c45
 
 
 
 
 
 
 
 
 
 
 
 
be82a8a
fab8c45
 
 
 
 
 
 
 
be82a8a
fab8c45
be82a8a
 
564c01c
be82a8a
 
 
 
 
 
 
 
 
 
fab8c45
 
 
 
 
 
 
 
be82a8a
 
 
fab8c45
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
import gradio as gr
import re
from typing import List, Tuple
import html

def create_collapsible_html(thinking: str) -> str:
    """Create HTML for a collapsible thinking section"""
    escaped_thinking = html.escape(thinking)
    return f'''
    <div class="thinking-container" style="margin: 8px 0;">
        <details class="thinking-details" style="
            border: 1px solid #ddd;
            border-radius: 4px;
            padding: 8px;
            background-color: #f8f9fa;
            cursor: pointer;
        ">
            <summary style="
                font-weight: 500;
                display: flex;
                align-items: center;
                gap: 8px;
            ">
                <span style="transform: rotate(-90deg); display: inline-block;">➤</span>
                Show thoughts
            </summary>
            <div style="
                margin-top: 8px;
                padding: 8px;
                border-top: 1px solid #eee;
                white-space: pre-wrap;
            ">{escaped_thinking}</div>
        </details>
    </div>
    '''

def process_message(message: str) -> str:
    """Process a message to convert thinking sections into collapsible elements"""
    def replace_thinking(match):
        thinking_content = match.group(1).strip()
        return create_collapsible_html(thinking_content)
    
    # Replace thinking sections with collapsible elements
    processed = re.sub(r'<think>(.*?)</think>', replace_thinking, message, flags=re.DOTALL)
    
    # Remove output tags if present
    processed = re.sub(r'<output>(.*?)</output>', r'\1', processed, flags=re.DOTALL)
    
    return processed

class CustomChatInterface(gr.ChatInterface):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        
    def render_message(self, message: str, is_user: bool) -> str:
        if not is_user:
            return process_message(message)
        return message

demo = CustomChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(value="You must always include <think> ... </think> <output> </output> tokens.", label="System message"),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p (nucleus sampling)",
        ),
    ],
    css="""
    .thinking-details[open] summary span {
        transform: rotate(0deg) !important;
    }
    .thinking-details summary::-webkit-details-marker {
        display: none;
    }
    """
)

if __name__ == "__main__":
    demo.launch()