from state import ModelState

GEMMA3_270M = ModelState(vocab_size=256000, num_layers=9, hidden_size=1152, intermediate_size=4608)
GEMMA3_1B = ModelState(vocab_size=262208, num_layers=26, hidden_size=2304, intermediate_size=9216)
GEMMA3_4B = ModelState(vocab_size=262208, num_layers=28, hidden_size=3072, intermediate_size=12288)
GEMMA3_12B = ModelState(vocab_size=262208, num_layers=42, hidden_size=4608, intermediate_size=18432)
GEMMA3_27B = ModelState(vocab_size=262208, num_layers=46, hidden_size=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
}