Method __init__

tensorrt_llm/models/chatglm/model.py:177–205

(self, config: ChatGLMConfig)

class ChatGLMModel(Module):

    def __init__(self, config: ChatGLMConfig):
        super().__init__()
        self.chatglm_version = config.chatglm_version
        norm_cls = RmsNorm if config.rmsnorm else LayerNorm

        self.vocab_embedding = Embedding(config.vocab_size,
                                         config.hidden_size,
                                         dtype=config.dtype)

        if config.chatglm_version == 'glm':
            self.position_embedding = Embedding(
                config.max_position_embeddings + 1,
                config.hidden_size,
                dtype=config.dtype,
            )
            self.block_embedding = Embedding(
                config.max_position_embeddings + 1,
                config.hidden_size,
                dtype=config.dtype,
            )

        self.layers = DecoderLayerList(ChatGLMDecoderLayer, config)

        self.ln_f = norm_cls(
            normalized_shape=config.hidden_size,
            eps=config.norm_epsilon,
            elementwise_affine=True,
            dtype=config.dtype,
        )

    def forward(
        self,