fla/models/__init__.py

# -*- coding: utf-8 -*-

from fla.models.abc import ABCConfig, ABCForCausalLM, ABCModel
from fla.models.bitnet import BitNetConfig, BitNetForCausalLM, BitNetModel
from fla.models.delta_net import DeltaNetConfig, DeltaNetForCausalLM, DeltaNetModel
from fla.models.forgetting_transformer import (
    ForgettingTransformerConfig,
    ForgettingTransformerForCausalLM,
    ForgettingTransformerModel
)
from fla.models.gated_deltanet import GatedDeltaNetConfig, GatedDeltaNetForCausalLM, GatedDeltaNetModel
from fla.models.gated_deltaproduct import GatedDeltaProductConfig, GatedDeltaProductForCausalLM, GatedDeltaProductModel
from fla.models.gla import GLAConfig, GLAForCausalLM, GLAModel
from fla.models.gsa import GSAConfig, GSAForCausalLM, GSAModel
from fla.models.hgrn import HGRNConfig, HGRNForCausalLM, HGRNModel
from fla.models.hgrn2 import HGRN2Config, HGRN2ForCausalLM, HGRN2Model
from fla.models.lightnet import LightNetConfig, LightNetForCausalLM, LightNetModel
from fla.models.linear_attn import LinearAttentionConfig, LinearAttentionForCausalLM, LinearAttentionModel
from fla.models.mamba import MambaConfig, MambaForCausalLM, MambaModel
from fla.models.mamba2 import Mamba2Config, Mamba2ForCausalLM, Mamba2Model
from fla.models.nsa import NSAConfig, NSAForCausalLM, NSAModel
from fla.models.retnet import RetNetConfig, RetNetForCausalLM, RetNetModel
from fla.models.rwkv6 import RWKV6Config, RWKV6ForCausalLM, RWKV6Model
from fla.models.rwkv7 import RWKV7Config, RWKV7ForCausalLM, RWKV7Model
from fla.models.samba import SambaConfig, SambaForCausalLM, SambaModel
from fla.models.transformer import TransformerConfig, TransformerForCausalLM, TransformerModel
from fla.models.transformer_top import TOPTransformerConfig, TOPTransformerForCausalLM, TOPTransformerModel
from fla.models.transformer_mtp import MTPTransformerConfig, MTPTransformerForCausalLM, MTPTransformerModel

__all__ = [
    'ABCConfig', 'ABCForCausalLM', 'ABCModel',
    'BitNetConfig', 'BitNetForCausalLM', 'BitNetModel',
    'DeltaNetConfig', 'DeltaNetForCausalLM', 'DeltaNetModel',
    'ForgettingTransformerConfig', 'ForgettingTransformerForCausalLM', 'ForgettingTransformerModel',
    'GatedDeltaNetConfig', 'GatedDeltaNetForCausalLM', 'GatedDeltaNetModel',
    'GLAConfig', 'GLAForCausalLM', 'GLAModel',
    'GSAConfig', 'GSAForCausalLM', 'GSAModel',
    'HGRNConfig', 'HGRNForCausalLM', 'HGRNModel',
    'HGRN2Config', 'HGRN2ForCausalLM', 'HGRN2Model',
    'LightNetConfig', 'LightNetForCausalLM', 'LightNetModel',
    'LinearAttentionConfig', 'LinearAttentionForCausalLM', 'LinearAttentionModel',
    'MambaConfig', 'MambaForCausalLM', 'MambaModel',
    'Mamba2Config', 'Mamba2ForCausalLM', 'Mamba2Model',
    'NSAConfig', 'NSAForCausalLM', 'NSAModel',
    'RetNetConfig', 'RetNetForCausalLM', 'RetNetModel',
    'RWKV6Config', 'RWKV6ForCausalLM', 'RWKV6Model',
    'RWKV7Config', 'RWKV7ForCausalLM', 'RWKV7Model',
    'SambaConfig', 'SambaForCausalLM', 'SambaModel',
    'TransformerConfig', 'TransformerForCausalLM', 'TransformerModel',
    'TOPTransformerConfig', 'TOPTransformerForCausalLM', 'TOPTransformerModel',
    'MTPTransformerConfig', 'MTPTransformerForCausalLM', 'MTPTransformerModel',
    'GatedDeltaProductConfig', 'GatedDeltaProductForCausalLM', 'GatedDeltaProductModel',
]