class CNN(nn.Module):
    def __init__(self, num_classes):
        super().__init__()
        # TODO: Instantiate the layers your CNN will have!
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, stride=2, padding=1)
        self.conv2 = nn.Conv2d(in_channels=16, out_channels=16, kernel_size=7, stride=2, padding=1)
        self.linear = nn.Linear(in_features=16*54*54, out_features=num_classes)
        # Optional: You can also define your ReLU layers here if you'd like
        # Otherwise, if you decide to use F.relu you can just remove these lines
        self.relu1 = nn.ReLU()
        self.relu2 = nn.ReLU()

    def forward(self, x):
        # TODO: 1. Pass the input through the first conv layer
        conv1_out = self.conv1(x)

        # TODO: 2. Apply ReLU
        relu1_out = self.relu1(conv1_out)

        # TODO: 3. Pass to the 2nd conv layer
        conv2_out = self.conv2(relu1_out)

        # TODO: 4. Apply ReLU
        relu2_out = self.relu2(conv2_out)

        # TODO: 5. Flatten before linear layer
        flat = relu2_out.view(relu2_out.size(0), -1)

        # TODO: 6. Pass to linear layer
        out = self.linear(flat)

        return out

class myDataset(Dataset):

    def __init__(self, ...):
        raise NotImplementedError
        # save data
        # initialize preprocessing
        # set up configurations

    def __len__(self):
        raise NotImplementedError
        # return total number(rows) of dataset

    def __getitem__(self, idx):
        raise NotImplementedError
        # return a single data and its label from dataset
        # preprocessing will also be done here

class MiniImageNetDataset(Dataset):
    def __init__(
        self,
        dataset,
        class_id_to_name: dict[int, str],
        split: str = "train",
        transform: transforms.Compose = None,
    ):
        """
        Args:
            dataset (datasets.Dataset): HuggingFace dataset object
                Each sample in the dataset has:
                - An "image" field that contains the image data
                - A "label" field that contains the numeric class ID
            class_id_to_name (dict): dictionary mapping numeric class IDs to class names
            split (str): "train" or "val" or "test"
        """
        # TODO: Implement the __init__ method of the MiniImageNetDataset
        self.split = split
        self.dataset = dataset
        self.transform = transform
        self.class_id_to_name = class_id_to_name

    def __len__(self):
        # TODO: Implement the __len__ method of the MiniImageNetDataset
        return len(self.dataset)

    def __getitem__(self, index):
        """
        Returns a single sample from the dataset at the given index.
        Args:
            index (int): index of the item to get
        Returns:
            tuple: (image, label) where image is a tensor of shape C x H x W and label is a tensor of shape 1
        """

        # TODO: Implement the __getitem__ method of the MiniImageNetDataset
        # Don't forget to...
        # - Cast the item and label to torch.Tensor with the correct dtype
        # - Permute the image into the shape (C, H, W)
        # - Apply the transform if it's not None
        sample = self.dataset[index]
        img = sample['image']
        label = sample['label']

        img = img.convert("RGB")
        
        if self.transform is not None:
            img=self.transform(img)

        label = torch.tensor(label,dtype=torch.long)
        return img, label

    def show_images(self, num_images: int = 10):
        """
        Visualize images in a grid layout.

        Args:
            num_images (int): Number of images to display per class

        Returns:
            tuple: (fig, ax) - Matplotlib figure and axes objects
        """
        # TODO: Implement the show_images method
        # Make sure to call fig.show() or plt.show() to display your grid of images at the end!
        
        # 所有类别 ID，按从小到大排序，方便排版
        class_ids = sorted(self.class_id_to_name.keys())
        n_classes = len(class_ids)

        # 创建子图：每一行一个类别，每行 num_images 张图
        fig, axes = plt.subplots(n_classes, num_images, figsize=(num_images * 2, n_classes * 2))

        # 当只有 1 行或 1 列时，axes 不是二维数组，做个统一处理
        if n_classes == 1:
            axes = [axes]
        if num_images == 1:
            axes = [[ax] for ax in axes]

        # 对每个类别画图
        for row, cls in enumerate(class_ids):
            # 找到属于该类别的样本下标
            idxs = [i for i, y in enumerate(self.dataset["label"]) if y == cls]
            # 只取前 num_images 个
            idxs = idxs[:num_images]

            for col, idx in enumerate(idxs):
                img = self.dataset[idx]["image"]  # 通常是 PIL.Image
                ax = axes[row][col]
                ax.imshow(img)
                ax.axis("off")
                # 每一行第一个子图上写类别名字
                if col == 0:
                    ax.set_title(self.class_id_to_name[cls])

        plt.tight_layout()
        plt.show()  # 或 fig.show()
        return fig, axes

train_transform = transforms.Compose([
    transforms.Resize(256),                  # Resize the shorter side to 256
    transforms.RandomCrop(224),              # Random crop into shape 224×224
    transforms.RandomHorizontalFlip(),       # Random horizontal flip
    transforms.ToTensor(),                   # Convert to tensor, scale to [0,1], and permute into shape (C, H, W)
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225]) # Normalize
])

val_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

train_dataset = MiniImageNetDataset(mini_imagenet_train, class_id_to_name, split="train", transform=train_transform)
val_dataset = MiniImageNetDataset(mini_imagenet_val, class_id_to_name, split="val", transform=val_transform)

train_loss = loss.item()
preds = torch.argmax(y_hat, dim=1)
train_correct += (preds == y).sum().item()

Step11: 每个epoch结束之后计算`val_dataset`上的准确率, 这段代码是在内层循环外面的

```python
model.eval()
with torch.no_grad():
    for x,y in val_dataloader:
        x = x.to(device, dtype = torch.float)
        y = y.to(device, dtype = torch.long)

        y_hat = model(x)
        loss = criterion(y_hat, y)
        val_loss += loss.item()
        preds = torch.argmax(y_hat, dim=1)
        val_correct += (preds == y).sum().item()

val_loss /= len(val_dataloader)

# === SETUP ===
model.to(device)

# === EPOCH LOOP ===
for epoch in range(num_epochs):
    train_loss = 0.0
    train_correct = 0.0
    
    # === TRAINING PHASE ===
    model.train()  # Enable training mode
    
    for x, y in train_dataloader:
        # 1. Move data to device
        x, y = x.to(device), y.to(device)
        
        # 2. Reset gradients
        optimizer.zero_grad()
        
        # 3. Forward pass
        y_hat = model(x)
        
        # 4. Compute loss
        loss = criterion(y_hat, y)
        
        # 5. Backward pass (compute gradients)
        loss.backward()
        
        # 6. Update weights
        optimizer.step()
        
        # 7. Track metrics (optional)
        # ...
    
    # === VALIDATION PHASE ===
    model.eval()  # Disable training-specific layers
    with torch.no_grad():  # Don't compute gradients
        for x, y in val_dataloader:
            # Forward pass only, no backward pass!
            y_hat = model(x)
            loss = criterion(y_hat, y)
            # Track validation metrics

# Common practice to check if GPU is available
if torch.accelerator.is_available():
    device = torch.accelerator.current_accelerator()
else:
    device = "cpu"
print(f"Using device: {device}")
assert str(device) in ["cuda", "mps"], "Please make sure you are using an accelerator (e.g. GPU or MPS)"


def train(
    model, optimizer, criterion, num_epochs, train_dataloader, val_dataloader=None
):
    """
    Args:
        model: the model to train
        optimizer: the optimizer to use
        criterion: the loss function to use
        num_epochs: the number of epochs to train for
        train_dataloader: the dataloader for the training set
        val_dataloader: the dataloader for the validation set
    """
    # === SETUP ===
    model.to(device)

    # == (Optional) Perform a validation loop before training as a baseline ==
    if val_dataloader:
        model.eval()  # Set model to evaluation mode
        num_correct = 0.0
        val_loss = 0.0
        with torch.no_grad():  # Don't compute gradients
            for x, y in val_dataloader:
                x = x.to(device, dtype=torch.float)
                y = y.to(device, dtype=torch.long)
                y_hat = model(x)

                loss = criterion(y_hat, y)  # pred, target
                val_loss += loss.item()

                preds = torch.argmax(y_hat, dim=1)
                num_correct += (preds == y).sum().item()

        val_acc = num_correct / len(val_dataloader.dataset)
        tqdm.write(f"Initial validation loss: {val_loss:.4f}")
        tqdm.write(f"Initial validation accuracy: {val_acc:.4f}")

    # Lists to store metrics across epochs
    train_accuracies = []
    train_losses = []
    val_accuracies = []
    val_losses = []

    epoch_pbar = tqdm(range(num_epochs), desc="Training Progress", position=0)

    # === EPOCH LOOP ===
    for epoch in epoch_pbar:
        # === TRAINING PHASE ===
        model.train()  # Set model to training mode

        # Initialize metrics for this epoch
        train_loss = 0.0
        train_correct = 0.0

        # === INNER LOOP (iterate over training batches) ===
        for x, y in train_dataloader:
            # 1. Move data to device
            x = x.to(device, dtype=torch.float)
            y = y.to(device, dtype=torch.long)

            # 2. Reset gradients
            optimizer.zero_grad()

            # 3. Forward pass
            y_hat = model(x)

            # 4. Compute loss
            loss = criterion(y_hat, y)

            # 5. Backward pass (compute gradients)
            loss.backward()

            # 6. Update weights
            optimizer.step()

            # 7. Track metrics
            train_loss += loss.item()  # Accumulate loss (convert to Python float)
            preds = torch.argmax(
                y_hat, dim=1
            )  # Get predicted class for each item in the batch (highest prob = highest confidence)
            train_correct += (preds == y).sum().item()  # Count correct predictions

        # === END OF INNER LOOP ===

        # Compute average metrics for the epoch
        avg_train_loss = train_loss / len(
            train_dataloader
        )  # Average loss = total loss / total number of batches
        train_acc = (
            train_correct / len(train_dataloader.dataset)
        )  # Average accuracy = accuracy / total number of data points seen across the whole epoch (i.e. in the entire dataset)
        train_losses.append(avg_train_loss)
        train_accuracies.append(train_acc)

        print(f"\nEpoch {epoch + 1} - Training accuracy: {train_acc:.3f}")
        print(f"Epoch {epoch + 1} - Training loss: {train_loss:.3f}")

        # === END OF TRAINING PHASE ===

        # === VALIDATION PHASE ===
        if val_dataloader:
            model.eval()  # Set model to evaluation mode
            val_loss = 0.0
            val_correct = 0.0

            with torch.no_grad():  # Don't compute gradients
                # === INNER LOOP (iterate over validation batches) ===
                for x, y in val_dataloader:
                    # 1. Move data to device
                    x = x.to(device, dtype=torch.float)
                    y = y.to(device, dtype=torch.long)

                    # 2. Forward pass only, no backward pass!
                    y_hat = model(x)

                    # 3. Compute loss
                    loss = criterion(y_hat, y)

                    # 4. Track validation metrics
                    val_loss += loss.item()
                    preds = torch.argmax(y_hat, dim=1)
                    val_correct += (preds == y).sum().item()

                # === END OF INNER LOOP ===

                # Compute average metrics for the epoch
                avg_val_loss = val_loss / len(val_dataloader)
                val_acc = val_correct / len(val_dataloader.dataset)

                val_losses.append(avg_val_loss)
                val_accuracies.append(val_acc)

            # === END OF VALIDATION PHASE ===

        # === END OF EPOCH LOOP ===

        # Print metrics for the epoch
        print(f"Epoch {epoch + 1} - Validation accuracy: {val_acc:.3f}")
        print(f"Epoch {epoch + 1} - Validation loss: {val_loss:.3f}")

        epoch_pbar.set_postfix(
            {
                "train_loss": f"{train_loss:.3f}",
                "train_acc": f"{train_acc:.3f}",
                "val_loss": f"{val_loss:.3f}",
                "val_acc": f"{val_acc:.3f}",
            }
        )

    return train_accuracies, train_losses, val_accuracies, val_losses

def plot_metrics(train_losses, val_losses, train_accuracies=None, val_accuracies=None, num_epochs=None, title=""):
    """
    Plots the training loss, training accuracy, validation loss, and validation accuracy.
    Args:
        train_losses: list of training losses
        val_losses: list of validation losses
        train_accuracies: list of training accuracies
        val_accuracies: list of validation accuracies
        num_epochs: number of epochs
        title: title of the plot
    """
    fig, axes = plt.subplots(2, 2, figsize=(12, 8))

    # 训练损失
    axes[0, 0].plot(train_losses, label='Training Loss')
    axes[0, 0].set_title('Training Loss')
    axes[0, 0].set_xlabel('Epoch')
    axes[0, 0].set_ylabel('Loss')

    axes[0, 1].plot(train_accuracies, label='Training Accuracy')
    axes[0, 1].set_title('Training Accuracy')
    axes[0, 1].set_xlabel('Epoch')
    axes[0, 1].set_ylabel('Accuracy')

    axes[1, 0].plot(val_losses, label='Validation Loss')
    axes[1, 0].set_title('Validation Loss')
    axes[1, 0].set_xlabel('Epoch')
    axes[1, 0].set_ylabel('Loss')

    axes[1, 1].plot(val_accuracies, label='Validation Accuracy')
    axes[1, 1].set_title('Validation Accuracy')
    axes[1, 1].set_xlabel('Epoch')
    axes[1, 1].set_ylabel('Accuracy')

    plt.tight_layout()
    plt.suptitle(title)
    plt.show()
    
    
    
plot_metrics(cnn_train_losses, cnn_val_losses, cnn_train_accuracies, cnn_val_accuracies, num_epochs=20, title="CNN Training Metrics")

                    输入 x (in_channels)
                        │
          ┌─────────────┴─────────────┐
          │                           │
          ▼                           │
    ┌─────────┐                       │
    │ Conv1   │ ──→ BatchNorm1 ──→ ReLU
    └────┬────┘                       │
         │                            │
    ┌────┴────┐                       │
    │ Conv2   │ ──→ BatchNorm2        │
    └────┬────┘                       │
         │                            │
    ┌────┴────────────────────────────┐
    │              +                  │  ← 残差相加
    └────┬────────────────────────────┘
         │
      ReLU (out)

# 输入: conv_output shape = [B=4, C=3, H=2, W=2]
# B=batch size, C=channels, H=height, W=width

x = torch.tensor([
    # batch 0
    [[[1, 2],      # channel 0
      [3, 4]],
     [[5, 6],      # channel 1  
      [7, 8]],
     [[9, 10],     # channel 2
      [11, 12]]],
    # batch 1
    [[[2, 3],
      [4, 5]],
     [[6, 7],
      [8, 9]],
     [[10, 11],
      [12, 13]]],
    # batch 2...
    # batch 3...
])

┌─────────────────────────────────────────────────────────────┐
│  BatchNorm2d 对每个通道独立计算：                            │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│  Channel 0 (所有batch):                                     │
│  ┌─────────┬─────────┬─────────┬─────────┐                 │
│  │ batch0  │ batch1  │ batch2  │ batch3  │                 │
│  │ [[1,2], │ [[2,3], │ ...     │ ...     │                 │
│  │  [3,4]] │  [4,5]] │         │         │                 │
│  └─────────┴─────────┴─────────┴─────────┘                 │
│        │           │                               × 4    │
│        └───────────┴───────────────────────→               │
│                 展平成: [1,2,3,4, 2,3,4,5, ...]            │
│                                                             │
│  μ₀ = mean([1,2,3,4,2,3,4,5, ...])                         │
│  σ₀² = var([1,2,3,4,2,3,4,5, ...])                         │
│                                                             │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│  Channel 1:                                                 │
│  ┌─────────┬─────────┬─────────┬─────────┐                 │
│  │ [[5,6], │ [[6,7], │ ...     │ ...     │                 │
│  │  [7,8]] │  [8,9]] │         │         │                 │
│  └─────────┴─────────┴─────────┴─────────┘                 │
│        │           │                               × 4    │
│        └───────────┴───────────────────────→               │
│                 展平成: [5,6,7,8,6,7,8,9, ...]             │
│                                                             │
│  μ₁ = mean([5,6,7,8,6,7,8,9, ...])                         │
│  σ₁² = var([5,6,7,8,6,7,8,9, ...])                         │
│                                                             │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│  Channel 2:  同理计算 μ₂, σ₂²                               │
│                                                             │
└─────────────────────────────────────────────────────────────┘

结果：3个通道 → 3个均值(μ₀, μ₁, μ₂) + 3个方差(σ₀², σ₁², σ₂²)

class ResidualBlock(nn.Module):
    def __init__(
        self,
        in_channels=3,
        out_channels=64,
        kernel_size=3,
        initial_downsample=False, # whether to downsample the input
        verbose=False, # whether to print debug statements
    ):
        super().__init__()
        self.verbose = verbose


        # TODO: Define the first conv layer
        self.conv1 = nn.Conv2d(in_channels=in_channels, out_channels=out_channels,
        kernel_size=kernel_size,padding=1,stride=1 if not initial_downsample else 2)
        # TODO: Define the first batchnorm layer
        self.batchnorm1 = nn.BatchNorm2d(num_features=out_channels)
        # TODO: Define the second conv layer
        self.conv2 = nn.Conv2d(in_channels=out_channels,out_channels=out_channels,kernel_size=kernel_size,
        padding=1,stride=1)
        # TODO: Define the second batchnorm layer
        self.batchnorm2 = nn.BatchNorm2d(num_features=out_channels)
        # TODO: Define the residual connection
        self.residual_connection = nn.Conv2d(in_channels=in_channels,out_channels=out_channels,kernel_size=1,
        stride=1 if not initial_downsample else 2)

    def forward(self, x):
        # TODO: Implement the forward pass
        conv1 = self.conv1
        if self.verbose:
            print(f"conv1 shape: {conv1.shape}")
        bn1 = self.batchnorm1
        if self.verbose:
            print(f"bn1 shape: {bn1.shape}")

        relu=nn.ReLU()
        z1 = relu(bn1(conv1(x)))

        conv2 = self.conv2
        if self.verbose:
            print(f"conv2 shape: {conv2.shape}")
        bn2 = self.batchnorm2
        if self.verbose:
            print(f"bn2 shape: {bn2.shape}")

        residual = self.residual_connection
        if self.verbose:
            print(f"residual shape: {residual.shape}")

        out = bn2(conv2(z1))+residual(x)
        # Don't forget to apply ReLU activation after adding the residual
        out = relu(out)
        return out

class ResNet18(nn.Module):
    def __init__(self, num_classes, verbose=False):

        # TODO: Instantiate all the layers/stages of ResNet-18
        # Hint: don't forget to call super().__init__() first!
        super().__init__()
        self.conv1=nn.Conv2d(in_channels=3,out_channels=64,kernel_size=7,stride=2,padding=3)
        self.batchnorm1=nn.BatchNorm2d(num_features=64)
        self.relu=nn.ReLU()
        self.maxpool1=nn.MaxPool2d(kernel_size=3,stride=2,padding=1)

        self.residualblock1=ResidualBlock(in_channels=64,out_channels=64,kernel_size=3)

        self.residualblock2=ResidualBlock(in_channels=64,out_channels=64,kernel_size=3)

        self.residualblock3=ResidualBlock(in_channels=64,out_channels=128,kernel_size=3,initial_downsample=True)

        self.residualblock4=ResidualBlock(in_channels=128,out_channels=128,kernel_size=3)

        self.residualblock5=ResidualBlock(in_channels=128,out_channels=256,kernel_size=3,initial_downsample=True)

        self.residualblock6=ResidualBlock(in_channels=256,out_channels=256,kernel_size=3)

        self.residualblock7=ResidualBlock(in_channels=256,out_channels=512,kernel_size=3,initial_downsample=True)

        self.residualblock8=ResidualBlock(in_channels=512,out_channels=512,kernel_size=3)

        self.avgpool=nn.AdaptiveAvgPool2d(output_size=(1,1))

        self.flatten=nn.Flatten()

        self.fc=nn.Linear(in_features=512,out_features=num_classes)
        
        self.model=nn.Sequential(self.conv1,self.batchnorm1,self.relu,self.maxpool1,self.residualblock1,self.residualblock2,
        self.residualblock3,self.residualblock4,self.residualblock5,self.residualblock6,
        self.residualblock7,self.residualblock8,self.avgpool,self.flatten,self.fc)
        
    def forward(self, x):

        # TODO: Implement the forward pass of ResNet-18
        out=self.model(x)

        return out

# TODO: Instantiate a ResNet18 model
resnet = ResNet18(num_classes=10)

# TODO: Set the number of epochs
num_epochs = 50

# TODO: Instantiate an optimizer
optimizer = AdamW(resnet.parameters(),lr=0.001)

# TODO: Instantiate a loss criterion
criterion = nn.CrossEntropyLoss()

# TODO: Train the model

resnet_train_accuracies, resnet_train_losses, resnet_val_accuracies, resnet_val_losses = train(
    model=resnet,
    optimizer=optimizer,
    criterion=criterion,
    num_epochs=num_epochs,
    train_dataloader=train_dataloader,
    val_dataloader=val_dataloader
)

import math
def scaled_dot_product_attention(Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor, mask=None, verbose=False):
    """
    Q: matrix of shape (B, N, d_k)
    K: matrix of shape (B, N, d_k)
    V: matrix of shape (B, N, d_v)
    mask: boolean matrix of shape (B, N, N). Values where mask is True will be masked out

    B is the batch size
    N is the number of query/key/value vectors
    d_k is the dimensions of the query/key vectors
    d_v is the dimensions of the value vectors
    """
    assert Q.shape[-1] == K.shape[-1], "Q and K must have same d_k dimension"
    assert K.shape[-2] == V.shape[-2], "K and V must have same sequence length"


    # TODO: Implement the attention function
    return torch.matmul((softmax(torch.matmul(Q,K.transpose(-2,-1))/math.sqrt(Q.shape[-1]))),V)

class AttentionHead(nn.Module):
    def __init__(self, d_model=512, d_k=64, verbose=False):

        # TODO: Implement the __init__ method
        # Don't forget to call super().__init__()!
        super().__init__()
        self.W_q=nn.Linear(d_model,d_k)
        self.W_k=nn.Linear(d_model,d_k)
        self.W_v=nn.Linear(d_model,d_k)
        self.verbose=verbose


    def forward(self, x, encoder_output=None, mask=None):
        """
        x is the input to use for the queries, keys, and values
        encoder_output is the output from the encoder (used for cross-attention)
        mask is the mask to use for the attention
        """


        # TODO: Implement the forward method
        Q = self.W_q(x)
        K = self.W_k(x)
        V = self.W_v(x)

        if encoder_output is not None:
            Q = self.W_q(x)
            K = self.W_k(encoder_output)
            V = self.W_v(encoder_output)

        out = scaled_dot_product_attention(Q,K,V,mask=mask,verbose=self.verbose)

        return out

class MultiHeadAttention(nn.Module):
    def __init__(
        self,
        num_heads=8,
        d_model=512,  # dimension of the embeddings
        verbose=False,
    ):
        """
        Args:
            num_heads: number of attention heads
            d_model: dimension of the embeddings
            verbose: whether to print debug information
        """

        # TODO: Implement the __init__ method
        super().__init__()
        self.d_k = d_model//num_heads
        self.heads = nn.ModuleList([AttentionHead(d_model=d_model,d_k=self.d_k,verbose=verbose) for _ in range(num_heads)])
        self.W_out = nn.Linear(d_model,d_model)

    def forward(self, x, encoder_output=None, mask=None):

        # TODO: Implement the forward method
        out = torch.cat([head(x,encoder_output,mask)for head in self.heads],dim=-1)
        out = self.W_out(out)

        return out

class EncoderLayer(nn.Module):
    def __init__(self, d_model=512, num_heads=8, verbose=False):

        # TODO: Implement the __init__ method
        super().__init__()
        self.multiheadlayer=MultiHeadAttention(num_heads=num_heads,d_model=d_model,verbose=verbose)
        self.layer_norm1=nn.LayerNorm(d_model)
        self.ffn=nn.Sequential(
            nn.Linear(d_model,d_model),
            nn.ReLU(),
            
        )

    def forward(self, x, mask=None):

        # TODO: Implement the forward method
        # 1. Self-attention
        out = self.multiheadlayer(x,mask=mask)
        # 2. Residual connection
        out = x + out
        # 3. LayerNorm
        out = self.layer_norm1(out)
        norm1_out = out
        # 4. Feedforward network
        out=self.ffn(out)
        # 5. Residual connection
        out=out+norm1_out
        # 6. LayerNorm
        out=self.layer_norm1(out)
        # 7. Return the output
        return out

class DecoderLayer(nn.Module):
    def __init__(self, d_model=512, num_heads=8, verbose=False):

        # TODO: Implement the __init__ method
        super().__init__()
        self.masked_multiheadlayer1=MultiHeadAttention(num_heads=num_heads,d_model=d_model,verbose=verbose)
        self.layernorm1=nn.LayerNorm(d_model)
        self.masked_multiheadlayer2=MultiHeadAttention(num_heads=num_heads,d_model=d_model,verbose=verbose)
        self.ffn=nn.Sequential(
            nn.Linear(d_model,4*d_model),
            nn.ReLU(),
            nn.Linear(4*d_model,d_model)
        )
        self.layernorm2=nn.LayerNorm(d_model)

    def forward(self, x, encoder_output):

        # TODO: Implement the forward method
        mask = torch.triu(torch.ones(x.shape[1], x.shape[1], dtype=torch.bool, device=x.device), diagonal=1)
        # Create a look-ahead mask

        # 1. Masked self-attention
        out = self.masked_multiheadlayer1(x,mask=mask)
        # 2. Residual connection
        out = x + out
        # 3. LayerNorm
        out = self.layernorm1(out)
        norm1 = out
        # 4. Cross-attention
        out = self.masked_multiheadlayer2(out,encoder_output=encoder_output)
        # 5. Residual connection
        out = out + norm1
        # 6. LayerNorm
        out = self.layernorm2(out)
        norm2=out
        # 7. Position-wise feed forward network
        out = self.ffn(out)
        # 8. Residual connection
        out = out + norm2
        # 9. LayerNorm
        out = self.layernorm2(out)
        # 10. Return the final normalized output
        return out
        # 10. Return the final normalized output

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=1024):
        """
        max_len = max number of tokens in the sequence
        """
        super().__init__()
        self.pos = torch.zeros((max_len, d_model))
        numerator = torch.arange(end=max_len, dtype=torch.float32).reshape(-1, 1) # shape (seq_len, 1)
        denominator = torch.pow(10000, torch.arange(start=0, end=d_model, step=2, dtype=torch.float32) / d_model).reshape(1, -1) # shape (1, d_model / 2)

        quotient = numerator / denominator # shape (seq_len, d_model). quotient[i, j] = numerator[i, 0] / denominator[0, j]
        self.pos[:, 0::2] = torch.sin(quotient) # assign result of trig to all even indices for every token positions
        self.pos[:, 1::2] = torch.cos(quotient) # assign result of trig to all odd indices for every token position

    def forward(self, x):
        """
        Input x is of shape (B, seq_len, d_embedding)
        """
        out = x + self.pos[:x.shape[1], :].to(device=x.device)
        return out

class TransformerEncoder(nn.Module):
    def __init__(
        self, vocab_size=1024, d_model=512, num_layers=6, num_heads=8, verbose=False
    ):

        # TODO: Implement the __init__ method
        super().__init__()
        self.embedding=nn.Embedding(num_embeddings=vocab_size,embedding_dim=d_model)
        self.positional_encoding=PositionalEncoding(d_model=d_model)
        self.encoder_layers=nn.ModuleList([EncoderLayer(d_model=d_model,num_heads=num_heads,verbose=verbose) for _ in range(num_layers)])
        self.num_layers=num_layers
        self.num_heads=num_heads

    def forward(self, x):

        # 1. Embed the input sequence of tokens
        x = self.embedding(x)
        # 2. Add positional encodings
        x = self.positional_encoding(x)
        # 3. Sequentially pass the result through each EncoderLayer
        for layer in self.encoder_layers:
            x = layer(x)
        # 4. Return the final tensor
        return x

class TransformerDecoder(nn.Module):
    def __init__(self, vocab_size=1024, d_model=512, num_layers=6, num_heads=8, verbose=False):

        # TODO: Implement the __init__ method
        super().__init__()
        self.embedding=nn.Embedding(num_embeddings=vocab_size,embedding_dim=d_model)
        self.positional_encoding=PositionalEncoding(d_model=d_model)
        self.num_layers=num_layers
        self.decoder_layers=nn.ModuleList([DecoderLayer(d_model=d_model,num_heads=num_heads,verbose=verbose) for _ in range(num_layers)])
        self.linear=nn.Linear(d_model,vocab_size)

    def forward(self, x, encoder_output):

        # 1. Embed the input sequence of tokens
        x = self.embedding(x)

        # 2. Add positional encodings
        x = self.positional_encoding(x)

        # 3. Sequentially pass the result through each DecoderLayer
        for layer in self.decoder_layers:
            x = layer(x,encoder_output)

        # 4. Apply the linear layer to predict the next token in our vocabulary
        x = self.linear(x)

        # 5. Return the final tensor
        return x

class Transformer(nn.Module):
    def __init__(
        self,
        vocab_size=1024,
        d_model=512,
        num_encoder_layers=6,
        num_decoder_layers=6,
        num_heads=8,
        verbose=False,
    ):

        # TODO: Implement the __init__ method
        super().__init__()
        self.encoder=TransformerEncoder(vocab_size=vocab_size,d_model=d_model,num_layers=num_encoder_layers,num_heads=num_heads,verbose=verbose)
        self.decoder=TransformerDecoder(vocab_size=vocab_size,d_model=d_model,num_layers=num_decoder_layers,num_heads=num_heads,verbose=verbose)

    def forward(self, x):

        # TODO: Implement the forward method
        # 1. Pass `x` through the encoder
        encoder_output = self.encoder(x)
        # 2. Pass `x` and the encoder's output through the decoder
        x = self.decoder(x,encoder_output)
        # 3. Return the output of the decoder
        return x

def create_sequences(stories, word_to_id, max_seq_length=20, pad_token=PAD_TOKEN, start_token=START_TOKEN, end_token=END_TOKEN):
    """
    Create input/target sequences from paragraphs

    Args:
        stories (list): List of stories to create sequences from
        word_to_id (dict): Dictionary mapping words to their token IDs
        max_seq_length (int): Maximum sequence length to create
        pad_token (str): Token for the pad token
        start_token (str): Token for the start of a sequence
        end_token (str): Token for the end of a sequence

    Returns:
        inputs (tensor): Tensor of input sequences
        targets (tensor): Tensor of target sequences
    """
    inputs = []
    targets = []

    # List of paragraphs
    paragraphs = []

    # TODO: Split the stories into paragraphs and append them to the `paragraphs` list
    for story in stories:
        story_paragraphs = story.split("\n\n")
        paragraphs.extend(story_paragraphs)


    # TODO: Iterate over each paragraph, extract full words, add `<start>` and `<end>` tokens, convert to IDs, and build sequences
    # Make sure to skip paragraphs that are too short to create any sequences
    for paragraph in paragraphs:
        full_words = extract_full_words(paragraph.lower())
        full_words.insert(0,start_token)
        full_words.append(end_token)
        if len(full_words) < max_seq_length:
            continue

        for i in range(len(full_words)-max_seq_length):
            input_seq = full_words[i:i+max_seq_length]
            target_seq = full_words[i+1:i+max_seq_length+1]

            input_seq = [word_to_id.get(word,0) for word in input_seq]
            target_seq = [word_to_id.get(word,0) for word in target_seq]
            inputs.append(input_seq)
            targets.append(target_seq)
            
            

    # TODO: Convert the inputs and targets to tensors of type `torch.long`
    inputs = torch.tensor(inputs,dtype=torch.long)
    targets = torch.tensor(targets,dtype=torch.long)
    return inputs, targets

# Generate sequences
max_seq_length = 20
inputs, targets = create_sequences(tiny_stories_text, word_to_id, max_seq_length)

print(f"Total sequences created: {len(inputs)}")
print(f"Each input sequence length: {len(inputs[0])}")
print(f"Each target sequence length: {len(targets[0])}")

# TODO: Split into training and validation datasets (80/20 split)
train_inputs = inputs[:2000]
train_targets = targets[:2000]
val_inputs = inputs[2000:2200]
val_targets = targets[2000:2200]

print(f"Number of training sequences: {len(train_inputs)}")
print(f"Number of validation sequences: {len(val_inputs)}")

# TODO: Create TensorDatasets
train_dataset = TensorDataset(train_inputs,train_targets)
val_dataset = TensorDataset(val_inputs,val_targets)

# TODO: Create DataLoaders
batch_size = 32
train_loader = DataLoader(train_dataset,batch_size=batch_size,shuffle=True)
val_loader = DataLoader(val_dataset,batch_size=batch_size,shuffle=True)

print(f"Training batches: {len(train_loader)}")
print(f"Validation batches: {len(val_loader)}")

batch = next(iter(train_loader)) # Get the first batch from the training DataLoader

train_input, train_target = batch # Unpack the batch into input and target tensors
print(f"First 10 tokens of training input in first batch: {train_input[0][:10]}")
print(f"First 10 tokens of training target in first batch: {train_target[0][:10]}")

# Transformer training loop
def train_transformer(
    model,
    device,
    optimizer,
    criterion,
    num_epochs,
    train_dataloader,
    val_dataloader,
    vocab,
):
    # === SETUP ===
    # Move model to device
    model.to(device)

    # Lists to store metrics across epochs
    train_losses = []
    val_losses = []

    # === EPOCH LOOP ===
    for epoch in tqdm(range(num_epochs)):
        # === TRAINING PHASE ===
        model.train()  # Set model to training mode

        # Initialize metrics for this epoch
        total_train_loss = 0.0

        # === INNER LOOP (iterate over training batches) ===
        for batch_inputs, batch_targets in tqdm(train_dataloader, desc="Training"):
            # Move inputs and targets to device
            batch_inputs = batch_inputs.to(device)
            batch_targets = batch_targets.to(device)

            # Reset the gradients
            optimizer.zero_grad()

            # Forward pass
            logits = model(batch_inputs)  # [batch_size, seq_len, vocab_size]

            # Reshape for loss calculation
            loss = criterion(logits.reshape(-1, len(vocab)), batch_targets.reshape(-1))
            # Compute loss
            loss.backward()

            # Optional: Add gradient clipping to prevent exploding gradients
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

            # Update the model parameters
            optimizer.step()

            # Accumulate the loss
            total_train_loss += loss.item()

        # === END OF INNER LOOP ===

        # Compute the average loss for this epoch
        avg_train_loss = total_train_loss / len(train_dataloader)
        train_losses.append(avg_train_loss)

        # === END OF TRAINING PHASE ===

        # === VALIDATION PHASE ===
        model.eval()  # Set model to evaluation mode
        total_val_loss = 0.0

        with torch.no_grad():
            # === INNER LOOP (iterate over validation batches) ===
            for batch_inputs, batch_targets in tqdm(val_dataloader, desc="Validation"):
                # Move inputs and targets to device
                batch_inputs = batch_inputs.to(device)
                batch_targets = batch_targets.to(device)

                # Forward pass
                logits = model(batch_inputs)

                # Compute loss
                loss = criterion(
                    logits.reshape(-1, len(vocab)), batch_targets.reshape(-1)
                )

                # Accumulate the loss
                total_val_loss += loss.item()
        # === END OF INNER LOOP ===

        # Compute the average loss for this epoch
        avg_val_loss = total_val_loss / len(val_dataloader)
        val_losses.append(avg_val_loss)

        # === END OF VALIDATION PHASE ===

        # Print epoch results
        print(f"\nEpoch {epoch + 1} Results:")
        print(f"  Training Loss: {avg_train_loss:.4f}")
        print(f"  Validation Loss: {avg_val_loss:.4f}")

    # === END OF EPOCH LOOP ===

    return train_losses, val_losses

# TODO: Instantiate a Transformer model

vocab_size=len(vocab)

transformer = Transformer(vocab_size=vocab_size,
                                    d_model=512,
                                    num_encoder_layers=6,
                                    num_decoder_layers=6,
                                    num_heads=8,
                                    verbose=False).to(device)

# TODO: Set the optimizer
optimizer = AdamW(transformer.parameters(),lr=0.0001)

# TODO: Set the loss criterion
criterion = nn.CrossEntropyLoss()

# TODO: Set the number of epochs
num_epochs= 50

# TODO: Train the model
transformer_train_losses, transformer_val_losses = train_transformer(
    model=transformer,
    optimizer=optimizer,
    criterion=criterion,
    num_epochs=num_epochs,
    train_dataloader=train_loader,
    val_dataloader=val_loader,
    device=device,
    vocab=vocab
)

# TODO: Plot the training and validation loss across the epochs

def plot_metrics_transformer(train_losses, val_losses, num_epochs=None, title=""):
    """
    Plots the training loss, training accuracy, validation loss, and validation accuracy.
    Args:
        train_losses: list of training losses
        val_losses: list of validation losses
        train_accuracies: list of training accuracies
        val_accuracies: list of validation accuracies
        num_epochs: number of epochs
        title: title of the plot
    """
    fig, axes = plt.subplots(1, 2, figsize=(12, 8))

    # 训练损失
    axes[0].plot(train_losses, label='Training Loss')
    axes[0].set_title('Training Loss')
    axes[0].set_xlabel('Epoch')
    axes[0].set_ylabel('Loss')


    axes[1].plot(val_losses, label='Validation Loss')
    axes[1].set_title('Validation Loss')
    axes[1].set_xlabel('Epoch')
    axes[1].set_ylabel('Loss')



    plt.tight_layout()
    plt.suptitle(title)
    plt.show()
    
plot_metrics_transformer(train_losses=transformer_train_losses,val_losses=transformer_val_losses,
num_epochs=50,title="Transformer Training Metrics")

print(f"Final Training Loss: {transformer_train_losses[-1]:.4f}")
print(f"Final Validation Loss: {transformer_val_losses[-1]:.4f}")