MCPcopy Create free account
hub / github.com/THUDM/LongWriter / LlamaModel

Class LlamaModel

train/patch/modeling_llama.py:915–1155  ·  view source on GitHub ↗

Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`LlamaDecoderLayer`] Args: config: LlamaConfig

Source from the content-addressed store, hash-verified

913 LLAMA_START_DOCSTRING,
914)
915class LlamaModel(LlamaPreTrainedModel):
916 """
917 Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`LlamaDecoderLayer`]
918
919 Args:
920 config: LlamaConfig
921 """
922
923 def __init__(self, config: LlamaConfig):
924 super().__init__(config)
925 self.padding_idx = config.pad_token_id
926 self.vocab_size = config.vocab_size
927
928 self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
929 self.layers = nn.ModuleList(
930 [LlamaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
931 )
932 self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
933 self.rotary_emb = LlamaRotaryEmbedding(config=config)
934 self.gradient_checkpointing = False
935
936 # Initialize weights and apply final processing
937 self.post_init()
938
939 def get_input_embeddings(self):
940 return self.embed_tokens
941
942 def set_input_embeddings(self, value):
943 self.embed_tokens = value
944
945 @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING)
946 def forward(
947 self,
948 input_ids: torch.LongTensor = None,
949 attention_mask: Optional[torch.Tensor] = None,
950 position_ids: Optional[torch.LongTensor] = None,
951 past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
952 inputs_embeds: Optional[torch.FloatTensor] = None,
953 use_cache: Optional[bool] = None,
954 output_attentions: Optional[bool] = None,
955 output_hidden_states: Optional[bool] = None,
956 return_dict: Optional[bool] = None,
957 cache_position: Optional[torch.LongTensor] = None,
958 ) -> Union[Tuple, BaseModelOutputWithPast]:
959 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
960 output_hidden_states = (
961 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
962 )
963 use_cache = use_cache if use_cache is not None else self.config.use_cache
964 return_dict = return_dict if return_dict is not None else self.config.use_return_dict
965
966 if (input_ids is None) ^ (inputs_embeds is not None):
967 raise ValueError(
968 "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
969 )
970
971 if self.gradient_checkpointing and self.training and use_cache:
972 logger.warning_once(

Callers 4

__init__Method · 0.85
__init__Method · 0.85
__init__Method · 0.85
__init__Method · 0.85

Calls

no outgoing calls

Tested by

no test coverage detected