PebbleLM-117M / src /model /__init__.py
nameissakthi's picture
Add model architecture code
27871e7
"""SLM Model Components."""
from .config import SLMConfig
from .transformer import SLMForCausalLM, SLMModel, SLMOutput
from .kv_cache import KVCache
from .normalization import RMSNorm
from .rope import RotaryEmbedding
from .attention import MultiHeadAttention, create_causal_mask
from .ffn import FeedForward
from .decoder import DecoderBlock
__all__ = [
"SLMConfig",
"SLMForCausalLM",
"SLMModel",
"SLMOutput",
"KVCache",
"RMSNorm",
"RotaryEmbedding",
"MultiHeadAttention",
"create_causal_mask",
"FeedForward",
"DecoderBlock",
]