easydel.modules.llava.__init__

class easydel.modules.llava.__init__.LlavaConfig(vision_config=None, text_config=None, image_token_index=32000, projector_hidden_act='gelu', vision_feature_select_strategy='default', vision_feature_layer=-2, image_seq_length=576, multimodal_projector_bias=True, **kwargs)

Bases: EasyDeLBaseConfig

This is the configuration class to store the configuration of a [LlavaForConditionalGeneration]. It is used to instantiate an Llava model according to the specified arguments, defining the model architecture. Instantiating a configuration with the defaults will yield a similar configuration to that of the Llava-9B.

e.g. [llava-hf/llava-9b](https://huggingface.co/llava-hf/llava-9b)

Configuration objects inherit from [PretrainedConfig] and can be used to control the model outputs. Read the documentation from [PretrainedConfig] for more information.

Parameters

• vision_config (Union[AutoConfig, dict], optional, defaults to CLIPVisionConfig) – The config object or dictionary of the vision backbone.

• text_config (Union[AutoConfig, dict], optional, defaults to LlamaConfig) – The config object or dictionary of the text backbone.

• image_token_index (int, optional, defaults to 32000) – The image token index to encode the image prompt.

• projector_hidden_act (str, optional, defaults to “gelu”) – The activation function used by the multimodal projector.

• vision_feature_select_strategy (str, optional, defaults to “default”) – The feature selection strategy used to select the vision feature from the vision backbone. Can be one of “default” or “full”.

• vision_feature_layer (Union[int, List[int]], optional, defaults to -2) – The index of the layer to select the vision feature. If multiple indices are provided, the vision feature of the corresponding indices will be concatenated to form the vision features.

• image_seq_length (int, optional, defaults to 576) – Sequence length of one image embedding.

• multimodal_projector_bias (bool, optional, defaults to True) – Whether to use bias in the multimodal projector.

get_partition_rules(*args, **kwargs)

Get the partition rules for distributed training by combining the partition rules from both the text and vision configurations.

This method retrieves the partition rules from the text_config and vision_config components and combines them to create a comprehensive set of rules for the entire multimodal model.

Parameters

• *args – Variable length argument list to be passed to the text and vision configs.

• **kwargs – Arbitrary keyword arguments to be passed to the text and vision configs.

Returns

A combined tuple of partition rules from both text and vision configurations.

Return type

tuple

model_type: str = 'llava'

sub_configs: dict[str, 'PretrainedConfig'] = {'text_config': <class 'easydel.modules.auto.auto_configuration.AutoEasyDeLConfig'>, 'vision_config': <class 'easydel.modules.auto.auto_configuration.AutoEasyDeLConfig'>}

class easydel.modules.llava.__init__.LlavaForConditionalGeneration(*args: Any, **kwargs: Any)

Bases: EasyDeLBaseModule

get_image_features(pixel_values: Union[Array, ndarray, bool, number]) → Union[Array, ndarray, bool, number]

init_cache(batch_size, max_length, starts=None, shardings=None, pad_token_id=None)

Initializes and returns a standard (non-paged) Key-Value cache.

This method first creates the necessary metadata using create_cache_metadata and then calls TransformerCache.init_cache to allocate and initialize the cache tensors based on the model’s configuration, dtype, sharding, quantization settings, and provided batch size and maximum length.

Parameters

• batch_size (int) – The batch size for the cache.

• max_length (int) – The maximum sequence length the cache needs to support.

• starts (int | None) – Optional starting positions for the cache sequences. If provided, influences the initial state. Defaults to None (usually 0).

• shardings (dict | None) – Optional dictionary specifying sharding configurations. (Note: This argument appears unused in the current implementation shown).

• pad_token_id (int | None) – The ID of the padding token. If None, it’s inferred.

Returns

An initialized standard TransformerCache object.

Return type

TransformerCache

loss_type = 'ForCausalLM'

prepare_inputs_for_generation(input_ids: Union[Array, ndarray, bool, number], max_length: int, pad_token_id: int, starts: int | None = None, pixel_values: Optional[Union[Array, ndarray, bool, number]] = None, attention_mask: Optional[Union[Array, ndarray, bool, number]] = None)

Sets up the initial inputs required for starting autoregressive generation.

This function initializes the Key-Value cache (past_key_values) using init_cache, calculates the initial position_ids based on the input attention_mask (or assumes a contiguous range if no mask is provided), and prepares an extended attention_mask suitable for caching. It ensures inputs are placed on the correct devices/shards.

Parameters

• input_ids (chex.Array) – The initial sequence of token IDs. Shape (batch_size, seq_length).

• max_length (int) – The maximum sequence length that the KV cache should support.

• pad_token_id (int) – The ID used for padding tokens. Used to calculate starts if not provided.

• starts (int | None) – Optional pre-calculated starting positions (number of leading pads). If None, calculated using compute_prefill_length.

• shardings (dict | None) – Optional sharding configuration passed to init_cache.

• attention_mask (tp.Optional[chex.Array]) – An optional mask indicating which tokens should be attended to. Shape (batch_size, seq_length).

• token_type_ids (tp.Optional[chex.Array]) – Optional segment IDs for models that use them.

Returns

A dictionary containing the prepared inputs, typically including:

• ”past_key_values”: The initialized KV cache.

• ”attention_mask”: The extended attention mask for generation.

• ”position_ids”: The calculated initial position IDs.

• ”token_type_ids”: (Optional) Prepared token type IDs.

This dictionary is then passed through prepare_inputs_for_call.

Return type

dict

update_inputs_for_generation(model_outputs, model_kwargs)

Updates the keyword arguments for the next generation step.

Specifically, it takes the past_key_values from the model_outputs of the current step and updates the model_kwargs with them. It also increments the position_ids by one for the next token prediction.

Parameters

• model_outputs – The output object from the model’s forward pass in the previous step (should contain a past_key_values attribute).

• model_kwargs (dict) – The dictionary of keyword arguments used for the model call. This dictionary will be modified in-place or a new one returned.

Returns

The updated model_kwargs dictionary ready for the next generation step.

Return type

dict