Method __init__

lmdeploy/turbomind/checkpoint.py:229–239

(self, model_path: str, *,
                 mappings=(),
                 index_name: str | None = None,
                 file_pattern: str | None = None)

    """torch.load-backed checkpoint over ``*.bin`` / ``*.pt`` shards."""

    def __init__(self, model_path: str, *,
                 mappings=(),
                 index_name: str | None = None,
                 file_pattern: str | None = None):
        self._mappings = list(mappings)
        shards, _ = _gather_shards(model_path, index_name, file_pattern)
        self._data: dict[str, torch.Tensor] = {}
        for shard in shards:
            tmp = torch.load(shard, map_location='cpu', weights_only=True)
            for k, v in tmp.items():
                self._data[_apply_mappings(k, self._mappings)] = v

    def get(self, key: str, index=None):
        t = self._data[key]