Class Labels for Custom Datasets

🤗Datasets
3

Thanks! That worked and was a lot cleaner than my alternative solution.

# Creating a ClassLabel Object
df = dataset["train"].to_pandas()
labels = df['label'].unique().tolist()
ClassLabels = ClassLabel(num_classes=len(labels), names=labels)

# Mapping Labels to IDs
def map_label2id(example):
    example['label'] = ClassLabels.str2int(example['label'])
    return example

dataset = dataset.map(map_label2id, batched=True)

# Casting label column to ClassLabel Object
dataset = dataset.cast_column('label', ClassLabels)
3 Likes