from transformers import AutoTokenizer, AutoModelForSequenceClassification

model_name = "distilbert-base-uncased"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(
    model_name,
    num_labels=6,
    id2label={
        0: "Australopithecus",
        1: "Homo habilis",
        2: "Homo erectus",
        3: "Homo heidelbergensis",
        4: "Homo neanderthalensis",
        5: "Homo sapiens"
    },
    label2id={
        "Australopithecus": 0,
        "Homo habilis": 1,
        "Homo erectus": 2,
        "Homo heidelbergensis": 3,
        "Homo neanderthalensis": 4,
        "Homo sapiens": 5
    }
)

model.save_pretrained("human-evolution-classifier")
tokenizer.save_pretrained("human-evolution-classifier")