Update train.py

train.py

@@ -1,58 +1,29 @@
-from model import get_model
-import torch
-from transformers import BertTokenizer, Trainer, TrainingArguments
+from transformers import BertTokenizer, BertForSequenceClassification, 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
+# Load dataset
 dataset_name = "NicolaiSivesind/human-vs-machine"
 dataset = load_dataset(dataset_name)
 
+# Tokenizer
 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
+# Model
+model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=2)
+
+# Training Arguments
 training_args = TrainingArguments(
     output_dir="./results",
     num_train_epochs=3,
@@ -62,16 +33,28 @@ training_args = TrainingArguments(
     weight_decay=0.01,
     logging_dir='./logs',
     evaluation_strategy="steps",
-    save_steps=500, # Save model every 500 steps
+    save_steps=500,
     logging_steps=100,
 )
 
+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
+    }
+
 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
+    compute_metrics=compute_metrics
 )
 
 trainer.train()