integrations.liger.models.jamba

integrations.liger.models.jamba

Jamba model with LigerFusedLinearCrossEntropyLoss

Functions

Name	Description
lce_forward

lce_forward

integrations.liger.models.jamba.lce_forward(
    self,
    input_ids=None,
    attention_mask=None,
    position_ids=None,
    past_key_values=None,
    inputs_embeds=None,
    labels=None,
    use_cache=None,
    output_router_logits=None,
    logits_to_keep=0,
    **kwargs,
)

Parameters

Name	Type	Description	Default
labels	torch.LongTensor of shape (batch_size, sequence_length), optional	Labels for computing the masked language modeling loss. Indices should either be in [0, ..., config.vocab_size] or -100 (see input_ids docstring). Tokens with indices set to -100 are ignored (masked), the loss is only computed for the tokens with labels in [0, ..., config.vocab_size].	None
logits_to_keep	int or torch.Tensor, optional	Calculate logits for the last logits_to_keep tokens. If 0, calculate logits for all input_ids. Only last token logits are needed for generation, and calculating them only for that token can save memory, which becomes pretty significant for long sequences.	0

Returns:

Example:

>>> from transformers import AutoTokenizer, JambaForCausalLM

>>> model = JambaForCausalLM.from_pretrained("ai21labs/Jamba-v0.1")
>>> tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")

>>> prompt = "Hey, are you conscious? Can you talk to me?"
>>> inputs = tokenizer(prompt, return_tensors="pt")

>>> # Generate
>>> generate_ids = model.generate(inputs.input_ids, max_length=30)
>>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
"Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."