Imported from dzylol/graduation-project (src/models/AGENTS.md). Install upstream with npx skills add dzylol/graduation-project --skill models. Copyright stays with the author.

AGENTS.md - src/models/

Generated: 2026-04-20 (updated MLP head)

Core Bi-Mamba architecture. All model code lives here.

Structure

src/models/
├── __init__.py
├── bimamba.py                  # Manual SSM (primary, no external deps)
├── bimamba_with_mamba_ssm.py   # Wrapper using mamba-ssm package
├── bimamba_with_mamba_ssm_architecture.md  # Architecture tutorial
└── AGENTS.md

Key Classes

Fusion Modes

• gate (default): sigmoid(W_fwd) * fwd + (1-sigmoid(W_fwd)) * bwd
• concat: W * concat(fwd, bwd)
• add: fwd + bwd

Pooling Types

• mean (default): global average pooling
• max: global max pooling
• cls: [CLS] token pooling

Model Types

Conventions (THIS MODULE)

• SSM parameters: d_state=16, d_conv=4, expand=2
• dt_rank: auto = ceil(d_model / 16)
• Activation: SiLU (nn.SiLU)
• A_log: nn.Parameter(log(A)) where A = torch.arange(1, d_state+1)

Head Architecture

Regression (task_type=regression):

MLP head: nn.Sequential(
    nn.Linear(d_model, d_model // 2),
    nn.ReLU(),
    nn.Dropout(dropout),
    nn.Linear(d_model // 2, num_labels),
)

Classification (task_type=classification):

Linear head: nn.Linear(d_model, num_labels)

Anti-Patterns (THIS MODULE)

• NEVER use as any or @ts-ignore — type safety is required
• NEVER use bare dicts — use dataclasses