|
import gradio as gr |
|
from transformers import AutoTokenizer |
|
import json |
|
import os |
|
from huggingface_hub import login |
|
|
|
|
|
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 is a helper to learn more about [Chat Templates](https://huggingface.co/docs/transformers/main/en/chat_templating). |
|
""" |
|
|
|
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 get_template_names(model_name): |
|
try: |
|
tokenizer = AutoTokenizer.from_pretrained(model_name) |
|
if isinstance(tokenizer.chat_template, dict): |
|
return list(tokenizer.chat_template.keys()) |
|
else: |
|
return [] |
|
except Exception as e: |
|
return ["Default"] |
|
|
|
def update_template_dropdown(model_name): |
|
template_names = get_template_names(model_name) |
|
if template_names: |
|
return gr.Dropdown.update(choices=template_names, value=template_names[0]) |
|
return gr.Dropdown.update(choices=[], value=None) |
|
|
|
def apply_chat_template(model_name, test_conversation, add_generation_prompt, cleanup_whitespace, template_name, hf_token, kwargs): |
|
try: |
|
if hf_token: |
|
login(token=hf_token) |
|
|
|
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) |
|
|
|
if template_name and tokenizer.chat_template: |
|
template = tokenizer.chat_template.get(template_name, None) |
|
else: |
|
template = None |
|
|
|
formatted = tokenizer.apply_chat_template( |
|
conversation, |
|
chat_template=template, |
|
tokenize=False, |
|
add_generation_prompt=add_generation_prompt, |
|
tools=default_tools |
|
) |
|
return formatted |
|
except Exception as e: |
|
return f"Error: {str(e)}" |
|
|
|
with gr.Blocks() as demo: |
|
gr.Markdown(description_text) |
|
|
|
model_name_input = gr.Textbox(label="Model Name", placeholder="Enter model name", value=default_model) |
|
template_dropdown = gr.Dropdown(label="Template Name", choices=[], interactive=True) |
|
conversation_input = gr.TextArea(value=demo_conversation, lines=6, 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") |
|
hf_token_input = gr.Textbox(label="Hugging Face Token (optional)", placeholder="Enter your HF token", type="password") |
|
kwargs_input = gr.JSON(label="Additional kwargs", value=default_tools, visible=False) |
|
output = gr.TextArea(label="Formatted conversation", interactive=False) |
|
|
|
update_button = gr.Button("Update Template List") |
|
format_button = gr.Button("Format Conversation") |
|
|
|
update_button.click(fn=update_template_dropdown, inputs=model_name_input, outputs=template_dropdown) |
|
|
|
format_button.click( |
|
fn=apply_chat_template, |
|
inputs=[ |
|
model_name_input, |
|
conversation_input, |
|
add_generation_prompt_checkbox, |
|
cleanup_whitespace_checkbox, |
|
template_dropdown, |
|
hf_token_input, |
|
kwargs_input |
|
], |
|
outputs=output |
|
) |
|
|
|
demo.launch() |
|
|
