app.py

from cProfile import label
import logging
from pydoc import text
from turtle import title

import tiktoken
from transformers import AutoTokenizer

import gradio as gr

import os
from dotenv import load_dotenv
load_dotenv()

HF_TOKEN = os.getenv("HF_TOKEN") or None

logger = logging.getLogger(__name__)  # noqa


def load_test_phrases(filename):
    with open(f"./data/{filename}", "r", encoding="utf-8") as file:
        return file.read().splitlines()


models = ["salmankhanpm/gemma3-tokenizer-telugu",    # Gemma 3 Telugu Expansion
          "HuggingFaceTB/SmolLM2-135M-Instruct",     # SmolLM2
          "meta-llama/Llama-3.2-3B-Instruct",        # LLAMA-3
          "Telugu-LLM-Labs/Telugu-Llama2-7B-v0-Instruct", # LLama 2 Finetuned for Improving Telugu
          "CohereForAI/aya-23-8B",                   # AYA
          "google/gemma-3-4b-it",                    # GEMMA 3
          "google/gemma-3-270m-it",                  # Gemma 3 270m
          "sarvamai/sarvam-1",                       # SarvamAI
          "gpt-4o",                                  # GPT4o
          "openai/gpt-oss-20b",                      # OpenAI GPT OSS 
          "Qwen/Qwen3-4B",                           # Qwen
          "TWO/sutra-mlt256-v2"]                     # SUTRA

test_phrase_set = [
    "ఐదు వాక్యాలలో న్యూట్రాన్ స్కాటరింగ్ గురించి నాకు వివరణ ఇవ్వండి",
    "నటుడు బిల్ ముర్రే ఇంటి చిరునామా 445 నార్త్ బెడ్ఫోర్డ్ డ్రైవ్, లాస్ ఏంజిల్స్, CA 90049.",
    "హే, వీబ్ అంటే ఏమిటో మీకు తెలుసా?",
    "నాకు మరింత వివరాలు ఇవ్వండి",
    "కుక్కలు మరియు పిల్లుల మధ్య రెండు తేడాలు వాటి ప్రవర్తన మరియు వాటి శారీరక లక్షణాలు. కుక్కలు సాధారణంగా మరింత సామాజికంగా మరియు శిక్షణ పొందగలవు, అయితే పిల్లులు మరింత స్వతంత్రంగా ఉంటాయి. అదనంగా, పిల్లులతో పోలిస్తే కుక్కలు సాధారణంగా పెద్ద పరిమాణంలో ఉంటాయి.",
    "మేము టెక్స్ట్ మరియు వాయిస్‌కి మించి ఆంగ్లానికి మించిన అన్ని భాషల కోసం AIని రూపొందిస్తున్నాము. మేము అందరి కోసం AIని నిర్మిస్తున్నాము.",
]

test_phrase_set_long_1 = load_test_phrases('multilingualphrases01.txt')
test_phrase_set_long_2 = load_test_phrases('multilingualphrases02.txt')
test_phrase_set_long_3 = load_test_phrases('multilingualphrases03.txt')

# Cache tokenizers at the top level to avoid reloading on every call
hf_tokenizers = {}
for model in models:
    print("Loading tokenizer for", model)
    if  model != 'gpt-4o':
        hf_tokenizers[model] = AutoTokenizer.from_pretrained(model)
    else:
        hf_tokenizers[model] = tiktoken.encoding_for_model(model)


def generate_tokens_as_table(text):
    table = []
    for model in models:
        tokenizer = hf_tokenizers[model]
        if  model != 'gpt-4o':
            tokens = tokenizer.encode(text, add_special_tokens=False)
        else:
            tokens = tokenizer.encode(text)
        decoded = [tokenizer.decode([t]) for t in tokens]
        table.append([model] + decoded)
    return table

def baseline_tokeizer_extractor(text):
    word_count = len(text.split(' '))
    base_tokenizer = hf_tokenizers[models[-1]]
    base_tokens = base_tokenizer.encode(text)
    base_tokens_count = len(base_tokens)
    vocab_size = base_tokenizer.vocab_size

    return [models[-1], vocab_size, word_count, base_tokens_count, f"{base_tokens_count / word_count:0.2f}" ]


def generate_tokenizer_table(text):
    if not text:
        return []

    token_counts = {model: 0 for model in models}
    vocab_size = {model: 0 for model in models}

    for model in models:
        tokenizer = hf_tokenizers[model]
        if  model != 'gpt-4o':
            vocab_size[model] = tokenizer.vocab_size
            token_counts[model] = len(tokenizer.encode(text, add_special_tokens=True))
        else:
            vocab_size[model] = tokenizer.n_vocab
            token_counts[model] = len(tokenizer.encode(text))

    word_count = len(text.split(' '))



    base_tokenizer = baseline_tokeizer_extractor(text)

    output = []
    for m in models:
        # if m == models[-1]:
        #     row = [m, vocab_size[m], word_count, token_counts[m], f"{token_counts[m] / word_count:0.2f}", "1.0"]
        # else:
        row = [m, vocab_size[m], word_count, token_counts[m], f"{token_counts[m] / word_count:0.2f}", f"{token_counts[m] / base_tokenizer[3]:0.2f}"]

        output.append(row)

    return output



def generate_split_token_table(text):
    if not text:
        return gr.Dataframe()

    table = generate_tokenizer_table(text)
    return gr.Dataframe(
        table,
        headers=['tokenizer', 'v size', '#word', '#token', '#tokens/word', "NSL Value"],
        datatype=["str", "number", "str"],
        row_count=len(models),
        col_count=(6, "fixed"),
    )

def generate_baseline_tokenizer_insights(text):
    word_count = len(text.split(' '))
    base_tokenizer = hf_tokenizers[models[-1]]
    base_tokens = base_tokenizer.encode(text)
    base_tokens_count = len(base_tokens)

    # Table should be a list of rows (list of lists)
    table = [[models[-1], base_tokenizer.vocab_size, word_count, base_tokens_count, f"{base_tokens_count / word_count:0.2f}"]]

    return gr.Dataframe(
        table,
        headers=['tokenizer', 'v size', '#words', '#token', '#tokens/word'],
        datatype=["str", "number", "number", "number", "str"],
        row_count=1,
        col_count=(5, "fixed"),
    )

def generate_tokens_table(text):
    table = generate_tokens_as_table(text)
    cols = len(table[0])
    return gr.Dataframe(
        table,
        headers=['model'] + [str(i) for i in range(cols - 1)],
        row_count=2,
        col_count=(cols, "fixed"),
    )


with gr.Blocks() as sutra_tokenize:
    gr.Markdown(
        """
        ## Tokenize a sentence with various tokenizers and inspect how it's broken down.
        """)
    examples = test_phrase_set
    textbox = gr.Textbox(label="Input Text")
    with gr.Row():
        submit_button = gr.Button("Submit")
    gr.Examples(examples=examples, inputs=[textbox])
    example_display = gr.Textbox(label="Selected Example", interactive=False)
    textbox.change(lambda x: x, inputs=[textbox], outputs=[example_display])
    baseline_tokenizer = gr.Dataframe(label="Baseline Tokenizer")
    output_detailed = gr.Dataframe(label="Detailed Token Table")
    output_tokens = gr.Dataframe(label="Token Table")

    submit_button.click(lambda x: x, inputs=[textbox], outputs=[textbox])
    submit_button.click(generate_baseline_tokenizer_insights, inputs=[textbox], outputs=[baseline_tokenizer])
    submit_button.click(generate_split_token_table, inputs=[textbox], outputs=[output_detailed])
    submit_button.click(generate_tokens_table, inputs=[textbox], outputs=[output_tokens])



if __name__ == '__main__':
    with gr.Blocks(analytics_enabled=False) as demo:
        with gr.Row():
            gr.Markdown(
                """
                # Telugu Tokenizer Sentence Inspector & Evaluation.
                ### Using Two SUTRA Tokenizer as Baseline.
                #### Paper :
                ### [EVALUATING TOKENIZER PERFORMANCE OF LARGE LANGUAGE MODELS ACROSS OFFICIAL INDIAN LANGUAGES](https://arxiv.org/pdf/2411.12240v2) \n
                ### [PERFORMANCE EVALUATION OF TOKENIZERS IN LARGE LANGUAGE MODELS FOR THE ASSAMESE LANGUAGE](https://arxiv.org/pdf/2410.03718)
                ####

                """
            )
        with gr.Row():
            gr.TabbedInterface(
                interface_list=[sutra_tokenize],
                tab_names=["Tokenize Text"]
            )

demo.queue(default_concurrency_limit=5).launch(
    server_name="0.0.0.0",
    allowed_paths=["/"],
)