Changing train.py

train.py

@@ -2,6 +2,8 @@ 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
 
 # Load dataset dynamically or from a config
 dataset_name = "NicolaiSivesind/human-vs-machine"
@@ -9,6 +11,18 @@ 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)