micm_nlp.models.architectures

Model classes that HuggingFace does not ship.

CustomT5ForConditionalGeneration subclasses T5ForConditionalGeneration to add optional FlashAttention, enabled with flash_attn=True and silently unavailable when flash_attn is not installed.

This module stays small on purpose: every other backbone is selected by name from transformers through model.pretrained.cls, so adding one should need no code here. If it does, that is usually a signal the change belongs in the consumer repo.

Attributes

FlashSelfAttention

Classes

CustomT5ForConditionalGeneration

T5 with optional FlashAttention in place of its self-attention.

Module Contents

class micm_nlp.models.architectures.CustomT5ForConditionalGeneration(config: transformers.T5Config, flash_attn: bool = False)

Bases: transformers.T5ForConditionalGeneration

T5 with optional FlashAttention in place of its self-attention.

When flash_attn=True every self-attention module in the encoder and decoder is replaced by a FlashSelfAttention, with the pretrained q/k/v weights and biases copied across – so the swap does not change what the model computes, only how fast it does it. The decoder’s replacements are causal, the encoder’s are not.

flash_attn is optional at import: if the package is missing the class still loads, and only asking for it raises.

Parameters:

• config – the T5 config.

• flash_attn – swap in FlashAttention.

Raises:

RuntimeError – if flash_attn is requested but flash-attn is not installed – silently falling back would hide the reason a run is slower than expected.

forward(input_ids=None, attention_mask=None, inputs_embeds=None, decoder_input_ids=None, decoder_attention_mask=None, labels=None, task_ids=None, **kwargs)

Forward pass, dropping the arguments stock T5 does not accept.

task_ids is accepted so the signature matches the task-conditioned models, and ignored here. num_items_in_batch is popped from kwargs: newer Trainer versions pass it for loss averaging and T5 rejects it.

micm_nlp.models.architectures.FlashSelfAttention = None