Create trainer.py

trainer.py

@@ -0,0 +1,40 @@
+# Import Hugging Face libraries and datasets
+from transformers import Trainer, TrainingArguments, BertForSequenceClassification, BertTokenizer
+from datasets import load_dataset
+
+# Load dataset
+dataset = load_dataset("imdb")
+
+# Preprocess data: Tokenizer
+tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+
+def tokenize_function(examples):
+    return tokenizer(examples['text'], padding='max_length', truncation=True)
+
+# Apply tokenization
+tokenized_datasets = dataset.map(tokenize_function, batched=True)
+
+# Load pre-trained model
+model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=2)
+
+# Training Arguments
+training_args = TrainingArguments(
+    output_dir='/results',               # Save path in the Space environment
+    evaluation_strategy="epoch",         # Eval after each epoch
+    learning_rate=2e-5,                  # Learning rate
+    per_device_train_batch_size=8,       # Training batch size
+    per_device_eval_batch_size=8,        # Eval batch size
+    num_train_epochs=3,                  # Number of epochs
+    weight_decay=0.01                    # Weight decay for regularization
+)
+
+# Trainer setup
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=tokenized_datasets["train"],
+    eval_dataset=tokenized_datasets["test"]
+)
+
+# Train the model (this will run in the cloud, no need for local machine)
+trainer.train()