Various bug fixes

app.py

@@ -4,19 +4,15 @@ import gradio as gr
 from datasets import load_dataset
 from transformers import AutoTokenizer
 
-def ReturnTokens(dataset, tokenizer="openai-community/gpt2", split="train"):
-    global tokens_
-    tokenizer=AutoTokenizer.from_pretrained(tokenizer)
-    dataset=load_dataset(dataset, split=split)
-    tokens_=0
-    def CountTokens(Example):
-        global tokens_
-        for k,i in enumerate(Example):
-            tokens_+=len(tokenizer.tokenize(i))
-    categories=[i for i in dataset[0].keys()]
-    for cat in categories:
-        CountTokens(dataset[cat])
-    return tokens_
+def ReturnTokens(dataset_name, tokenizer_name="openai-community/gpt2", split="train"):
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+    dataset = load_dataset(dataset_name, split=split)
+    def count_tokens(examples):
+        return sum(len(tokenizer.tokenize(example)) for example in examples)
+    total_tokens = 0
+    for field in dataset[0].keys():
+        total_tokens += count_tokens(dataset[field])
+    return total_tokens
 
 with gr.Blocks(title="Dataset token counter") as app:
     gr.Markdown("# Token Counter")