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