File size: 1,707 Bytes
29f25be | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | """Architecture dispatch shared by training and checkpoint inference."""
from vimeml.training.model import GPTConfig, TinyGPT
from vimeml.training.model_v2 import GPTV2Config, TinyGPTV2
ARCHITECTURES = {
"tiny_gpt_v1": (GPTConfig, TinyGPT, "vimeml_tiny_gpt_v1"),
"tiny_gpt_v2": (GPTV2Config, TinyGPTV2, "vimeml_tiny_gpt_v2"),
}
def specification(architecture):
if architecture not in ARCHITECTURES:
raise ValueError(f"Unsupported architecture: {architecture}")
return ARCHITECTURES[architecture]
def configuration_for(architecture, values):
config_type, _, _ = specification(architecture)
return config_type(**values)
def create_model(architecture, config):
config_type, model_type, _ = specification(architecture)
if type(config) is not config_type:
raise ValueError("Model configuration does not match architecture.")
return model_type(config)
def checkpoint_format(architecture):
return specification(architecture)[2]
def model_from_checkpoint(saved):
architecture = next(
(
name
for name, (_, _, format_name) in ARCHITECTURES.items()
if format_name == saved.get("format")
),
None,
)
if architecture is None:
raise ValueError("Unsupported checkpoint format.")
if (
saved.get("architecture", architecture) != architecture
or saved.get("config", {}).get("architecture", architecture) != architecture
):
raise ValueError("Checkpoint architecture differs from its format.")
model = create_model(architecture, configuration_for(architecture, saved["model_config"]))
model.load_state_dict(saved["model"])
return model
|