算是HAPPY-LLM的作业吧，不过代码跟原文基本没关系。
自己对着图和公式实现的，有些知道原理但不熟悉torch，所以不会实现的部分问了AI。
即使前面学原理比较认真，手搓过程也遇到很多困难。到现在关于代码也有一些发散的疑惑没有解决，不过我决定暂时略过了，或许在以后的实践中会有所感悟。

'''
代码实现顺序跟随输入到输出顺序
显而易见的废话很多，费曼学习法实践+仅供参考
'''
import torch
from torch import nn
import numpy as np
'''
First
首先是输入部分inputs，对于原始语料，如“我爱你”，模型外的处理过程如下：
1.分词，分为“我，爱，你”
2.根据词表转化为id。比如我这儿有一个词表[1：你，2：我，3：爱]
  那这个序列就变成[2,3,1]
  Input的一般维度是这样(batch_size, len)
  就是说有batch_size个序列，每个序列是类似[2,3,1]的内容
'''

'''
Second
下面进行input embedding的实现，可以直接用torch中的nn.embedding函数
一般这个函数传入两个参数vocab和d_model,分别是词表大小和词嵌入维度
    embedding_input = nn.Embedding(vocab, d_model)
词表大小顾名思义，可以理解为字典词汇量的多少
词嵌入维度就是把每个id/id对应的token，转化为有d_model个维度的向量。
论文中d_model = 512
这一步之后输出的维度(batch_size, len, d_model)
'''

'''
Third
然后实现positional encoding
这一步简单来讲就是位置编码后，加在上一步转化完毕的向量上。
假设"我爱你"对应的embedding结果是[2,3,1] -> [[0,0,1],[0,1,0],[0,1,1]]
假设"我爱你"对应的PE结果是[[0,0,1.1],[0,0,1.2],[0,0,1.3]]
位置编码后的向量就是上面两个对应相加。所以位置编码信息的维度和上面的d_model是相等的
这一步输出的维度(batch_size, len, d_model)

下面具体实现代码
'''
class PositionalEncoding(nn.Model):
    def __init__(self, d_model, dropout, seq_len):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)
        pe = torch.zeros(seq_len, d_model) 

        # 直观计算位置编码
        # for pos in torch.arange(seq_len):
        #     for i in torch.arange(int(d_model/2)):
        #         fenmu = torch.power(10000, 2*i/d_model)
        #         pe[pos,2*i] = torch.sin(pos/fenmu)
        #         pe[pos,2*i+1] = torch.cos(pos/fenmu)        
        
        # 计算位置编码
        max_len = seq_len
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-torch.log(10000.0) / d_model)).unsqueeze(0)
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)

        pe = pe.unsqueeze(0) #对pe进行升维，因为要加在词向量上(batch, d_model),它缺batch
        self.register_buffer('pe', pe) #将张量 pe 注册为模型的“缓冲区”（buffer）​，使其成为模型的一部分，但不会被当作可训练参数

    def forward(self, x):
        _, seq_len, _ = x.shape
        return x + self.pe[:,:seq_len,:]
    
'''
Fourth
然后就到了实现encoder的部分。
图中N=6，左边每个方框里是一个encoder layer，6个连一起就是encoder
FFN可以用torch自带的
还需要实现多头注意力
先实现注意力，再多头，最后encoder
'''       
def attention(query, key, value, mask=None):
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2. -1)) / np.sqrt(d_k)

    if mask is not None:
        scores = scores.masked_fill(mask, -1e9)

    p_attn = torch.softmax(scores, dim=-1)

    return torch.matmul(p_attn, value)


class multihead(nn.Module):
    '''
    将查询(Q)、键(K)、值(V)通过不同的线性变换投影到多个子空间
    在每个子空间中计算缩放点积注意力
    将所有头的输出拼接并通过线性变换得到最终输出
    '''
    def __init__(self,d_model=512,num_heads=8,dropout=0.1):
      super().__init__()
      assert d_model % num_heads == 0, "模型维度必须能被注意力头数整除"
      self.d_model = d_model
      self.n_heads = num_heads
      self.d_head = d_model // num_heads

      # 线性变换层，用于Q,K,V和最终输出
      self.w_q = nn.Linear(d_model, d_model)  # 查询变换
      self.w_k = nn.Linear(d_model, d_model)  # 键变换
      self.w_v = nn.Linear(d_model, d_model)  # 值变换
      self.w_o = nn.Linear(d_model, d_model)  # 输出变换

      self.dropout = nn.Dropout(dropout)

    def forward(self, q, k, v, mask = None):
        batch_size = q.size(0)

        # 投射q,k,v。自注意力的话，输入的q，k，v是相同的序列
        # (batch_size, seq_len, d_model) -> (batch_size, seq_len, num_heads, d_k) -> (batch_size, num_heads, seq_len, d_k)
        q = self.w_q(q).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        k = self.w_k(k).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        v = self.w_v(v).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)

        # 计算attention
        output = attention(q,k,v,mask)
        # 前面交换1，2是方便矩阵乘法，乘完要换回来。
        # transpose后内存不连续，无法用.view，用contiguous()使连续
        # 用view合并多头为(batch_size, seq_len, d_model)，也就是原始的输入形状
        output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
        output = self.w_o(output)
        return output


class encoder_layer(nn.Module):
    def __init__(self, d_model=512, n_heads=8, dim_feedforward=2048, dropout=0.1):
        super().__init__()
        # MULTIHEADS
        self.self_attn = multihead(d_model, n_heads, dropout)

        # FFN
        self.feed_forward = nn.Sequential(
            nn.Linear(d_model, dim_feedforward),
            nn.ReLU(),
            nn.Linear(dim_feedforward, d_model)
        )
        # NORM
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        
        # Dropout层
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

        # ADD直接对输出加输入就能实现

    def forward(self, x, mask = None):
        attn_op = self.self_attn(x,x,x,mask)
        x = x + attn_op #add
        x = self.norm1(x)

        FFN_op = self.feed_forward(x)
        x = x + FFN_op
        x= self.norm2(x)

        return x


class encoder(nn.Module):
    def __init__(self, n_layers = 6, d_model=512, n_heads=8, dim_feedforward=2048, dropout=0.1, max_len=5000):
        super().__init__()
        self.d_model = d_model

        # 位置编码
        self.PE =  PositionalEncoding(d_model, dropout, max_len)

        # 堆叠多个编码器层
        self.layers = nn.ModuleList([
            encoder_layer(d_model, n_heads, dim_feedforward, dropout)
            for _ in range(n_layers)
        ])

        self.norm = nn.LayerNorm(d_model)

    def forward(self, x):
        src = self.PE(x)
        for layer in self.layers:
          src = layer(src)
        src = self.norm(src)
        return src

'''
Fifth
接下来实现decoder
要算一下mask，填入上面实现的attention，就是masked多头
'''
# 要用的东西基本上已经全部实现了
# 再写一个生成mask的就行
class decoder_layer(nn.Module):
    def __init__(self, n_layers = 6, d_model=512, n_heads=8, dim_feedforward=2048, dropout=0.1, max_len=5000):
        super().__init__()
        self.attn1 = multihead()
        self.attn2 = multihead()
        self.feed_forward = nn.Sequential(
            nn.Linear(d_model, dim_feedforward),
            nn.ReLU(),
            nn.Linear(dim_feedforward, d_model)
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)

    def forward(self, tgt, memory):
        """
        参数:
            tgt: 目标序列 (batch_size, tgt_len, d_model)
            memory: 编码器输出 (batch_size, src_len, d_model)
            tgt_mask: 目标序列掩码 (防止看到未来token)
            memory_mask: 编码器输出掩码 (如padding mask)
            
        返回:
            解码后的序列 (batch_size, tgt_len, d_model)
        """
        _,tgt_len,_ = tgt.size()
        # 创建下三角掩码（防止看到未来token）
        tgt_mask = torch.tril(torch.ones(tgt_len, tgt_len)).bool()

        # 第一子层：多头自注意力（带目标序列掩码）
        attn1_output = self.attn1(tgt, tgt, tgt, tgt_mask)
        tgt = tgt + self.dropout1(attn1_output)
        tgt = self.norm1(tgt)
        
        # 第二子层：多头交叉注意力
        attn2_output = self.attn2(tgt, memory, memory)
        tgt = tgt + self.dropout2(attn2_output)
        tgt = self.norm2(tgt)
        
        # 第三子层：前馈网络
        ff_output = self.feed_forward(tgt)
        tgt = tgt + self.dropout3(ff_output)
        tgt = self.norm3(tgt)
        
        return tgt
        
    

class decoder(nn.Module):
    def __init__(self, n_layers = 6, d_model=512, n_heads=8, dim_feedforward=2048, dropout=0.1, max_len=5000):
        super().__init__()
        # 位置编码
        self.PE =  PositionalEncoding(d_model, dropout, max_len)
        self.layers = nn.ModuleList([
            decoder_layer()
            for _ in range(n_layers)
        ])
        self.norm = nn.LayerNorm()


    def forward(self,tgt,memory):
        tgt = self.PE(tgt)
        for layer in self.layers:
            tgt = layer(tgt,memory)
        tgt = self.norm(tgt)
        return tgt  


'''
sixth
把前面实现的所有部分连起来，在forward里过一个linear
计算loss的时候会自动softmax
'''
class transformer(nn.Module):
    def __init__(self, src_vocab_size, tgt_vocab_size, num_layers=6, d_model=512,
                 num_heads=8, dim_feedforward=2048, dropout=0.1, max_len=5000):
        super().__init__()
        self.srcEmbedding = nn.Embedding(src_vocab_size,d_model)
        self.tgtEmbedding = nn.Embedding(tgt_vocab_size,d_model)
        self.encoder = encoder(num_layers,d_model,num_heads,dim_feedforward,dropout,max_len)
        self.decoder = decoder(num_layers,d_model,num_heads,dim_feedforward,dropout,max_len)
        self.output_proj = nn.Linear(d_model, tgt_vocab_size)
        self._reset_parameters()

    def _reset_parameters(self):
      for p in self.parameters():
          if p.dim() > 1:
              nn.init.xavier_uniform_(p)

    def forward(self,src,tgt):
        src = self.srcEmbedding(src)
        tgt = self.tgtEmbedding(tgt)
        memory = self.encoder(src)
        tgt = self.decoder(tgt,memory)
        out = self.output_proj(tgt)
        
        # 还要算loss，明天再写