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")