File size: 2,662 Bytes
020c9d7 15c3ca6 020c9d7 6c014d0 5dbb891 020c9d7 5dbb891 d4939c3 020c9d7 d4939c3 020c9d7 6c014d0 020c9d7 d4939c3 020c9d7 15c3ca6 d4939c3 2518646 15c3ca6 d4939c3 15c3ca6 020c9d7 d4939c3 020c9d7 d4939c3 020c9d7 d4939c3 020c9d7 |
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 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 |
from model import get_model
import torch
from transformers import BertTokenizer, Trainer, TrainingArguments
from datasets import load_dataset
import numpy as np
from sklearn.metrics import accuracy_score, precision_recall_fscore_support
from torch.utils.data import DataLoader
from sklearn.utils.class_weight import compute_class_weight
# Other imports and code remain the same...
# Compute class weights
class_weights = compute_class_weight(
'balanced', classes=np.unique(train_dataset['labels']), y=train_dataset['labels'])
class_weights = torch.tensor(class_weights, dtype=torch.float)
# Update the model's classifier with class weights
model.classifier.weight.data = class_weights
# Load dataset dynamically or from a config
dataset_name = "NicolaiSivesind/human-vs-machine"
dataset = load_dataset(dataset_name)
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
def compute_metrics(pred):
labels = pred.label_ids
preds = np.argmax(pred.predictions, axis=1)
precision, recall, f1, _ = precision_recall_fscore_support(labels, preds, average='binary')
acc = accuracy_score(labels, preds)
return {
'accuracy': acc,
'f1': f1,
'precision': precision,
'recall': recall
}
def tokenize_function(examples):
# Add any specific preprocessing steps if necessary
return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=512)
def get_tokenizer():
try:
return BertTokenizer.from_pretrained('./trained_model')
except Exception:
return BertTokenizer.from_pretrained('bert-base-uncased')
tokenized_dataset = dataset.map(tokenize_function, batched=True)
tokenized_dataset = tokenized_dataset.rename_column("original_label_name", "labels")
tokenized_dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'labels'])
train_dataset = tokenized_dataset["train"]
eval_dataset = tokenized_dataset["validation"]
model = get_model()
# Make training arguments configurable
training_args = TrainingArguments(
output_dir="./results",
num_train_epochs=3,
per_device_train_batch_size=8,
per_device_eval_batch_size=8,
warmup_steps=500,
weight_decay=0.01,
logging_dir='./logs',
evaluation_strategy="steps",
save_steps=500, # Save model every 500 steps
logging_steps=100,
)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=train_dataset,
eval_dataset=eval_dataset,
compute_metrics=compute_metrics # Define this function to compute additional metrics
)
trainer.train()
model.save_pretrained("./trained_model")
tokenizer.save_pretrained("./trained_model")
|