File size: 781 Bytes
b6fc63b | 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 | 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") |