mindformers.models.llama.LlamaForCausalLM

class mindformers.models.llama.LlamaForCausalLM(config: Optional[LlamaConfig] = None)

Provide llama training loss or logits through network. Args:

config (LlamaConfig): The config of llama model.

Inputs:

input_ids(Tensor): the tokenized inputs with datatype int32, Tensor of shape \((batch, seq\_length)\). labels(Tensor): the tokenized labels with datatype int32, Tensor of shape \((batch, seq\_length)\). input_position(Tensor): current position, used by model.predict. position_ids(Tensor): Reserved param, not used. attention_mask(Tensor): Reserved param, not used. input_embeds(Tensor): Reserved param, not used. init_reset(bool, optional): A bool tensor with shape [1], used to clear the past key parameter and

past value parameter used in the incremental prediction. Default True.

batch_valid_length(Tensor): the past calculated the index with datatype int32, used for incremental

prediction. Tensor of shape \((batch_size,)\). Default None.

Returns:

Tensor, the loss or logits of the network.

Examples:

>>> from mindformers.models.llama import LlamaConfig, LlamaForCausalLM
>>> config = LlamaConfig(batch_size=2)
>>> network = LlamaForCausalLM(config=config)