Skip to content

metadata

checkpoint.metadata.metadata

체크포인트 파일의 버전 및 metadata를 읽고 쓰는 API를 제공합니다.

CHECKPOINT_VERSION_LENGTH module-attribute 

CHECKPOINT_VERSION_LENGTH = 7

CHECKPOINT_SPEC module-attribute 

CHECKPOINT_SPEC = {b'2212210': {'version': b'2212210', 'metadata_size': 1024 * 1024, 'weight_position': CHECKPOINT_VERSION_LENGTH + 1024 * 1024}, b'2401040': {'version': b'2401040', 'metadata_size': 12 * 1024 * 1024, 'weight_position': CHECKPOINT_VERSION_LENGTH + 12 * 1024 * 1024}}

handle_unpickling_error 

handle_unpickling_error(function)
Source code in SaigeToolkit/checkpoint/metadata/metadata.py 
16
17
18
19
20
21
22
23
24
def handle_unpickling_error(function):
    @wraps(function)
    def wrapper(*args, **kwargs):
        try:
            return function(*args, **kwargs)
        except pickle.UnpicklingError:
            raise InvalidModelFileError

    return wrapper

read_checkpoint_version 

read_checkpoint_version(checkpoint_path: str) -> ByteString

체크포인트 파일의 버전을 읽어옵니다.

Parameters:

  • checkpoint_path (str) –

    checkpoint_path

Returns:

  • ByteString ( ByteString ) –

    체크포인트 파일의 버전

Raises:

  • ModelFileNotFoundError –

    checkpoint_path에 해당하는 파일이 없는 경우

Source code in SaigeToolkit/checkpoint/metadata/metadata.py 
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
@handle_unpickling_error
def read_checkpoint_version(checkpoint_path: str) -> ByteString:
    """체크포인트 파일의 버전을 읽어옵니다.

    Args:
        checkpoint_path (str): checkpoint_path

    Returns:
        ByteString: 체크포인트 파일의 버전

    Raises:
        ModelFileNotFoundError: checkpoint_path에 해당하는 파일이 없는 경우
    """
    if not os.path.isfile(checkpoint_path):
        raise ModelFileNotFoundError

    with open(checkpoint_path, "rb") as f:
        return f.read(CHECKPOINT_VERSION_LENGTH)

read_metadata 

read_metadata(checkpoint_path: str) -> Union[Dict, bool]

체크포인트 파일의 metadata 섹션 데이터를 읽어옵니다.

Parameters:

  • checkpoint_path (str) –

    checkpoint_path

Returns:

  • Union[Dict, bool] –

    Union[Dict, bool]: 체크포인트 버전이 맞지 않는 경우 False 값을 리턴

Raises:

  • ModelFileNotFoundError –

    checkpoint_path에 해당하는 파일이 없는 경우

Source code in SaigeToolkit/checkpoint/metadata/metadata.py 
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
@handle_unpickling_error
def read_metadata(checkpoint_path: str) -> Union[Dict, bool]:
    """체크포인트 파일의 metadata 섹션 데이터를 읽어옵니다.

    Args:
        checkpoint_path (str): checkpoint_path

    Returns:
        Union[Dict, bool]: 체크포인트 버전이 맞지 않는 경우 False 값을 리턴

    Raises:
        ModelFileNotFoundError: checkpoint_path에 해당하는 파일이 없는 경우
    """
    if not os.path.isfile(checkpoint_path):
        raise ModelFileNotFoundError
    with open(checkpoint_path, "rb") as f:
        version = f.read(CHECKPOINT_VERSION_LENGTH)
        return False if version not in CHECKPOINT_SPEC else pickle.load(f)

write_metadata 

write_metadata(checkpoint_path: str, metadata: Dict) -> bool

체크포인트 파일의 metadata 섹션 데이터를 수정해 저장합니다.

Parameters:

  • checkpoint_path (str) –

    checkpoint_path

  • metadata (Dict) –

    metadata

Returns:

  • bool ( bool ) –

    체크포인트 버전이 맞지 않는 경우 False 값을 리턴

Raises:

  • ModelFileNotFoundError –

    checkpoint_path에 해당하는 파일이 없는 경우

Source code in SaigeToolkit/checkpoint/metadata/metadata.py 
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
def write_metadata(checkpoint_path: str, metadata: Dict) -> bool:
    """체크포인트 파일의 metadata 섹션 데이터를 수정해 저장합니다.

    Args:
        checkpoint_path (str): checkpoint_path
        metadata (Dict): metadata

    Returns:
        bool: 체크포인트 버전이 맞지 않는 경우 False 값을 리턴

    Raises:
        ModelFileNotFoundError: checkpoint_path에 해당하는 파일이 없는 경우
    """
    if not os.path.isfile(checkpoint_path):
        raise ModelFileNotFoundError
    with open(checkpoint_path, "r+b") as f:
        version = f.read(CHECKPOINT_VERSION_LENGTH)
        if version not in CHECKPOINT_SPEC:
            return False
        bytes = pickle.dumps(metadata)
        metadata_size = CHECKPOINT_SPEC[version]["metadata_size"]
        if len(bytes) >= metadata_size:
            raise ValueError("metadata size exceeds the limit")
        f.write(bytes)
        return True