Kindly proof read the code below:
```
# Install the GPyOpt library
!pip install GPyOpt

# Download the tiny Shakespeare dataset
!curl -O https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt

# Imports
import os
import json
import numpy as np
import torch
import GPyOpt
import time
import math
import torch.nn as nn
import torch.optim as optim
from torch.nn import functional as F
from sklearn.cluster import KMeans

# Define hyperparameters
batch_size = 64
block_size = 256
max_iters = 5000
eval_interval = 500
device = 'cuda' if torch.cuda.is_available() else 'cpu'
eval_iters = 200
n_clusters = 4
total_experts = 8
aux_loss_weight = 0.0001

GLOBAL_ITER = 1  # Declare at the top of the script

# Set random seed for reproducibility
torch.manual_seed(1337)

# Load data and define data processing functions
with open('input.txt', 'r', encoding='utf-8') as f:
    text = f.read()

# Create a character-level vocabulary
chars = sorted(list(set(text)))
vocab_size = len(chars)

# Create mappings from characters to indices and vice versa
stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for i, ch in enumerate(chars)}

# Encode and decode functions to convert between text and indices
encode = lambda s: [stoi[c] for c in s]
decode = lambda l: ''.join([itos[i] for i in l])

# Convert entire text into a tensor of indices
data = torch.tensor(encode(text), dtype=torch.long)

# Split the data into training and validation sets
n = int(0.9 * len(data))
train_data = data[:n]
val_data = data[n:]

# Function to generate a batch of data for training or validation
def get_batch(split):
    data = train_data if split == 'train' else val_data
    ix = torch.randint(len(data) - block_size, (batch_size,))
    x = torch.stack([data[i:i + block_size] for i in ix])
    y = torch.stack([data[i + 1:i + block_size + 1] for i in ix])
    x, y = x.to(device), y.to(device)
    return x, y

# Function to estimate the model's loss on training and validation data
@torch.no_grad()
def estimate_loss(model):
    out = {}
    model.eval()  # Set model to evaluation mode
    for split in ['train', 'val']:
        total_loss = 0.0
        total_ce_loss = 0.0
        total_aux_loss = 0.0
        for k in range(eval_iters):
            X, Y = get_batch(split)
            # Unpack all outputs from the model
            logits, loss, ce_loss, aux_loss_value = model(X, Y)
            total_loss += loss.item()
            total_ce_loss += ce_loss.item()
            total_aux_loss += aux_loss_value.item()
        avg_loss = total_loss / eval_iters
        avg_ce_loss = total_ce_loss / eval_iters
        avg_aux_loss = total_aux_loss / eval_iters
        out[split] = {
            'avg_loss': avg_loss,
            'avg_ce_loss': avg_ce_loss,
            'avg_aux_loss': aux_loss_weight * avg_aux_loss
        }
    model.train()  # Set model back to training mode
    return out

def compute_metrics(model):
    model.eval()
    total_ce_loss = 0.0
    correct_predictions = 0
    total_tokens = 0

    with torch.no_grad():
        for _ in range(eval_iters):
            inputs, targets = get_batch('val')
            inputs = inputs.to(device)
            targets = targets.to(device)

            logits, _, ce_loss, _ = model(inputs, targets)
            total_ce_loss += ce_loss.item()

            predictions = torch.argmax(logits, dim=-1)
            correct_predictions += (predictions == targets).sum().item()
            total_tokens += targets.numel()

    avg_ce_loss = total_ce_loss / eval_iters

    try:
        perplexity = math.exp(avg_ce_loss)
    except OverflowError:
        perplexity = float('inf')  # Handle overflow by setting perplexity to infinity
    accuracy = correct_predictions / total_tokens
    return {'perplexity': perplexity, 'accuracy': accuracy}

class LiquidExpert(nn.Module):
    def __init__(self, input_dim, output_dim, hidden_dim, tau=0.1):
        super(LiquidExpert, self).__init__()
        self.hidden_dim = hidden_dim
        self.tau = tau

        # Define weights
        self.W = nn.Parameter(torch.randn(hidden_dim, hidden_dim) * 0.1)
        self.U = nn.Parameter(torch.randn(hidden_dim, input_dim) * 0.1)
        self.b = nn.Parameter(torch.zeros(hidden_dim))

        # Output layer
        self.output_layer = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        """
        x: Input tensor of shape (B, T, input_dim)
        """
        B, T, _ = x.shape
        h = torch.zeros(B, self.hidden_dim, device=x.device)  # Initial hidden state

        for t in range(T):  # Process each timestep
            u_t = x[:, t, :]  # Input at timestep t
            h_dot = -h + torch.tanh(self.W @ h.T + self.U @ u_t.T + self.b.unsqueeze(1))
            h = h + self.tau * h_dot.T

        output = self.output_layer(h)  # Final output
        return output

class ClusteredSwitchLiquid(nn.Module):
    """Clustered Mixture-of-Experts (MoE) with teams of specialized experts."""
    def __init__(self, n_clusters, total_experts, n_embd, dropout):
        super().__init__()
        self.n_clusters = n_clusters
        self.total_experts = total_experts

        # Initialize experts directly, there is no expert pool per cluster
        self.experts = nn.ModuleList([LiquidExpert(n_embd, dropout) for _ in range(self.total_experts)])

        # For managing expert distributions
        self.expert_distributions = nn.Parameter(torch.ones(n_clusters, total_experts) / total_experts)

    def forward(self, x):
        B, T, D = x.shape

        # Cluster assignment (may be better cached outside of the forward pass)
        if not hasattr(self, 'cluster_assignments'):
            cluster_assignments = KMeans(n_clusters=self.n_clusters).fit_predict(x.view(-1, D).detach().cpu().numpy())
            self.cluster_assignments = torch.tensor(cluster_assignments, device=x.device).view(B, T)  # Cache the assignments

        outputs_per_cluster = []
        for cluster_id in range(self.n_clusters):
            # Get indices for inputs belonging to this cluster
            assigned_indices = (self.cluster_assignments == cluster_id).nonzero(as_tuple=True)[0]
            if len(assigned_indices) == 0:
                continue  # Skip if no inputs are assigned to this cluster

            cluster_inputs = x[assigned_indices]  # Select relevant inputs

            # Get expert distribution for this cluster
            num_experts = torch.round(self.expert_distributions[cluster_id] * self.total_experts).long().tolist()
            active_experts = [self.experts[i] for i in range(self.total_experts) if i < num_experts]

            # Split the inputs for each expert
            even_split_inputs = torch.chunk(cluster_inputs, len(active_experts), dim=0)
            cluster_outputs = []
            for idx, expert in enumerate(active_experts):
                expert_out = expert(even_split_inputs[idx])  # Process data through the expert
                cluster_outputs.append(expert_out)

            # Combine outputs from all experts in this cluster
            outputs_per_cluster.append(sum(cluster_outputs))

        # Combine outputs from all clusters
        if outputs_per_cluster:  # Avoid error if all assignments result in empty
            final_output = torch.cat(outputs_per_cluster)
        else:
            final_output = torch.zeros(B, T, D, device=x.device)

        return final_output

class LiquidHead(nn.Module):
    """One head of self-attention with liquid dynamics."""

    def __init__(self, head_size, dropout):
        super().__init__()
        self.key = nn.Linear(n_embd, head_size, bias=False)
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)
        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))
        self.dropout = nn.Dropout(dropout)

        # Liquid neural dynamics parameters
        self.alpha_k = nn.Parameter(torch.randn(head_size))
        self.beta_k = nn.Parameter(torch.randn(head_size))
        self.alpha_q = nn.Parameter(torch.randn(head_size))
        self.beta_q = nn.Parameter(torch.randn(head_size))

    def forward(self, x):
        B, T, C = x.shape
        k = self.key(x)  # (B, T, head_size)
        q = self.query(x)  # (B, T, head_size)

        # Apply liquid dynamics to key and query
        k = self.alpha_k * k + self.beta_k * torch.tanh(k)
        q = self.alpha_q * q + self.beta_q * torch.tanh(q)

        # Compute scaled dot-product attention
        wei = q @ k.transpose(-2, -1) * C**-0.5  # (B, T, T)
        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))  # Apply causal mask
        wei = F.softmax(wei, dim=-1)  # (B, T, T)
        wei = self.dropout(wei)

        v = self.value(x)  # (B, T, head_size)
        out = wei @ v  # (B, T, head_size)
        return out

class MultiHeadAttention(nn.Module):
    """Multiple heads of self-attention in parallel."""

    def __init__(self, num_heads, head_size, dropout):
        super().__init__()
        self.heads = nn.ModuleList([LiquidHead(head_size, dropout) for _ in range(num_heads)])
        self.proj = nn.Linear(head_size * num_heads, n_embd)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        out = torch.cat([h(x) for h in self.heads], dim=-1)  # Concatenate all heads
        out = self.dropout(self.proj(out))
        return out

# Define the modified Block_ClusteredSwitch class
class Block_ClusteredSwitch(nn.Module):
    """Transformer block with Clustered Mixture-of-Experts FeedForward layer."""
    def __init__(self, n_embd, n_head, dropout):
        super().__init__()
        self.ln1 = nn.LayerNorm(n_embd)
        self.ln2 = nn.LayerNorm(n_embd)
        self.sa = MultiHeadAttention(n_head, (n_embd // n_head), dropout)
        self.ffwd = ClusteredSwitchLiquid(n_clusters, total_experts, n_embd, dropout)

    def forward(self, x):
        # Self-Attention with residual connection
        x = x + self.sa(self.ln1(x))  # (B, T, D)

        # Clustered MoE FeedForward with auxiliary loss
        x = self.ffwd(self.ln2(x))  # (B, T, D)

        # Return updated tensor (auxiliary loss omitted for simplicity)
        return x

# Define the modified Clustered_ST class
class Clustered_ST(nn.Module):
    """Clustered Switch Transformer with Dynamic Mixture-of-Experts."""
    def __init__(self, n_embd, n_head, n_layer, dropout):
        super().__init__()
        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
        self.position_embedding_table = nn.Embedding(block_size, n_embd)
        self.blocks = nn.ModuleList([Block_ClusteredSwitch(n_embd, n_head, dropout) for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(n_embd)  # Final layer norm
        self.lm_head = nn.Linear(n_embd, vocab_size)

        self.apply(self._init_weights)

    def _init_weights(self, module):
        """Initialize weights with normal distribution."""
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        tok_emb = self.token_embedding_table(idx)  # (B, T, n_embd)
        pos_emb = self.position_embedding_table(torch.arange(T, device=idx.device))  # (T, n_embd)
        x = tok_emb + pos_emb  # (B, T, n_embd)

        for block in self.blocks:
            x = block(x)  # (B, T, n_embd)

        x = self.ln_f(x)  # (B, T, n_embd)
        logits = self.lm_head(x)  # (B, T, vocab_size)

        if targets is None:
            loss = None
            ce_loss = None
            aux_loss_value = None
        else:
            B, T, C = logits.shape
            logits_flat = logits.view(B * T, C)
            targets_flat = targets.view(B * T)
            ce_loss = F.cross_entropy(logits_flat, targets_flat)
            aux_loss_value = 0  # Including auxiliary loss logic later if needed
            loss = ce_loss + aux_loss_weight * aux_loss_value  # Total loss with auxiliary term

        return logits, loss, ce_loss, aux_loss_value

    def generate(self, idx, max_new_tokens):
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -block_size:]
            outputs = self(idx_cond)
            logits = outputs[0]
            logits = logits[:, -1, :]  # Focus on the last time step
            probs = F.softmax(logits, dim=-1)  # Convert to probabilities
            idx_next = torch.multinomial(probs, num_samples=1)  # Sample the next token
            idx = torch.cat((idx, idx_next), dim=1)  # Append to the sequence
        return idx

# Define the objective function for Bayesian Optimization
def objective_function(hyperparameters):
    global GLOBAL_ITER  # Ensure this variable is being modified
    results = []

    for idx, params in enumerate(hyperparameters):
        # Extract standard hyperparameters
        learning_rate, n_embd, n_head, n_layer, weight_decay, dropout = params[:6]
        expert_distribution = params[6:]  # Extract expert distribution as integers

        # Ensure valid expert allocation
        if sum(expert_distribution) != total_experts:
            raise ValueError(f"Invalid expert allocation: {expert_distribution}, does not sum to {total_experts}")

        # Initialize model
        model = Clustered_ST(
            n_embd=int(n_embd),
            n_head=int(n_head),
            n_layer=int(n_layer),
            dropout=dropout
        ).to(device)

        # Apply expert distribution
        model.blocks[0].ffwd.expert_distributions.data = torch.tensor(expert_distribution, dtype=torch.float32).to(device)

        optimizer = optim.AdamW(
            model.parameters(),
            lr=learning_rate,
            weight_decay=weight_decay
        )

        best_acc = 0
        best_step = 0

        # Train the model
        for step in range(max_iters + 1):
            if step % eval_interval == 0:
                metrics = compute_metrics(model)
                if metrics['accuracy'] > best_acc:
                    best_acc = metrics['accuracy']
                    best_step = step

            xb, yb = get_batch('train')
            optimizer.zero_grad(set_to_none=True)
            loss = model(xb, yb)[1]
            loss.backward()
            optimizer.step()

        results.append(1 - best_acc)  # BO minimizes, so use 1 - accuracy
        print(f"BO Iteration {GLOBAL_ITER}: Best accuracy {best_acc:.4f} at step {best_step}/{max_iters}")
        GLOBAL_ITER += 1

    return np.array(results).reshape(-1, 1)

def constraint_sum_experts(params):
    """Ensure the sum of experts allocated equals the total available."""
    expert_allocations = params[6:]  # Extract expert allocations
    return np.array([np.sum(expert_allocations) - total_experts])  # Enforce total_experts

def save_best_hyperparameters(best_params, filepath="best_hyperparams.json"):
    with open(filepath, "w") as f:
        json.dump(best_params, f)

def load_best_hyperparameters(filepath="best_hyperparams.json"):
    try:
        with open(filepath, "r") as f:
            return json.load(f)
    except FileNotFoundError:
        return None

def episodic_bayesian_optimization(bo_iterations=10, bounds=None, prev_best_params=None, save_path="best_hyperparams.json"):
    if bounds is None:
        bounds = [
            {'name': 'learning_rate', 'type': 'continuous', 'domain': (1e-5, 1e-2)},
            {'name': 'n_embd', 'type': 'discrete', 'domain': (128, 256, 512, 1024)},
            {'name': 'n_head', 'type': 'discrete', 'domain': (2, 4, 8, 16)},
            {'name': 'n_layer', 'type': 'discrete', 'domain': (2, 4, 8, 12)},
            {'name': 'weight_decay', 'type': 'continuous', 'domain': (1e-6, 1e-2)},
            {'name': 'dropout', 'type': 'continuous', 'domain': (0.1, 0.5)},
            # Expert allocation per cluster, ensuring at least 1 expert per cluster
            {'name': f'expert_cluster_{i}', 'type': 'discrete', 'domain': tuple(range(1, total_experts + 1))}
            for i in range(n_clusters)
        ]

    constraints = [{'name': 'expert_sum_constraint', 'constraint': constraint_sum_experts}]

    initial_guess = np.array([
        [
            prev_best_params['learning_rate'],
            prev_best_params['n_embd'],
            prev_best_params['n_head'],
            prev_best_params['n_layer'],
            prev_best_params['weight_decay'],
            prev_best_params['dropout']
        ]
    ]) if prev_best_params else None

    print(f"Starting Bayesian Optimization for {bo_iterations} iterations...\n")

    bo = GPyOpt.methods.BayesianOptimization(
        f=objective_function,
        domain=bounds,
        X=initial_guess,
    )

    bo.run_optimization(max_iter=bo_iterations)

    best_hyperparams = {
        'learning_rate': float(bo.x_opt[0]),
        'n_embd': int(bo.x_opt[1]),
        'n_head': int(bo.x_opt[2]),
        'n_layer': int(bo.x_opt[3]),
        'weight_decay': float(bo.x_opt[4]),
        'dropout': float(bo.x_opt[5]),
    }

    save_best_hyperparameters(best_hyperparams, filepath=save_path)
    print(f"\nBest hyperparameters: {best_hyperparams}")

    best_accuracy = 1 - bo.fx_opt
    print(f"Best accuracy: {best_accuracy}")

    return best_hyperparams, best_accuracy

# Define a function to append results to a log file
def log_results(episode_number, best_params, best_accuracy, log_filepath="bo_log.json"):
    entry = {
        "episode_number": episode_number,
        "best_hyperparameters": best_params,
        "accuracy": best_accuracy
    }

    if not os.path.exists(log_filepath):
        with open(log_filepath, "w") as f:
            json.dump([], f)

    with open(log_filepath, "r") as f:
        data = json.load(f)

    data.append(entry)

    with open(log_filepath, "w") as f:
        json.dump(data, f, indent=2)

# Run the Bayesian Optimization (current episode: 1)
best_params, best_accuracy = episodic_bayesian_optimization(bo_iterations=3, save_path="best_HPs.json")
log_results(1, best_params, best_accuracy, log_filepath="results_log.json")
```