yagnik12's picture
Create train.py
736c7ec verified
raw
history blame
1.5 kB
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments
import evaluate
# 1. Load BiScope dataset
dataset = load_dataset("HanxiGuo/BiScope_Data")
# 2. Tokenizer
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)
# 3. Load model
model = AutoModelForSequenceClassification.from_pretrained(MODEL, num_labels=2)
# 4. Metrics
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)
# 5. Training args
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, # ✅ Upload to HF
hub_model_id="your-username/biscope-detector"
)
# 6. Trainer
trainer = Trainer(
model=model,
args=training_args,
train_dataset=encoded_dataset["train"],
eval_dataset=encoded_dataset["validation"],
tokenizer=tokenizer,
compute_metrics=compute_metrics,
)
# 7. Train & Push
trainer.train()
trainer.push_to_hub()