Spaces:
Sleeping
Sleeping
| 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") | |
| 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 users location."}},"required":["location","format"]}}}] | |
| # render the tool use prompt as a string: | |
| 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 "None" | |
| def update_template_dropdown(model_name): | |
| template_names = get_template_names(model_name) | |
| if template_names: | |
| return gr.update(choices=template_names, value=None) | |
| def apply_chat_template(model_name, test_conversation, add_generation_prompt, cleanup_whitespace, template_name, hf_token, kwargs): | |
| try: | |
| login(token=hf_token) | |
| tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| except: | |
| return f"model {model_name} could not be loaded or invalid HF token" | |
| try: | |
| outputs = [] | |
| conversation = json.loads(test_conversation) | |
| template = tokenizer.chat_template.get(template_name) if template_name else None | |
| print(kwargs) | |
| 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 str(e) | |
| with gr.Blocks() as demo: | |
| model_name_input = gr.Textbox(label="Model Name", placeholder="Enter model name") | |
| 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") | |
| kwargs_input = gr.JSON(label="Additional kwargs", value=default_tools, render=False) | |
| output = gr.TextArea(label="Formatted conversation") | |
| model_name_input.change(fn=update_template_dropdown, inputs=model_name_input, outputs=template_dropdown) | |
| gr.Interface( | |
| description=description_text, | |
| 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() |