BUG FIXES + FASTER

app.py

@@ -5,14 +5,26 @@ from datasets import load_dataset
 from transformers import AutoTokenizer
 
 def ReturnTokens(dataset_name, tokenizer_name="openai-community/gpt2", split="train"):
+    # Initialize tokenizer
     tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+    
+    # Load dataset
     dataset = load_dataset(dataset_name, split=split)
-    def count_tokens(examples):
-        return sum(len(tokenizer.tokenize(example)) for example in examples)
-    total_tokens = 0
+    
+    # Function to count tokens in a single example
+    def count_tokens_in_example(example):
+        total_tokens = 0
+        tokenized = tokenizer.batch_encode_plus(example)
+        for i in tokenized:
+            total_tokens+=len(i)
+        return total_tokens
+    
+    tokens_=0
+
     for field in dataset[0].keys():
-        total_tokens += count_tokens(dataset[field])
-    return total_tokens
+        tokens_+=count_tokens_in_example(dataset[field])
+    
+    return tokens_
 
 with gr.Blocks(title="Dataset token counter") as app:
     gr.Markdown("# Token Counter")