from state import ModelState

GEMMA3_270M = ModelState(
    vocab_size=256000, num_layers=9, hidden_dim=1152, intermediate_size=4608
)
GEMMA3_1B = ModelState(
    vocab_size=262208, num_layers=26, hidden_dim=2304, intermediate_size=9216
)
GEMMA3_4B = ModelState(
    vocab_size=262208, num_layers=28, hidden_dim=3072, intermediate_size=12288
)
GEMMA3_12B = ModelState(
    vocab_size=262208, num_layers=42, hidden_dim=4608, intermediate_size=18432
)
GEMMA3_27B = ModelState(
    vocab_size=262208, num_layers=46, hidden_dim=6144, intermediate_size=24576
)

DEFAULTS = {
    "Gemma3 270M": GEMMA3_270M,
    "Gemma3 1B": GEMMA3_1B,
    "Gemma3 4B": GEMMA3_4B,
    "Gemma3 12B": GEMMA3_12B,
    "Gemma3 27B": GEMMA3_27B,
}