File size: 1,313 Bytes
6d07233 410a5d4 6d07233 410a5d4 6d07233 410a5d4 6d07233 410a5d4 6d07233 410a5d4 6d07233 410a5d4 6d07233 410a5d4 6d07233 |
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 29 30 31 32 33 34 35 36 37 38 39 |
import json
from datasets import load_dataset
from setfit import SetFitModel, SetFitTrainer
from huggingface_hub import upload_file
LABELS = ["pre-1900","1900–1945","1946–1990","1991–2008","2009–2015","2016–2018","2019–2022","2023–present"]
name2id = {n:i for i,n in enumerate(LABELS)}
ds = load_dataset("json", data_files={"train":"train.jsonl","val":"val.jsonl"})
# Assert all labels are present in train
seen = set([row["label"] for row in ds["train"]])
assert seen.issuperset(LABELS), f"Train set missing labels: {set(LABELS)-seen}"
ds = ds.map(lambda x: {"label": name2id[x["label"]]})
model = SetFitModel.from_pretrained("sentence-transformers/paraphrase-mpnet-base-v2", num_labels=len(LABELS))
trainer = SetFitTrainer(
model=model,
train_dataset=ds["train"],
eval_dataset=ds["val"],
metric="accuracy",
num_iterations=20,
num_epochs=2,
batch_size=16
)
trainer.train()
print("Eval:", trainer.evaluate())
# push model
repo_id = "DelaliScratchwerk/text-period-setfit"
trainer.push_to_hub(repo_id)
print("Pushed to:", repo_id)
# push labels.json alongside the model
with open("labels.json","w") as f:
json.dump(LABELS, f)
upload_file(path_or_fileobj="labels.json", path_in_repo="labels.json", repo_id=repo_id, repo_type="model")
print("Uploaded labels.json")
|