39 lines
846 B
Python
Executable File
39 lines
846 B
Python
Executable File
|
|
def get_model(cfg, training=True):
|
|
name = cfg.name
|
|
|
|
if not cfg.experimental:
|
|
from .ar_nar import AR_NAR
|
|
model = AR_NAR(
|
|
n_tokens=cfg.tokens,
|
|
d_model=cfg.dim,
|
|
n_heads=cfg.heads,
|
|
n_layers=cfg.layers,
|
|
n_experts=cfg.experts,
|
|
|
|
p_dropout=cfg.dropout,
|
|
|
|
l_padding = cfg.input_alignment,
|
|
|
|
training = training,
|
|
config = cfg,
|
|
)
|
|
model._cfg = cfg
|
|
else:
|
|
from .experimental import Model as Experimental
|
|
model = Experimental(
|
|
d_model=cfg.dim,
|
|
n_layers=cfg.layers,
|
|
n_heads=cfg.heads,
|
|
p_dropout=cfg.dropout,
|
|
|
|
config = cfg,
|
|
)
|
|
|
|
print(f"{name} ({next(model.parameters()).dtype}): {sum(p.numel() for p in model.parameters() if p.requires_grad)} parameters")
|
|
|
|
return model
|
|
|
|
def get_models(models, training=True):
|
|
return { model.full_name: get_model(model, training=training) for model in models }
|