File size: 2,715 Bytes
3194a67 | 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 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 | from __future__ import annotations
import json
import random
from pathlib import Path
from typing import Any
import numpy as np
import torch
from .config import Config
from .model import MultiMDMTransformer
from .tokenizer import CharacterTokenizer
def set_seed(seed: int) -> None:
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
def resolve_device(requested: str) -> torch.device:
if requested == "auto":
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
if requested == "cuda" and not torch.cuda.is_available():
print("CUDA unavailable; falling back to CPU.")
return torch.device("cpu")
return torch.device(requested)
def build_model(config: Config, tokenizer: CharacterTokenizer) -> MultiMDMTransformer:
return MultiMDMTransformer(
vocab_size=tokenizer.vocab_size,
clean_vocab_size=tokenizer.clean_vocab_size,
seq_len=config.seq_len,
num_masks=config.num_masks,
d_model=config.d_model,
nhead=config.nhead,
num_layers=config.num_layers,
dim_feedforward=config.dim_feedforward,
dropout=config.dropout,
)
def save_json(value: Any, path: str | Path) -> None:
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
with path.open("w", encoding="utf-8") as handle:
json.dump(value, handle, indent=2, ensure_ascii=False)
def save_checkpoint(
path: str | Path,
model: torch.nn.Module,
optimizer: torch.optim.Optimizer,
config: Config,
tokenizer: CharacterTokenizer,
step: int,
metrics: dict[str, Any],
) -> None:
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
torch.save(
{
"model_state_dict": model.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
"config": config.to_dict(),
"vocab": tokenizer.vocab,
"num_masks": tokenizer.num_masks,
"step": step,
"metrics": metrics,
},
path,
)
def load_checkpoint(
path: str | Path, device: str | torch.device = "cpu"
) -> tuple[MultiMDMTransformer, CharacterTokenizer, Config, dict[str, Any]]:
device = torch.device(device)
checkpoint = torch.load(path, map_location=device, weights_only=False)
config = Config.from_dict(checkpoint["config"])
tokenizer = CharacterTokenizer(checkpoint["vocab"], checkpoint["num_masks"])
model = build_model(config, tokenizer).to(device)
model.load_state_dict(checkpoint["model_state_dict"])
model.eval()
return model, tokenizer, config, checkpoint
|