Trying class weights

train.py

@@ -4,7 +4,18 @@ 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)