from ray import tune

# Define a dictionary to represent the search space for hyperparameters
config = {
    "feature_size": tune.uniform(128, 512),  # Tune feature size between 128 and 512
    "num_layers": tune.randint(1, 4),  # Tune number of Transformer layers (1 to 4)
    "dropout": tune.uniform(0.05, 0.2),  # Tune dropout rate between 0.05 and 0.2
    "lr": tune.loguniform(1e-4, 1e-2),  # Tune learning rate using log-uniform sampling
}

class TransAmTrainable(tune.Trainable):

    def setup(self, config):
        self.model = TransAm(feature_size=config["feature_size"],
                              num_layers=config["num_layers"],
                              dropout=config["dropout"])
        self.model.to(device)
        self.criterion = nn.MSELoss()
        self.optimizer = torch.optim.SGD(self.model.parameters(), lr=config["lr"])

    def train_batch(self, train_data):
        data, targets = get_batch(train_data)  # Assuming your `get_batch` function exists
        optimizer.zero_grad()
        output, _ = self.model(data)
        loss = self.criterion(output, targets)
        loss.backward()
        optimizer.step()
        return loss.item()  # Return the training loss

    def validation_loss(self, val_data):
        self.model.eval()  # Set the model to evaluation mode
        with torch.no_grad():
            val_loss = 0.0
            for data, target in val_data:  # Iterate through validation data
                output, _ = self.model(data)
                val_loss += self.criterion(output, target).item()
            val_loss /= len(val_data)  # Average validation loss
        return val_loss  # Return the validation loss

    def train(self, train_data, val_data):
        # Train the model for a single epoch
        train_loss = self.train_batch(train_data)
        val_loss = self.validation_loss(val_data)
        # Update Weights & Biases with training and validation loss
        wandb.log({"train_loss": train_loss, "val_loss": val_loss})
        self.optimizer.step()
        return {"loss": val_loss}  # Return the validation loss for Ray Tune

from ray import tune
from wandb import init
import os
from ray.tune.suggest.hyperopt import HyperOptSearch

# Initialize Weights & Biases (replace with your project name)
init(project="Transformer Time Series Prediction")

def train_fn(config, checkpoint_dir=None):
    # Load model from checkpoint if provided (for resuming training)
    model = TransAmTrainable(config)  # Pass the config to the model
    if checkpoint_dir:
        model.load_checkpoint(checkpoint_dir)

    train_data, val_data = get_data()  # Load your training and validation data
    results = model.train(train_data, val_data)
    
    # Manually report the best checkpoint at the end of training
    if checkpoint_dir:
        model.save_checkpoint(checkpoint_dir)
        tune.report(loss=results["loss"], checkpoint=checkpoint_dir)
    else:
        tune.report(loss=results["loss"])


# Launch the Ray Tune experiment with PBT (Population Based Training)
analysis = tune.run(
    train_fn,
    config=config,  # Pass the configuration space
    metric="loss",  # Minimize validation loss
    mode="min",
    num_samples=16,  # Number of hyperparameter configurations to try
    # search_alg=tune.suggest.SigOptSearch(
    #     metric="loss",  # Optimize for validation loss
    #     mode="min"
    # ),
    search_alg=HyperOptSearch(
    metric="loss",  # Optimize for validation loss
    mode="min"
    ),
    scheduler=tune.schedulers.PopulationBasedTraining(
        time_attr="training_iteration",
        perturbation_interval=5,
        hyperparam_mutations={
            "lr": tune.uniform(1e-4, 1e-2),  # Learning rate between 1e-4 and 1e-2
            "optimizer": tune.choice(["adam", "sgd"])  # Choose between Adam and SGD
        }
    ),

    # Track experiment results in Weights & Biases
    name="TransAm_Hyperparameter_Tuning",
    local_dir="./ray_results",
)

# Print the best trial results
best_trial = analysis.best_trial
print("Best trial config:", best_trial.config)
print("Best trial final validation loss:", best_trial.last_result["loss"])

can you fix the above code