Please implement the model as described. Update model where the clusters represent a team of expert specializing in the same thing. In other words, each cluster is a category of specialization and all the experts in a given cluster are experts of that category. However the distribution of the experts across the clusters is optimizable based on the best performing distribution. The distribution values should be optimized with bayesian optimization keeping in mind a set number of total experts and clusters and it has to find the best distribution of experts across cluster given the predefined values of each `{f'cluster{i+1}': (1, total_experts - (n_clusters - 1)) for i in range(n_clusters)}`. Lastly replace the standard feedforward networks with liquid neural network so each expert as a liquid architect. The experts of a selected cluster should all be activated when tokens are routed to it, however the tokens are divided evenly across the experts and the output of each is combined. Use the below code as a template and generate the complete script with all these changes:
```
# 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


# 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 scri

# 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():
    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()

            # Assuming logits shape: (batch_size, seq_length, vocab_size)
            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}

# Define the ClusteredSwitchFeedForward class with dynamic expert allocation
class ClusteredSwitchFeedForward(nn.Module):
    """
    Clustered Mixture-of-Experts (MoE) FeedForward layer with dynamic expert allocation.
    Each cluster dynamically selects a subset of experts from the global pool.
    The total number of experts used across all clusters is fixed.
    """
    def __init__(self, top_k, n_embd, dropout):
        super().__init__()
        self.n_clusters = n_clusters
        self.total_experts = total_experts
        self.top_k = top_k  # Number of experts selected per cluster

        # Initialize a global pool of experts
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(n_embd, 4 * n_embd),
                nn.ReLU(),
                nn.Linear(4 * n_embd, n_embd),
                nn.Dropout(dropout)
            ) for _ in range(self.total_experts)
        ])

        # Gating networks for clusters to assign experts
        self.cluster_gates = nn.ModuleList([
            nn.Linear(n_embd, self.total_experts) for _ in range(self.n_clusters)
        ])

        # Auxiliary parameters for enforcing expert utilization constraints
        self.register_buffer("expert_usage", torch.zeros(self.total_experts))

    def forward(self, x):
        """
        Args:
            x: Tensor of shape (B, T, D)

        Returns:
            output: Tensor of shape (B, T, D) after processing through experts
            total_aux_loss: Scalar tensor representing the sum of auxiliary losses
        """
        B, T, D = x.shape

        # Flatten batch and time dimensions
        x_flat = x.view(B * T, D)  # (B*T, D)

        # Initialize lists to collect selected experts and their probabilities
        all_selected_experts = []
        all_selected_probs = []

        # Step 1: Compute gating scores and select top-k experts per cluster
        for gate in self.cluster_gates:
            # Compute gating scores for each token
            gate_scores = gate(x_flat)  # (B*T, E)
            gate_scores = gate_scores / self.total_experts ** 0.5  # Scaling for numerical stability

            # Select top-k experts and compute their probabilities
            topk_scores, topk_indices = torch.topk(gate_scores, self.top_k, dim=-1)  # (B*T, K)
            topk_probs = F.softmax(topk_scores, dim=-1)  # Convert scores to probabilities

            all_selected_experts.append(topk_indices)  # Append expert indices
            all_selected_probs.append(topk_probs)      # Append corresponding probabilities

        # Concatenate selections from all clusters
        expert_indices = torch.cat(all_selected_experts, dim=-1)  # (B*T, C*K)
        expert_probs = torch.cat(all_selected_probs, dim=-1)      # (B*T, C*K)

        # Step 2: Prepare inputs for experts by expanding x_flat
        x_expanded = x_flat.unsqueeze(1).expand(-1, self.n_clusters * self.top_k, -1)  # (B*T, C*K, D)
        x_expanded = x_expanded.contiguous().view(-1, D)           # (B*T*C*K, D)
        expert_indices = expert_indices.contiguous().view(-1)      # (B*T*C*K,)
        expert_probs = expert_probs.contiguous().view(-1, 1)       # (B*T*C*K, 1)

        # Step 3: Process inputs for all experts in a batch
        # Initialize tensor to hold expert outputs
        expert_outputs = torch.zeros_like(x_expanded, device=x.device)  # (B*T*C*K, D)

        # Get unique expert indices to iterate over
        unique_expert_indices = expert_indices.unique()

        # For each expert, process all assigned inputs
        for idx in unique_expert_indices:
            mask = expert_indices == idx  # Boolean mask for current expert
            if mask.any():
                inputs = x_expanded[mask]            # Inputs for this expert
                outputs = self.experts[idx](inputs)  # Process inputs
                expert_outputs[mask] = outputs       # Store outputs

        # Step 4: Weight the expert outputs by gating probabilities
        expert_outputs = expert_outputs * expert_probs  # (B*T*C*K, D)

        # Step 5: Aggregate outputs for each token
        expert_outputs = expert_outputs.view(B * T, self.n_clusters * self.top_k, D)
        output = expert_outputs.sum(dim=1)  # Sum over all experts per token
        output = output.view(B, T, D)       # Reshape back to (B, T, D)

        # Step 6: Compute auxiliary loss for balanced expert usage
        # Count how many times each expert was selected
        all_allocations = expert_indices     # (B*T*C*K,)
        expert_counts = torch.bincount(all_allocations, minlength=self.total_experts).float()  # (E,)

        # Compute ideal usage per expert
        total_selections = all_allocations.numel()
        ideal_usage = total_selections / self.total_experts

        # Compute variance from ideal usage
        usage_variance = ((expert_counts - ideal_usage) ** 2).mean()

        # Compute entropy of expert usage distribution
        usage_distribution = expert_counts / (total_selections + 1e-10)
        entropy = -torch.sum(usage_distribution * torch.log(usage_distribution + 1e-10))

        # Total auxiliary loss combines variance and entropy
        total_aux_loss = aux_loss_weight * (usage_variance + entropy)

        return output, total_aux_loss

# Define the Head class (unchanged)
class Head(nn.Module):
    """One head of self-attention."""

    def __init__(self, head_size, n_embd, 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)

    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)
        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

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

    def __init__(self, num_heads, head_size, n_embd, dropout):
        super().__init__()
        self.heads = nn.ModuleList([Head(head_size, n_embd, 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), n_embd, dropout)
        self.ffwd = ClusteredSwitchFeedForward(1, n_embd, dropout)  # Adjust top_k as needed

    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
        ffwd_out, loss_aux = self.ffwd(self.ln2(x))  # (B, T, D), scalar

        # FeedForward with residual connection
        x = x + ffwd_out  # (B, T, D)

        # Return updated tensor and auxiliary loss
        return x, loss_aux

# 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)
        aux_loss = 0.0
        for block in self.blocks:
            x, loss_aux = block(x)
            aux_loss += loss_aux  # Accumulate auxiliary loss
        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 = aux_loss  # Note: aux_loss_weight is applied later
            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):
    """
    Objective function for Bayesian Optimization.
    Evaluates the model with given hyperparameters and tracks the best accuracy.
    """
    results = []  # Stores the optimization results
    global GLOBAL_ITER  # Use the global counter for iteration tracking

    for idx, params in enumerate(hyperparameters):
        learning_rate, n_embd, n_head, n_layer, weight_decay, dropout = params

        # Dynamically initialize these parameters within the model context
        model = Clustered_ST(
            n_embd=int(n_embd),
            n_head=int(n_head),
            n_layer=int(n_layer),
            dropout=dropout
        ).to(device)

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

        best_acc = 0
        best_step = 0  # Initialize best step tracker

        # 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')  # Replace with your data fetching logic
            optimizer.zero_grad(set_to_none=True)
            loss = model(xb, yb)[1]  # Assuming the model returns (logits, loss)
            loss.backward()
            optimizer.step()

        # Use the best accuracy recorded during training
        results.append(1 - best_acc)  # BO minimizes, so use 1 - accuracy

        # Print the best accuracy and step for this set of hyperparameters
        print(f"BO Iteration {GLOBAL_ITER}: Best accuracy {best_acc:.4f} at step {best_step}/{max_iters}")
        GLOBAL_ITER += 1  # Increment the global iteration counter

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

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)},
        ]

    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 this message BEFORE any evaluations
    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}")

    # Compute best accuracy as 1 - objective value since we used 1 - accuracy
    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)
```