Update train.py

train.py

@@ -1,8 +1,8 @@
 from transformers import BertTokenizer, BertForSequenceClassification, Trainer, TrainingArguments
-from datasets import load_dataset
 import numpy as np
-import pandas as pd 
+import pandas as pd
 from datasets import Dataset
+from sklearn.model_selection import train_test_split
 from sklearn.metrics import accuracy_score, precision_recall_fscore_support
 
 # Load dataset
@@ -15,12 +15,15 @@ tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
 def tokenize_function(examples):
     return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=512)
 
-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'])
+# Convert DataFrames to Datasets and apply tokenization
+train_dataset = Dataset.from_pandas(train_df)
+eval_dataset = Dataset.from_pandas(eval_df)
 
-train_dataset = Dataset.from_pandas(train_df).map(tokenize_function, batched=True)
-eval_dataset = Dataset.from_pandas(eval_df).map(tokenize_function, batched=True)
+train_dataset = train_dataset.map(tokenize_function, batched=True)
+train_dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'labels'])
+
+eval_dataset = eval_dataset.map(tokenize_function, batched=True)
+eval_dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'labels'])
 
 # Model
 model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=2)
@@ -62,3 +65,4 @@ trainer = Trainer(
 trainer.train()
 model.save_pretrained("./trained_model")
 tokenizer.save_pretrained("./trained_model")
+