|
|
from datasets import load_dataset |
|
|
from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments |
|
|
import evaluate |
|
|
|
|
|
|
|
|
dataset = load_dataset("HanxiGuo/BiScope_Data") |
|
|
|
|
|
|
|
|
MODEL = "microsoft/deberta-v3-small" |
|
|
tokenizer = AutoTokenizer.from_pretrained(MODEL) |
|
|
|
|
|
def preprocess(examples): |
|
|
return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=256) |
|
|
|
|
|
encoded_dataset = dataset.map(preprocess, batched=True) |
|
|
|
|
|
|
|
|
model = AutoModelForSequenceClassification.from_pretrained(MODEL, num_labels=2) |
|
|
|
|
|
|
|
|
accuracy = evaluate.load("accuracy") |
|
|
|
|
|
def compute_metrics(eval_pred): |
|
|
logits, labels = eval_pred |
|
|
predictions = logits.argmax(axis=-1) |
|
|
return accuracy.compute(predictions=predictions, references=labels) |
|
|
|
|
|
|
|
|
training_args = TrainingArguments( |
|
|
output_dir="./results", |
|
|
evaluation_strategy="epoch", |
|
|
save_strategy="epoch", |
|
|
learning_rate=2e-5, |
|
|
per_device_train_batch_size=16, |
|
|
per_device_eval_batch_size=16, |
|
|
num_train_epochs=2, |
|
|
weight_decay=0.01, |
|
|
push_to_hub=True, |
|
|
hub_model_id="your-username/biscope-detector" |
|
|
) |
|
|
|
|
|
|
|
|
trainer = Trainer( |
|
|
model=model, |
|
|
args=training_args, |
|
|
train_dataset=encoded_dataset["train"], |
|
|
eval_dataset=encoded_dataset["validation"], |
|
|
tokenizer=tokenizer, |
|
|
compute_metrics=compute_metrics, |
|
|
) |
|
|
|
|
|
|
|
|
trainer.train() |
|
|
trainer.push_to_hub() |