This fine tuning script is working as expected. However, it took 200 hours to run based on the current parameters. Is it possible to make it run faster by tuning the parameters?
import logging
import torch
from datasets import load_dataset, Dataset, Features, Value
from transformers import AutoModelForSequenceClassification, AutoTokenizer, Trainer, TrainingArguments, AutoConfig

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Configuration
BASE_MODEL_NAME = "cl-nagoya/ruri-v3-pt-130m"
DATASET_NAME = "cl-nagoya/ruri-v3-dataset-reranker"
OUTPUT_MODEL_NAME = "cl-nagoya/ruri-v3-reranker-130m"
OUTPUT_DIR = "./ruri-v3-reranker-training"

def generate_pairs(dataset):
    """
    Generator function to yield (query, document, score) pairs from the dataset.
    """
    for example in dataset:
        query = example['anc']

        # Positive pair
        yield {
            "query": query,
            "document": example['pos'],
            "labels": float(example['score.pos'])
        }

        # Negative pairs
        for neg_doc, neg_score in zip(example['neg'], example['score.neg']):
            yield {
                "query": query,
                "document": neg_doc,
                "labels": float(neg_score)
            }

def main():
    # 1. Load Dataset
    raw_dataset = load_dataset(DATASET_NAME)

    # 2. Create Datasets using generator for train and test splits
    # Split the original dataset first to maintain data distribution
    dataset_splits = raw_dataset['train'].train_test_split(test_size=12608, seed=42)

    # Create Dataset objects from generators
    train_dataset = Dataset.from_generator(
        lambda: generate_pairs(dataset_splits['train']),
        features=Features({
            "query": Value("string"),
            "document": Value("string"),
            "labels": Value("float32")
        })
    )

    eval_dataset = Dataset.from_generator(
        lambda: generate_pairs(dataset_splits['test']),
        features=Features({
            "query": Value("string"),
            "document": Value("string"),
            "labels": Value("float32")
        })
    )

    logger.info(f"Original dataset size: {len(raw_dataset)}")
    logger.info(f"Transformed training pairs: {train_dataset.num_rows}")
    logger.info(f"Transformed evaluation pairs: {eval_dataset.num_rows}")

    # 3. Load Tokenizer and Model
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_NAME)

    # Configure for Cross-Encoder Regression
    config = AutoConfig.from_pretrained(BASE_MODEL_NAME)
    config.num_labels = 1
    config.problem_type = "regression"

    model = AutoModelForSequenceClassification.from_pretrained(
        BASE_MODEL_NAME,
        config=config
    )

    # Set padding token if missing (common for GPT-style models)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        model.config.pad_token_id = model.config.eos_token_id

    # 4. Preprocessing Function
    def preprocess_function(examples):
        tokenized_inputs = tokenizer(
            examples["query"],
            examples["document"],
            truncation=True,
            padding=True,
            max_length=512,  # Adjust based on GPU memory
        )
        return tokenized_inputs

    # Apply tokenization (batched for efficiency)
    tokenized_train = train_dataset.map(
        preprocess_function,
        batched=True,
        remove_columns=['query', 'document']
    )

    tokenized_eval = eval_dataset.map(
        preprocess_function,
        batched=True,
        remove_columns=['query', 'document']
    )

    # 5. Training Setup
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    logger.info(f"Training on device: {device}")

    training_args = TrainingArguments(
        output_dir=OUTPUT_DIR,
        eval_strategy="steps",
        save_strategy="steps",
        save_steps=31520,
        learning_rate=2e-5,
        per_device_train_batch_size=32,  # Adjust based on GPU memory
        per_device_eval_batch_size=32,
        num_train_epochs=5,
        weight_decay=0.01,
        load_best_model_at_end=True,
        metric_for_best_model="eval_loss",
        bf16=True,
        logging_dir='./logs',
        logging_steps=1576,
        report_to="none"
    )

    # 6. Initialize Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_train,
        eval_dataset=tokenized_eval,
        processing_class=tokenizer
    )

    # 7. Train
    logger.info("Starting fine-tuning...")
    trainer.train()

    # 8. Save the Model
    logger.info(f"Saving fine-tuned model to {OUTPUT_MODEL_NAME}")
    trainer.save_model(OUTPUT_MODEL_NAME)
    tokenizer.save_pretrained(OUTPUT_MODEL_NAME)

    logger.info("Fine-tuning complete.")

if __name__ == "__main__":
    main()