import tiktoken
import gradio as gr

color_map = {
    '0': '#ccbfee',
    '1': '#beedc6',
    '2': '#f4d7ab',
    '3': '#f4aeb1',
    '4': '#a4dcf3',
}

def gen_num():
    n = 0
    while True:
        yield str(n)
        n = (n + 1) % 5

def tokens_len(prompt, model):
    encoder = tiktoken.encoding_for_model(model)
    tokens = encoder.encode(prompt)
    return (
        f'{len(tokens)} tokens, {len(prompt)} chracters',
        [(encoder.decode([token]), category) for 
            token, category in zip(tokens, gen_num())])

web_tokenizer = gr.Interface(
    fn = tokens_len,
    inputs = [
        gr.Textbox(
            label='輸入提示文字',
            lines=6),
        gr.Dropdown(
            label='模型名稱',
            choices = [
                'gpt-4o-mini',
                'gpt-4o',
                'gpt-4',
                'gpt-3.5-turbo'
            ],
            value='gpt-4o-mini')
    ],
    outputs = [
        gr.Textbox(
            label='統計量'),
        gr.Highlightedtext(
            label='切割結果：', 
            color_map = color_map,
            show_inline_category=False)
    ]
)

web_tokenizer.launch(share=True)