File size: 3,336 Bytes
08fc296
 
 
 
 
 
9bf50e7
 
08fc296
9bf50e7
f5e5ce2
08fc296
 
 
 
 
 
 
 
7a435b7
08fc296
 
9bf50e7
 
 
 
 
 
 
 
 
 
 
 
08fc296
42f053a
08fc296
9bf50e7
 
 
08fc296
9bf50e7
 
 
08fc296
 
9bf50e7
 
 
 
 
42f053a
9bf50e7
08fc296
 
9bf50e7
08fc296
 
9bf50e7
 
7a435b7
 
 
 
 
 
 
 
9bf50e7
7a435b7
 
08fc296
42f053a
 
 
9bf50e7
08fc296
 
 
 
 
 
 
42f053a
08fc296
 
 
 
42f053a
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
86
87
88
89
import gradio as gr
from transformers import AutoTokenizer
import json
import os
from huggingface_hub import login

# Fetch HF Token
HUGGINGFACEHUB_API_TOKEN = os.environ.get("HF_TOKEN", "")

default_model = "meta-llama/Meta-Llama-3-8B-Instruct"

demo_conversation = """[
    {"role": "system", "content": "You are a helpful chatbot."},
    {"role": "user", "content": "Hi there!"},
    {"role": "assistant", "content": "Hello, human!"},
    {"role": "user", "content": "Can I ask a question?"}
]"""

description_text = """# Chat Template Viewer
### This space helps visualize chat formatting using Hugging Face models.
"""

default_tools = [{"type": "function", "function": {"name": "get_current_weather",
                "description": "Get the current weather",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "location": {"type": "string", "description": "The city and state, e.g. San Francisco, CA"},
                        "format": {"type": "string", "enum": ["celsius", "fahrenheit"],
                                   "description": "The temperature unit to use. Infer this from the user's location."}
                    },
                    "required": ["location", "format"]
                }
            }}]

def apply_chat_template(model_name, test_conversation, add_generation_prompt, cleanup_whitespace, hf_token, tools):
    try:
        if hf_token:
            login(token=hf_token)  # Ensure login is successful

        tokenizer = AutoTokenizer.from_pretrained(model_name)
    except Exception as e:
        return f"Error: Could not load model {model_name} or invalid HF token. {str(e)}"

    try:
        conversation = json.loads(test_conversation)

        formatted = tokenizer.apply_chat_template(
            conversation,
            tokenize=False,
            add_generation_prompt=add_generation_prompt,
            tools=tools
        )
        return formatted
    except Exception as e:
        return f"Error: {str(e)}"

with gr.Blocks() as demo:
    gr.Markdown(description_text)

    with gr.Row():
        with gr.Column():
            model_name_input = gr.Textbox(label="Model Name", placeholder="Enter model name", value=default_model)
            hf_token_input = gr.Textbox(label="Hugging Face Token (optional)", placeholder="Enter your HF token", type="password")
            conversation_input = gr.TextArea(value=demo_conversation, lines=8, label="Conversation")
            add_generation_prompt_checkbox = gr.Checkbox(value=False, label="Add generation prompt")
            cleanup_whitespace_checkbox = gr.Checkbox(value=True, label="Cleanup template whitespace")
            format_button = gr.Button("Format Conversation")

        with gr.Column():
            output = gr.TextArea(label="Formatted Conversation", interactive=False, lines=12)

    # Use gr.State() to pass default_tools correctly
    tools_state = gr.State(default_tools)

    format_button.click(
        fn=apply_chat_template,
        inputs=[
            model_name_input,
            conversation_input,
            add_generation_prompt_checkbox,
            cleanup_whitespace_checkbox,
            hf_token_input,
            tools_state  # Wrapped inside gr.State()
        ],
        outputs=output
    )

demo.launch()