from transformers import AutoTokenizer

from data_measurements import DataMeasurementSuite
from data_measurements.measurements import LabelDistribution, TextDuplicates, TextLengths


gpt2_tokenizer = AutoTokenizer.from_pretrained("gpt2")


def tokenizer(x):
    return gpt2_tokenizer(x)["input_ids"]


suite = DataMeasurementSuite(
    dataset="hate_speech18",
    feature="text",
    label="label",
    split="train",
    tokenizer=tokenizer,
    measurements=[
        TextDuplicates,
        TextLengths,
        LabelDistribution,
    ],
)

results = suite.run()

results["text_lengths"].to_figure().show()
results["label_distribution"].to_figure().show()
