|
import io |
|
import json |
|
import re |
|
|
|
import gradio as gr |
|
import matplotlib.pyplot as plt |
|
import pandas as pd |
|
from datasets import load_dataset |
|
from PIL import Image |
|
from transformers import AutoTokenizer |
|
|
|
tokenizers = [ |
|
"google/gemma-7b", |
|
"meta-llama/Llama-2-7b", |
|
"mistralai/Mistral-7B-v0.1", |
|
"facebook/opt-2.7b", |
|
"microsoft/phi-2", |
|
"THUDM/chatglm3-6b", |
|
"Qwen/Qwen1.5-7B-Chat", |
|
"bigscience/bloom-560m", |
|
"ise-uiuc/Magicoder-S-DS-6.7B", |
|
"google/flan-t5-base", |
|
"TinyLlama/TinyLlama-1.1B-Chat-v1.0", |
|
] |
|
|
|
|
|
def plot_histogram(data): |
|
plt.hist(data) |
|
plt.title("Histogram of number of tokens per dataset item") |
|
buf = io.BytesIO() |
|
plt.savefig(buf, format="png") |
|
buf.seek(0) |
|
im = Image.open(buf) |
|
return im |
|
|
|
|
|
def count(model_id, dataset_id, config, split, column, add_special_tokens=True): |
|
tokencounter = [] |
|
wordcounter = [] |
|
charcounter = [] |
|
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) |
|
if config == "": |
|
config is None |
|
dataset = load_dataset(dataset_id, config, split=split, trust_remote_code=True) |
|
pattern = r"[a-zA-Z]+" |
|
for item in dataset: |
|
tokens = tokenizer(item[column], add_special_tokens=add_special_tokens)["input_ids"] |
|
tokencounter.append(len(tokens)) |
|
charcounter.append(len(item[column])) |
|
|
|
words = re.findall(pattern, item[column]) |
|
wordcounter.append(len(words)) |
|
|
|
df = pd.DataFrame(tokencounter).describe().T |
|
df.insert(0, "type", "tokens") |
|
dfc = pd.DataFrame(charcounter).describe().T |
|
dfc.insert(0, "type", "chars") |
|
dfw = pd.DataFrame(wordcounter).describe().T |
|
dfw.insert(0, "type", "words") |
|
df.loc[-1] = dfw.values[0] |
|
df.index = df.index + 1 |
|
df.loc[-1] = dfc.values[0] |
|
df = df.round(1) |
|
df.drop("count", axis=1, inplace=True) |
|
|
|
return plot_histogram(tokencounter), df |
|
|
|
|
|
demo = gr.Interface( |
|
fn=count, |
|
title="Dataset token counts and distribution", |
|
inputs=[ |
|
gr.Dropdown(label="Tokenizer", choices=tokenizers, allow_custom_value=True), |
|
gr.Textbox(label="Dataset"), |
|
gr.Textbox(label="Config"), |
|
gr.Textbox(label="Split"), |
|
gr.Textbox(label="Column"), |
|
gr.Checkbox(label="Add special tokens", value=True), |
|
], |
|
outputs=[ |
|
gr.Image(), |
|
gr.Dataframe(label="Token, word and character counts per dataset item"), |
|
], |
|
examples=[ |
|
["mistralai/Mistral-7B-v0.1", "gsarti/flores_101", "eng", "dev", "sentence"], |
|
["mistralai/Mistral-7B-v0.1", "Muennighoff/flores200", "eng_Latn", "dev", "sentence"], |
|
["mistralai/Mistral-7B-v0.1", "wikitext", "wikitext-2-v1", "validation", "text"], |
|
["mistralai/Mistral-7B-v0.1", "hails/mmlu_no_train", "elementary_mathematics", "test", "question"], |
|
["mistralai/Mistral-7B-v0.1", "imdb", "", "test", "text"], |
|
["mistralai/Mistral-7B-v0.1", "gsm8k", "main", "test", "question"], |
|
["mistralai/Mistral-7B-v0.1", "locuslab/TOFU", "world_facts", "train", "question"], |
|
], |
|
cache_examples=False |
|
) |
|
|
|
demo.launch() |
|
|
