ez80ai/exportmodel.py at master · tornadus/ez80ai · GitHub

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
#!/usr/bin/env python3
"""
Export NEOCHAT model checkpoint to NumPy .npz format.

Handles dual bias export (fc4_bias + fc4_bias_start).

Usage:
    python3 exportmodel.py
    python3 exportmodel.py --model neochat_model.pt --output model.npz
"""

import argparse
import json
import os
import numpy as np
import torch

from train import (NeochatModel, CHARSET, NUM_CHARS, HIDDEN_SIZES, INPUT_SIZE,
                   filter_legacy_state)
import modelspec
from loadmodel import load_spec_from_model


def export_model(model_path, output_path):
    print(f"Loading model from {model_path}...")
    checkpoint = torch.load(model_path, weights_only=False, map_location='cpu')
    arch = checkpoint['architecture']
    charset = checkpoint.get('charset', CHARSET)
    num_chars = arch['num_classes']

    print(f"Architecture: input={arch['input_size']}, hidden={arch['hidden_sizes']}, output={num_chars}")
    print(f"Charset ({num_chars} chars): {repr(charset[:-1])} + EOS")

    model = NeochatModel(
        input_size=arch['input_size'],
        hidden_sizes=arch['hidden_sizes'],
        num_chars=num_chars,
    )
    model.load_state_dict(filter_legacy_state(checkpoint['model_state']))
    model.eval()

    params = model.get_quantized_params()

    export_data = {}
    for key, value in params.items():
        export_data[key] = value

    export_data['_architecture'] = np.array(json.dumps(arch).encode('utf-8'))
    export_data['_charset'] = np.array(charset.encode('utf-8'))
    export_data['_dual_bias_threshold'] = np.array(
        checkpoint.get('dual_bias_threshold', 3))
    # Carry the frozen resolved spec forward into the .npz so the build and the
    # faithfulness gate read the exact spec the model was trained with. Legacy
    # checkpoints (no baked spec) reconstruct an equivalent one from arch.
    spec = load_spec_from_model(model_path)
    export_data['_modelspec'] = np.array(modelspec.to_json(spec).encode('utf-8'))

    np.savez(output_path, **export_data)
    print(f"Exported to {output_path}")

    layer_names = sorted(set(k.replace('_weight', '').replace('_bias', '').replace('_start', '')
                             for k in params.keys()))
    for name in layer_names:
        w_key = f'{name}_weight'
        b_key = f'{name}_bias'
        if w_key in params:
            w = params[w_key]
            b = params[b_key]
            extra = ""
            if f'{name}_bias_start' in params:
                extra = f" + bias_start {params[f'{name}_bias_start'].shape}"
            print(f"  {name}: weight {w.shape}, bias {b.shape}{extra}")


def main():
    default_model = os.path.join(os.path.dirname(__file__), 'neochat_model.pt')
    default_output = os.path.join(os.path.dirname(__file__), 'model.npz')

    parser = argparse.ArgumentParser(description='Export NEOCHAT model to NumPy format')
    parser.add_argument('--model', '-m', default=default_model, help='Input .pt checkpoint')
    parser.add_argument('--output', '-o', default=default_output, help='Output .npz file')
    args = parser.parse_args()

    export_model(args.model, args.output)


if __name__ == '__main__':
    main()