EmbeddedLLM/chat-template-generation
12
1import streamlit as st2from transformers import AutoTokenizer3import json4import tempfile5import os6import uuid7import copy8import shutil9 10st.set_page_config(layout="wide")11 12def sanitize_jinja2(jinja_lines):13 14 one_liner_jinja = ""15 for line in jinja_lines:16 one_liner_jinja += line.lstrip(" ").rstrip("\n")17 18 return one_liner_jinja19 20@st.cache_resource21def get_existing_templates():22 return [None] + os.listdir("./templates")23 24# if os.path.exists("./tmp"):25# if len(os.listdir("./tmp")) > 20:26# shutil.rmtree('./tmp')27 28# Initialization29if 'tokenizer_json' not in st.session_state:30 st.session_state['tokenizer_json'] = None31 32if 'tokenizer' not in st.session_state:33 st.session_state['tokenizer'] = None34 35if 'repo_normalized_name' not in st.session_state:36 st.session_state['repo_normalized_name'] = None37 38if 'repo_id' not in st.session_state:39 st.session_state['repo_id'] = None40 41if 'input_jinja_template' not in st.session_state:42 st.session_state['input_jinja_template'] = ""43 44if 'uuid' not in st.session_state:45 st.session_state['uuid'] = uuid.uuid4()46 os.makedirs(f"./tmp/{st.session_state['uuid']}")47 48if 'successful_template' not in st.session_state:49 st.session_state['successful_template'] = ''50 51if 'generated_prompt_w_add_generation_prompt' not in st.session_state:52 st.session_state['generated_prompt_w_add_generation_prompt'] = ''53 54if 'generated_prompt_wo_add_generation_prompt' not in st.session_state:55 st.session_state['generated_prompt_wo_add_generation_prompt'] = ''56 57if not os.path.exists("./tmp"):58 os.makedirs("./tmp")59 60title_description = """61Chat Template Generation: Make Chat Easier with Huggingface Tokenizer62"""63 64st.title(title_description)65st.markdown('This streamlit app is to serve as an easier way to check and push the chat template to your/exisiting huggingface repo')66 67list_of_templates = get_existing_templates()68with st.expander("Current predefined templates"):69 for model in list_of_templates[1:]:70 st.markdown(f"- {model}")71 st.info('More templates will be predefined for easier setup of chat template.', icon="ℹ️")72 73st.divider()74# custom_repo_tab, prebuilt_template_tab = st.tabs(["Specify Custom Repository Path", "Select Prebuilt Template"])75 76hf_model_repo_name = st.text_input("Hugging Face Model Repository To Update", value="tiiuae/falcon-7b", max_chars=None, key=None, type="default", 77 help=None, autocomplete=None, label_visibility="visible")78 79gen_button = st.button("Get Tokenizer Config")80 81if gen_button:82 with st.spinner(text="In progress...", cache=False):83 st.session_state['repo_id'] = hf_model_repo_name84 st.session_state['tokenizer'] = AutoTokenizer.from_pretrained(hf_model_repo_name)85 86 st.session_state['repo_normalized_name'] = hf_model_repo_name.replace("/", "_")87 st.session_state['tokenizer_json'] = f"./tmp/{st.session_state['uuid']}_{hf_model_repo_name}"88 # st.session_state['tokenizer'].save_pretrained(st.session_state['tokenizer_json'])89 90if st.session_state['tokenizer_json'] is not None:91 st.session_state['tokenizer'].save_pretrained(st.session_state['tokenizer_json'])92 with open(f"{st.session_state['tokenizer_json']}/tokenizer_config.json", "rb") as f:93 tokenizer_json = json.load(f)94 shutil.rmtree(st.session_state['tokenizer_json'])95 96 json_spec, col2 = st.columns(spec=[0.3, 0.7])97 98 99 with json_spec:100 st.markdown(f"### Tokenizer Config from {st.session_state['repo_normalized_name']}")101 st.json(json.dumps(tokenizer_json, indent=4))102 103 with col2:104 chat = [105 {"role": "system", "content": "You are a helpful assistant."},106 {"role": "user", "content": "Hello, how are you?"},107 {"role": "assistant", "content": "I'm doing great. How can I help you today?"},108 {"role": "user", "content": "I'd like to show off how chat templating works!"},109 ]110 st.markdown("### Example Conversation")111 st.json(json.dumps(chat, indent=4), expanded=False)112 113 prompt_template_col, prompt_template_output_col = st.columns(spec=[0.3, 0.7])114 115 with prompt_template_col:116 list_of_templates = get_existing_templates()117 selected_template = st.selectbox("Choose Existing Template or Leave Blank. (If template is None, it will check current tokenizer's `chat_template` and `default_chat_template` fields)", 118 options=list_of_templates, 119 index=0, placeholder="Choose a template (If template is None, it will check current tokenizer `chat_template` and `default_chat_template` fields)", disabled=False, label_visibility="visible")120 # add_generation_prompt_checkbox = st.checkbox("add_generation_prompt")121 generate_prompt_example_button = st.button("Generate Prompt", key="generate_prompt_example_button")122 123 # if selected_template is None:124 # st.session_state['input_jinja_template'] = st.text_area(125 # "Jinja Chat Template", value=st.session_state['input_jinja_template'], 126 # height=500, placeholder=None, disabled=False, label_visibility="visible")127 128 if selected_template is not None:129 with open(f"./templates/{selected_template}", "r") as f:130 jinja_lines = f.readlines()131 st.session_state['input_jinja_template'] = "".join(jinja_lines)132 133 if selected_template is None:134 st.session_state['input_jinja_template'] = st.session_state['tokenizer'].chat_template 135 if st.session_state['input_jinja_template'] is None:136 st.session_state['input_jinja_template'] = st.session_state['tokenizer'].default_chat_template137 138 139 st.session_state['input_jinja_template'] = st.text_area(140 "Jinja Chat Template", value=st.session_state['input_jinja_template'], 141 height=500, placeholder=None, disabled=False, label_visibility="visible")142 143 144 with prompt_template_output_col:145 # print(st.session_state['input_jinja_template'])146 if generate_prompt_example_button:147 with open(f"./tmp/{st.session_state['uuid']}/tmp_chat_template.json", "w") as fp:148 fp.write(st.session_state['input_jinja_template'])149 with open(f"./tmp/{st.session_state['uuid']}/tmp_chat_template.json", "r") as f:150 jinja_lines = copy.deepcopy(f.readlines())151 st.session_state['tokenizer'].chat_template = sanitize_jinja2(jinja_lines)152 # print(sanitize_jinja2(jinja_lines))153 os.remove(f"./tmp/{st.session_state['uuid']}/tmp_chat_template.json")154 st.session_state['generated_prompt_wo_add_generation_prompt'] = st.session_state['tokenizer'].apply_chat_template(chat, tokenize=False, add_generation_prompt= False)155 st.session_state['generated_prompt_w_add_generation_prompt'] = st.session_state['tokenizer'].apply_chat_template(chat, tokenize=False, add_generation_prompt= True)156 # print(generated_prompt_wo_add_generation_prompt)157 st.session_state['successful_template'] = copy.deepcopy(st.session_state['input_jinja_template'])158 # print(st.session_state['successful_template'])159 160 if len(st.session_state['successful_template']) > 0:161 st.text_area(162 "Generate Prompt with `add_generation_prompt=False`", value=st.session_state['generated_prompt_wo_add_generation_prompt'], 163 height=300, placeholder=None, disabled=True, label_visibility="visible", key="generated_prompt_wo_add_generation_prompt_text_area")164 165 st.text_area(166 "Generate Prompt with `add_generation_prompt=True`", value=st.session_state['generated_prompt_w_add_generation_prompt'], 167 height=300, placeholder=None, disabled=True, label_visibility="visible", key="generated_prompt_w_add_generation_prompt_text_area")168 169 access_token_no_cache = st.text_input("HuggingFace Access Token API with Write Access", type="password", key="access_token_no_cache")170 commit_message_text_input = st.text_input("Commit Message", key="commit_message_text_input")171 to_private_checkbox = st.checkbox("To Private Repo", key="to_private_checkbox")172 create_pr_checkbox = st.checkbox("Create PR (Check to contribute to others' model repository 🤗)", key="create_pr_checkbox")173 push_to_hub_button = st.button("Push to Hub", key="push_to_hub_button", use_container_width=True)174 st.session_state['tokenizer'].save_pretrained(st.session_state['tokenizer_json'])175 with open(f"{st.session_state['tokenizer_json']}/tokenizer_config.json", "r") as f:176 177 tokenizer_config_content = json.loads(f.read())178 shutil.rmtree(st.session_state['tokenizer_json'])179 180 st.download_button(181 label="Download tokenizer_config.json",182 data=json.dumps(tokenizer_config_content, indent=4),183 file_name='tokenizer_config.json',184 mime='application/json',185 use_container_width=True186 )187 st.download_button(188 label="Download chat_template.jinja2",189 data=st.session_state['successful_template'],190 file_name='chat_template.jinja2',191 mime='text/plain',192 use_container_width=True193 )194 if push_to_hub_button:195 with open(f"./tmp/{st.session_state['uuid']}/tmp_chat_template.json", "w") as fp:196 fp.write(st.session_state['successful_template'])197 with open(f"./tmp/{st.session_state['uuid']}/tmp_chat_template.json", "r") as f:198 successful_jinja_lines = f.readlines()199 st.session_state['tokenizer'].chat_template = sanitize_jinja2(successful_jinja_lines)200 try:201 with st.spinner(text="Pushing to hub ...", cache=False):202 st.session_state['tokenizer'].push_to_hub(203 repo_id=st.session_state['repo_id'], 204 commit_message=commit_message_text_input, 205 private=to_private_checkbox, 206 token=access_token_no_cache,207 create_pr=create_pr_checkbox)208 except Exception as e:209 st.write(f"Repo id: {st.session_state['repo_id']}")210 st.write(str(e))211 os.remove(f"./tmp/{st.session_state['uuid']}/tmp_chat_template.json")212