Changing train.py

.idea/.name

@@ -1 +1 @@
-model.py
+train.py

train.py

@@ -3,27 +3,31 @@ import torch
 from transformers import BertTokenizer, Trainer, TrainingArguments
 from datasets import load_dataset
 
-# If the dataset is gated/private, make sure you have run huggingface-cli login
-dataset = load_dataset("NicolaiSivesind/human-vs-machine")
-# Tokenizer and Model Initialization
-tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+# 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 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():
-    return BertTokenizer.from_pretrained('./trained_model')
-
+    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,
@@ -32,17 +36,19 @@ training_args = TrainingArguments(
     warmup_steps=500,
     weight_decay=0.01,
     logging_dir='./logs',
-    evaluation_strategy="steps"
+    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
+    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")
-