class BertGenerationEmbeddings(nn.Module):
# bert的word_embedding
class BertGenerationPreTrainedModel(PreTrainedModel):
# 一个用来保存和恢复checkpoint的抽象类，继承了transformers内部的独有类PreTrainedModel
class BertGenerationEncoder(BertGenerationPreTrainedModel):
# 最重要的类，可以根据args 变成 encoder or decoder
class BertGenerationOnlyLMHead(nn.Module):
# decoder的分类头，用来做vocab的分类的
class BertGenerationDecoder(BertGenerationPreTrainedModel):
# encoder的行为和decoder的行为不同，用一个decoder类定义decoder的行为

def __init__(self, config):
  super().__init__(config)
  # config是模型的参数类，定义了模型的一些超参
  self.config = config
  # embeddings 使用了自己定义的embeddings，encoder直接使用了bert内部的encoder
  self.embeddings = BertGenerationEmbeddings(config)
  # is_decoder在config被传入BertEncoder
  self.encoder = BertEncoder(config)

  # Initialize weights and apply final processing
  self.post_init()

class BertLayer(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.chunk_size_feed_forward = config.chunk_size_feed_forward
        self.seq_len_dim = 1
        # self-attention
        self.attention = BertAttention(config)
        self.is_decoder = config.is_decoder
        self.add_cross_attention = config.add_cross_attention
        # 增加了cross-attention
        if self.add_cross_attention:
            if not self.is_decoder:
                raise ValueError(f"{self} should be used as a decoder model if cross attention is added")
            # 发现cross-attention和self-attention使用的是同一个attention类说明进行了复用
            self.crossattention = BertAttention(config, position_embedding_type="absolute")
        # layer中的全连接层
        self.intermediate = BertIntermediate(config)
        self.output = BertOutput(config)

if self.is_decoder and encoder_hidden_states is not None:
    if not hasattr(self, "crossattention"):
        raise ValueError(
            f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers by setting `config.add_cross_attention=True`"
        )

    # cross_attn cached key/values tuple is at positions 3,4 of past_key_value tuple
    cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None
    cross_attention_outputs = self.crossattention(
        attention_output, # self-attn的输出做cross-attn的Q输入
        attention_mask,
        head_mask,
        encoder_hidden_states,# 来自encoder的隐藏层做K与V的输入
        encoder_attention_mask,
        cross_attn_past_key_value,
        output_attentions,
    )
    attention_output = cross_attention_outputs[0]
    outputs = outputs + cross_attention_outputs[1:-1]  # add cross attentions if we output attention weights

def forward(
    self,
    input_ids=None,
    attention_mask=None,
    position_ids=None,
    head_mask=None,# 前四个都是tokenizer的输出
    inputs_embeds=None,# 可以直接输入embeding而不是ids，input_ids,inputs_embeds只能输入一个
    encoder_hidden_states=None,# 作为decoder时会使用，Cross-Attn的KV
    encoder_attention_mask=None,# 作为decoder时会使用，用来屏蔽[PAD]
    past_key_values=None,# 作为decoder时会使用，预先算好的所有BertLayer的KV
    use_cache=None,# 作为encoder时，True会保留所有的BertLayer的KV然后供decoder使用也就是上面的arg
    output_attentions=None,
    output_hidden_states=None,
    return_dict=None,# 输出时调整输出格式的一些bool参数，详细可以去看doc
):

embedding_output = self.embeddings(
  input_ids=input_ids,
  position_ids=position_ids,
  inputs_embeds=inputs_embeds,
  past_key_values_length=past_key_values_length,
)

encoder_outputs = self.encoder(
  embedding_output,
  attention_mask=extended_attention_mask,
  head_mask=head_mask,
  encoder_hidden_states=encoder_hidden_states,# 当is_decoder=True时以下3个参数会派上用场
  encoder_attention_mask=encoder_extended_attention_mask,
  past_key_values=past_key_values,
  use_cache=use_cache,
  output_attentions=output_attentions,
  output_hidden_states=output_hidden_states,
  return_dict=return_dict,
)

def __init__(self, config):
  super().__init__(config)

  if not config.is_decoder:
    logger.warning("If you want to use `BertGenerationDecoder` as a standalone, add `is_decoder=True.`")
    # 复用了BertGenerationEncoder，通过config.is_decoder调整BertGenerationEncoder的行为
    self.bert = BertGenerationEncoder(config)
    self.lm_head = BertGenerationOnlyLMHead(config)

outputs = self.bert(
  input_ids,
  attention_mask=attention_mask,
  position_ids=position_ids,
  head_mask=head_mask,
  inputs_embeds=inputs_embeds,
  encoder_hidden_states=encoder_hidden_states,
  encoder_attention_mask=encoder_attention_mask,
  past_key_values=past_key_values,
  use_cache=use_cache,
  output_attentions=output_attentions,
  output_hidden_states=output_hidden_states,
  return_dict=return_dict,
)
# 输入的每个token的last hidden layer
sequence_output = outputs[0]
# hidden layer输入vocab的分类头计算每个token的预测值然后存在logits输出
prediction_scores = self.lm_head(sequence_output)

from transformers import BertGenerationTokenizer, BertGenerationDecoder, BertGenerationConfig
import torch

tokenizer = BertGenerationTokenizer.from_pretrained('google/bert_for_seq_generation_L-24_bbc_encoder')
config = BertGenerationConfig.from_pretrained("google/bert_for_seq_generation_L-24_bbc_encoder")
config.is_decoder = True
model = BertGenerationDecoder.from_pretrained('google/bert_for_seq_generation_L-24_bbc_encoder', config=config)

inputs = tokenizer("How are you", return_token_type_ids=False, return_tensors="pt")
tokenizer.pad_token
outputs = model(**inputs)

prediction_logits = outputs.logits
prediction_logits.shape

>>> normalizer.cc(51) LOG(INFO) precompiled_charsmap is empty. use identity normalization.
>>> Some weights of BertGenerationDecoder were not initialized from the model checkpoint at google/bert_for_seq_generation_L-24_bbc_encoder and are newly initialized: ['lm_head.bias', 'lm_head.decoder.weight', 'lm_head.decoder.bias']
>>>You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
>>> torch.Size([1, 3, 50358])

def log(func):
  def someprocess():
    print("before call")
    func()
  return someprocess

@log
def test():
  print("it's a test func")
  
before call
it's a test func

test = log(test)

from functools import partial

def log(func=None,*,post_message=None):
  def someprocess():
    print("before call")
    func()
    if post_message:
        print(post_message)
  # 如果以log(post_messag=..)调用，此处返回一个固定了post_message的参数的log装饰器   
  if func is None:
      return partial(log,post_message=post_message)
  # 如果以log方式调用则直接返回被装饰的函数
  return someprocess

@log
def test():
  print("it's a test func")

test() 

@log(post_message="after call")
def test2():
  print("it's a test func")

test2()

class SeqNet(nn.Module):
    def __init__(self):
        super().__init__()
        # 注意这里Sequential的参数是多个对象，而ModList的是一个数组
        self.encoder = nn.Sequential(
            nn.Linear(10,128),
            nn.ReLU(),
            nn.Linear(128,50)
        )
    def forward(self,x):
      # Sequential直接输入X则会按照定义的顺序依次运算
        return self.encoder(x)
class ModNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = nn.ModuleList([
            nn.Linear(10,128),
            nn.ReLU(),
            nn.Linear(128,50)
        ])
    def forward(self,x):
      # 如果按照以下方式输入则会报Not Implement error 说明如何运算需要自己定义
      # return self.encoder(x)
      # 需要我们自己定义整个运算过程
        for block in self.encoder:
            x = block(x)
        return x

import torch
a = torch.randn((3,5))
b = torch.randint(0,3,(5,)).bool()
print(a)
print(b)
a[:,b]

tensor([[ 1.2588,  0.1620, -0.3095,  2.0669, -0.3158],
        [-1.6065, -0.7930, -0.5658,  1.8601, -0.2592],
        [ 1.3226,  0.3040, -1.0924,  0.7583, -1.5444]])
tensor([ True, False, False,  True,  True])
tensor([[ 1.2588,  2.0669, -0.3158],
        [-1.6065,  1.8601, -0.2592],
        [ 1.3226,  0.7583, -1.5444]])

import torch
a = torch.randn((3,5))
b = torch.randint(0,3,(5,))
print(a)
print(b)
a[:,b]

tensor([[-1.0592,  0.9445, -0.6265,  0.2607, -0.2166],
        [ 1.8283,  0.9678,  0.6175,  2.0904,  0.2356],
        [ 0.1864,  1.0110,  0.0425, -1.3611,  1.1043]])
tensor([1, 1, 1, 2, 1])
tensor([[ 0.9445,  0.9445,  0.9445, -0.6265,  0.9445],
        [ 0.9678,  0.9678,  0.9678,  0.6175,  0.9678],
        [ 1.0110,  1.0110,  1.0110,  0.0425,  1.0110]])