File size: 5,456 Bytes
ec3cdee | 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 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 | import copy
import hashlib
import json
import shutil
from pathlib import Path
import torch
import torch.nn as nn
from safetensors.torch import load_file, save_file
from transformers import PreTrainedModel
from transformers.utils.hub import cached_file
from .configuration_mmtb import MMTBConfig
class MMTBBaseModel(PreTrainedModel):
config_class = MMTBConfig
base_model_prefix = "models"
def __init__(self, config):
super().__init__(config)
if config.task != self.task:
raise ValueError(f"Expected task {self.task}, received task {config.task}")
self.models = nn.ModuleList(self.build_models(config))
@classmethod
def from_pretrained(cls, pretrained_model_name_or_path, *model_args, config=None,
cache_dir=None, force_download=False, local_files_only=False,
token=None, revision="main", **kwargs):
if model_args:
raise TypeError("Positional model arguments are not supported")
subfolder = kwargs.pop("subfolder", "")
commit_hash = kwargs.pop("_commit_hash", None)
for key in ("_from_auto", "_from_pipeline", "trust_remote_code", "adapter_kwargs"):
kwargs.pop(key, None)
if config is None:
config = cls.config_class.from_pretrained(
pretrained_model_name_or_path, cache_dir=cache_dir, force_download=force_download,
local_files_only=local_files_only, token=token, revision=revision, subfolder=subfolder,
)
config = copy.deepcopy(config)
if "mode" in kwargs:
config.mode = kwargs.pop("mode")
device = kwargs.pop("device", "cpu")
dtype = kwargs.pop("dtype", None)
legacy_dtype = kwargs.pop("torch_dtype", None)
dtype = dtype if dtype is not None else legacy_dtype
if dtype is None or dtype == "auto":
dtype = getattr(config, "dtype", None) or torch.float32
if kwargs:
raise TypeError(f"Unsupported loading arguments: {', '.join(sorted(kwargs))}; use device= or model.to() for placement")
if device == "auto":
device = "cuda" if torch.cuda.is_available() else "cpu"
device = torch.device(device)
if dtype in (None, "auto"):
dtype = torch.float32
elif isinstance(dtype, str):
dtype = getattr(torch, dtype, None)
if dtype not in (torch.float32, torch.float16, torch.bfloat16, torch.float64):
raise ValueError("dtype must be a floating-point torch dtype or 'auto'")
if config.task != cls.task:
raise ValueError(f"Expected task {cls.task}, received task {config.task}")
indices = cls.checkpoint_indices(config)
resolved_revision = getattr(config, "_commit_hash", None) or commit_hash or revision
paths = [cached_file(
pretrained_model_name_or_path, f"weights/model_{index}.safetensors",
cache_dir=cache_dir, force_download=force_download, local_files_only=local_files_only,
token=token, revision=resolved_revision, subfolder=subfolder,
) for index in indices]
model = cls(config)
for net, path in zip(model.models, paths):
net.load_state_dict(load_file(path), strict=True)
net.to(device=device, dtype=dtype)
model.eval()
return model
def save_pretrained(self, save_directory, **kwargs):
if kwargs:
raise TypeError(f"Unsupported saving arguments: {', '.join(sorted(kwargs))}")
directory = Path(save_directory)
directory.mkdir(parents=True, exist_ok=True)
weights = directory / "weights"
weights.mkdir(exist_ok=True)
expected = {f"model_{index}.safetensors" for index in self.checkpoint_indices(self.config)}
if any(path.name not in expected for path in weights.iterdir()):
raise ValueError("Save to a fresh directory or one containing only the selected checkpoint files")
entries = []
for index, net in zip(self.checkpoint_indices(self.config), self.models):
target = weights / f"model_{index}.safetensors"
state = {name: tensor.detach().cpu().contiguous() for name, tensor in net.state_dict().items()}
save_file(state, str(target), metadata={"format": "pt"})
digest = hashlib.sha256()
with target.open("rb") as source:
for chunk in iter(lambda: source.read(8 * 1024 * 1024), b""):
digest.update(chunk)
entries.append({"name": f"model_{index}", "file": f"weights/{target.name}",
"bytes": target.stat().st_size, "sha256": digest.hexdigest()})
config = copy.deepcopy(self.config)
config.auto_map = {"AutoConfig": "configuration_mmtb.MMTBConfig", "AutoModel": "modeling_mmtb.MMTBModel"}
config.architectures = ["MMTBModel"]
config.dtype = self.dtype
config.save_pretrained(directory)
for filename in ("configuration_mmtb.py", "modeling_base.py", "modeling_mmtb.py", "predict.py"):
source = Path(__file__).resolve().parent / filename
target = directory / filename
if source.resolve() != target.resolve():
shutil.copyfile(source, target)
(directory / "model_manifest.json").write_text(json.dumps({"task": f"task{self.task}", "models": entries}, indent=2) + "\n")
|