模块化 Transformers¶

from torch import nn
from ..bert.configuration_bert import BertConfig
from ..bert.modeling_bert import (
    BertModel,
    BertEmbeddings,
    BertForMaskedLM
)

# RoBERTa 配置与 BERT 的配置相同
class RobertaConfig(BertConfig):
  model_type = 'roberta'

# 重新定义嵌入层以突出填充 ID 的差异，并重新定义位置嵌入
class RobertaEmbeddings(BertEmbeddings):
    def __init__(self, config):
        super().__init__(config)

        self.padding_idx = config.pad_token_id
        self.position_embeddings = nn.Embedding(
            config.max_position_embeddings, config.hidden_size, padding_idx=self.padding_idx
        )

# RoBERTa 模型与 BERT 模型相同，除了嵌入层。
# 我们在上面重新定义了嵌入层，因此这里不需要额外的工作
class RobertaModel(BertModel):
  def __init__(self, config):
    super().__init__(config)
    self.embeddings = RobertaEmbeddings(config)

# 头部现在只需要将模型内部重新定义为正确的 `RobertaModel`
class RobertaForMaskedLM(BertForMaskedLM):
  def __init__(self, config):
    super().__init__(config)
    self.model = RobertaModel(config)

class GemmaModel(LlamaModel):
    def __init__(self, config):
        super().__init__(config)
        del self.embed_tokens

class GemmaModel(PreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size

        self.layers = nn.ModuleList(
            [LlamaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
        )
        self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.rotary_emb = LlamaRotaryEmbedding(config=config)
        self.gradient_checkpointing = False

        # 初始化权重并应用最终处理
        self.post_init()

class GemmaTokenizer(LlamaTokenizer):
    # ...

    def get_spm_processor(self):
        raise AttributeError("Gemma 不需要此功能")

    def unk_token_length(self):
        raise AttributeError("Gemma 不需要此功能")

def my_new_function(*args, **kwargs):
  # 在这里做某事
  pass

class GemmaModel(LlamaModel):
    def forward(*args, **kwargs):
      # 调用函数
      example = my_new_function(*args, **kwargs)
      # 继续操作

class GemmaTokenizer(LlamaTokenizer, PretrainedTokenizerFast):
    def __init__(self, eos_token="</s>"):
        eos_token = AddedToken(eos_token)
        PretrainedTokenizerFast.__init__(self, eos_token)

class GemmaModel(nn.Module):
    def __init__(self):
        eos_token = AddedToken(eos_token)
        super().__init__(eos_token)

class GemmaVisionModel(CLIPModel):                                 
    pass