import torch
import torch.nn as nn

from .TransformersBlocks import TransformersBlocks
from .LayerNormal import LayerNormal

class GPTModel(nn.Module):
    def __init__(self,config):
        super().__init__()
        # 给输入批次添加词表维度的嵌入 查表操作
        self.tok_emb = nn.Embedding(config['vocab_size'],config['emb_dim'])
        # 给输入的批次添加对应的位置编码
        self.pos_emb = nn.Embedding(config['context_length'],config['emb_dim'])
        self.drop = nn.Dropout(config['drop_rate'])

        # 添加transformers blocks(按层数堆叠TransformerBlocks)
        # 注意添加range()
        self.transformer_blocks = nn.Sequential(*[
            TransformersBlocks(config)
            for _ in range(config['n_layers'])
        ])
        # 最终归一化
        self.final_norm = LayerNormal(config['emb_dim'])
        # 获取logits 讲emb_dim 转换为 vocab_size
        self.out_head = nn.Linear(
            config['emb_dim'],config['vocab_size'],
            bias = False
        )
    
    def forward(self,x):
        # (batch_size, seq_len)
        batch,num_token = x.shape
        # 加入嵌入向量维度 (batch_size, seq_len, emb_dim)
        tok_embs = self.tok_emb(x)
        # 加入位置编码 ( seq_len, emb_dim) 位置编码只使用一份
        pos_embs = self.pos_emb(
            # 每个编号代表一个位置
            # torch.arange() 默认会创建在 CPU 上，所以需要进行一致
            torch.arange(num_token,device=x.device)
        )
        x = tok_embs + pos_embs
        # drop
        x = self.drop(x)
        # 叠加transformers
        x = self.transformer_blocks(x)
        # 最终归一化
        x = self.final_norm(x)
        # 获取logits
        x = self.out_head(x)
        return x