有没有简短一点的Transformer模型源代码？

以下是一个简短的Transformer模型源代码：

import torch
import torch.nn as nn

class Transformer(nn.Module):
    def __init__(self, input_dim, output_dim, num_layers, num_heads, hidden_dim, dropout):
        super().__init__()
        self.encoder_layer = nn.TransformerEncoderLayer(input_dim, num_heads, hidden_dim, dropout)
        self.encoder = nn.TransformerEncoder(self.encoder_layer, num_layers)
        self.fc = nn.Linear(input_dim, output_dim)
        
    def forward(self, x):
        x = self.encoder(x)
        x = x.mean(dim=0)
        x = self.fc(x)
        return x

这个模型有一个编码器，它由多个Transformer编码器层组成。编码器的输出通过平均池化后，经过一个全连接层，得到最终的输出。这个模型的参数包括输入维度、输出维度、编码器层数、注意力头数、隐藏层维度和dropout概率