training_no_mps.py

# Training without the use of MPS

from datasets import load_dataset
import transformers
from transformers import AutoTokenizer, TrainingArguments, Trainer, AutoModelForSequenceClassification
from datetime import datetime


transformers.logging.set_verbosity_info()

# Load dataset
dataset = load_dataset('dair-ai/emotion')

# Initialize tokenizer
tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')

# Tokenization function that moves data to MPS device
def tokenize(e):
    return tokenizer(e['text'], padding='max_length', truncation=True, max_length=128)

# Tokenize the dataset
tokenized_dataset = dataset.map(tokenize, batched=True)

# Load model and move it to MPS device
model = AutoModelForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=6)

# Set up training arguments
training_args = TrainingArguments(
    output_dir=f"./results/training-run_{datetime.now()}",
    num_train_epochs=3,
    per_device_train_batch_size=8,  # Reduced batch size for memory efficiency
    per_device_eval_batch_size=32,  # Reduced eval batch size
    gradient_accumulation_steps=4,  # Increased gradient accumulation for effective larger batch processing
    warmup_steps=500,
    weight_decay=0.01,
    logging_dir='./logs',
    logging_steps=10,
)

# Initialize the Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset["train"],
    eval_dataset=tokenized_dataset["validation"],
)

trainer.train()