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 percentages")
    truncation_percentages = []
    recommended = None
    for max_len in TRUNCATION_LENGTHS:
        total_truncated_tokens = sum(max(length - max_len, 0) for length in sequence_lengths)
        truncation_percentage = total_truncated_tokens / total_tokens * 100
        truncation_percentages.append(truncation_percentage)
        if recommended is None and truncation_percentage < 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,
        "Percentage (%)": hist[0]/N_SAMPLES*100,
    })

    truncation_data = pd.DataFrame({
        "max_length": [str(value) for value in TRUNCATION_LENGTHS],
        "Percentage (%)": truncation_percentages,
    })

    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="Percentage (%)")
    truncation_percentage_plot = gr.BarPlot(None, title="Truncation percentage (how many tokens are discarded)", x="max_length", y="Percentage (%)")

    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_percentage_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 Percentage:** 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()