mii/tokenizers
1
0
Fork 0
mirror of https://github.com/mii443/tokenizers.git synced 2025-08-23 00:35:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fixing the benchmark. (#1583)

Browse Source
This commit is contained in:
Nicolas Patry
2024-08-01 10:36:53 +02:00
committed by GitHub
parent 35f338a7b8
commit c6f2c0b057
1 changed files with 29 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
bindings/python/benches/test_tiktoken.py
View File

@ -7,6 +7,7 @@ import tiktoken
from tokenizers import Tokenizer from tokenizers import Tokenizer
from huggingface_hub import hf_hub_download from huggingface_hub import hf_hub_download
from typing import Tuple, List from typing import Tuple, List
from multiprocessing import Process
MODEL_ID = "meta-llama/Meta-Llama-3.1-8B" MODEL_ID = "meta-llama/Meta-Llama-3.1-8B"
DATASET = "facebook/xnli" DATASET = "facebook/xnli"
@ -24,8 +25,8 @@ def format_byte_size(num_bytes: int) -> Tuple[str, str]:
return f"{num_bytes_f:.2f} PB", "PB" return f"{num_bytes_f:.2f} PB", "PB"
def benchmark_batch(model: str, documents: list[str]) -> None: def benchmark_batch(model: str, documents: list[str], num_threads: int) -> None:
num_threads = int(os.environ["RAYON_NUM_THREADS"]) os.environ["RAYON_NUM_THREADS"] = str(num_threads)
num_bytes = sum(map(len, map(str.encode, documents))) num_bytes = sum(map(len, map(str.encode, documents)))
readable_size, unit = format_byte_size(num_bytes) readable_size, unit = format_byte_size(num_bytes)
print(f"==============") print(f"==============")
@ -84,18 +85,25 @@ def test(model: str, dataset: str, dataset_config: str, threads: List[int]):
input_lengths = [(10, False), (10_000, False), (10_000, True)] # Example input lengths input_lengths = [(10, False), (10_000, False), (10_000, True)] # Example input lengths
for num_threads in threads: for num_threads in threads:
os.environ["RAYON_NUM_THREADS"] = str(num_threads)
os.environ["TOKENIZER_PARALLELISM"] = str(num_threads)
os.environ["RAYON_RS_NUM_THREADS"] = str(num_threads)
for length, fuse in input_lengths: for length, fuse in input_lengths:
documents = [] documents = []
for i, item in enumerate(dataset_xnli["train"]): for i, item in enumerate(dataset_xnli["train"]):
documents.append("".join(item["premise"].values()))
if i >= length: if i >= length:
break break
documents.append("".join(item["premise"].values()))
if fuse: if fuse:
documents=["".join(documents)] documents=["".join(documents)]
benchmark_batch(model, documents)
# Rayon thread pool is global to a process, we need to launch
# separate processes in order to accurately use the correct number of threads.
# Otherwise, we're simply running tokenizers in whatever tests comes first.
# tokenizers does NOT provide a method to change the number of threads during
# runtime.
p = Process(target=benchmark_batch, args=(model, documents, num_threads))
p.start()
p.join()
# benchmark_batch(model, documents, num_threads)
def main(): def main():