Add a token counter similar to automatic1111

It can now be found in the Default and Notebook tabs
This commit is contained in:
oobabooga 2023-08-20 19:13:01 -07:00
parent 120fb86c6a
commit 429cacd715
4 changed files with 38 additions and 14 deletions

View file

@ -18,7 +18,9 @@ def create_ui():
with gr.Row():
with gr.Column(scale=4):
with gr.Tab('Raw'):
shared.gradio['textbox-notebook'] = gr.Textbox(value='', elem_classes=['textbox', 'add_scrollbar'], lines=27)
with gr.Row():
shared.gradio['textbox-notebook'] = gr.Textbox(value='', elem_classes=['textbox', 'add_scrollbar'], lines=27)
shared.gradio['token-counter-notebook'] = gr.HTML(value="<span>0</span>", elem_classes=["token-counter"])
with gr.Tab('Markdown'):
shared.gradio['markdown_render-notebook'] = gr.Button('Render')
@ -45,9 +47,6 @@ def create_ui():
shared.gradio['save_prompt-notebook'] = gr.Button('💾', elem_classes=['refresh-button', 'refresh-button-small'])
shared.gradio['delete_prompt-notebook'] = gr.Button('🗑️', elem_classes=['refresh-button', 'refresh-button-small'])
shared.gradio['count_tokens-notebook'] = gr.Button('Count tokens')
shared.gradio['status-notebook'] = gr.Markdown('')
def create_event_handlers():
shared.gradio['Generate-notebook'].click(
@ -86,5 +85,5 @@ def create_event_handlers():
lambda x: x + '.txt', gradio('prompt_menu-notebook'), gradio('delete_filename')).then(
lambda: gr.update(visible=True), None, gradio('file_deleter'))
shared.gradio['count_tokens-notebook'].click(count_tokens, gradio('textbox-notebook'), gradio('status-notebook'), show_progress=False)
shared.gradio['textbox-notebook'].change(lambda x : f"<span>{count_tokens(x)}</span>", gradio('textbox-notebook'), gradio('token-counter-notebook'), show_progress=False)
shared.gradio['get_logits-notebook'].click(logits.get_next_logits, gradio('textbox-notebook'), gradio('logits-notebook'))