Spaces:
Running
Running
File size: 3,541 Bytes
f9089ef |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 |
import gradio as gr
from datasets import load_dataset
from trl import SFTTrainer, SFTConfig
from transformers import AutoTokenizer
import pandas as pd
import numpy as np
TRUNCATION_LENGTHS = [128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768]
SEED = 42
N_SAMPLES = 1000
CODE_TEMPLATE = """
training_args = SFTConfig(
...,
max_length={},
)"""
def benchmark(model_name, dataset_name):
print(f"Running benchmark for model: {model_name} on dataset: {dataset_name}...")
print("Loading dataset...")
dataset = load_dataset(dataset_name, split="train", streaming=True).shuffle(seed=SEED).take(N_SAMPLES)
print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(model_name)
print("Tokenizing dataset...")
config = SFTConfig(max_length=None, bf16=False)
tokenized_dataset = SFTTrainer._prepare_dataset(
None, dataset, tokenizer, config, packing=False, formatting_func=None, dataset_name="train"
)
print("Computing the sequence lengths and total tokens")
sequence_lengths = [len(sample["input_ids"]) for sample in tokenized_dataset]
total_tokens = sum(sequence_lengths)
print("Computing the truncation ratios")
truncation_ratios = []
recommended = None
for max_len in TRUNCATION_LENGTHS:
total_truncated_tokens = sum(max(length - max_len, 0) for length in sequence_lengths)
truncation_ratio = total_truncated_tokens / total_tokens * 100
truncation_ratios.append(truncation_ratio)
if recommended is None and truncation_ratio < 5.0:
recommended = max_len
hist = np.histogram(sequence_lengths, bins=50)
lengths_distribution = pd.DataFrame({
"max_length": (hist[1][:-1] + hist[1][1:])/2,
"Ratio (%)": hist[0]/N_SAMPLES*100,
})
truncation_data = pd.DataFrame({
"max_length": [str(value) for value in TRUNCATION_LENGTHS],
"Ratio (%)": truncation_ratios,
})
return lengths_distribution, truncation_data, CODE_TEMPLATE.format(recommended)
with gr.Blocks() as demo:
model_input = gr.Textbox(label="Model Name", value="Qwen/Qwen3-0.6B")
dataset_input = gr.Textbox(label="Dataset Name", value="trl-lib/tldr")
run_button = gr.Button("Run estimation")
lengths_plot = gr.BarPlot(None, title="Length distribution", x="max_length", y="Ratio (%)")
truncation_ratio_plot = gr.BarPlot(None, title="Truncation ratio (how many tokens are discarded)", x="max_length", y="Ratio (%)")
recommended_code = gr.Code(CODE_TEMPLATE.format("..."), language="python", label="Recommended configuration")
run_button.click(fn=benchmark, inputs=[model_input, dataset_input], outputs=[lengths_plot, truncation_ratio_plot, recommended_code])
with gr.Accordion("See details", open=False):
gr.Markdown("""
This tool helps you choose an appropriate `max_length` value for your SFT training (`SFTConfig`) by analyzing the tokenized dataset.
**How it works:**
- Randomly samples 1,000 examples from your dataset.
- Prepares and tokenizes the data exactly as `SFTTrainer` would.
- Generates two visualizations:
- **Sequence Length Distribution:** Shows how long your tokenized sequences are.
- **Truncation Ratio:** Estimates the percentage of tokens that would be discarded (truncated) for different `max_length` values.
- Recommends the smallest `max_length` where truncation affects less than 5% of the tokens.
Use this tool to balance efficiency and memory usage when setting your `max_length` parameter.
""")
demo.launch()
|