char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Daniel Beveniusmodel-conversion : merge inspect-org-model.py with tensor-info.py (#19823)72b44c0d2

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
9.3 KiB290 linesraw
1#!/usr/bin/env python3
2
3import argparse
4import json
5import os
6import re
7import struct
8import sys
9from pathlib import Path
10from typing import Optional
11from safetensors import safe_open
12
13
14MODEL_SAFETENSORS_FILE = "model.safetensors"
15MODEL_SAFETENSORS_INDEX = "model.safetensors.index.json"
16
17DTYPE_SIZES = {
18    "F64": 8, "I64": 8, "U64": 8,
19    "F32": 4, "I32": 4, "U32": 4,
20    "F16": 2, "BF16": 2, "I16": 2, "U16": 2,
21    "I8": 1, "U8": 1, "BOOL": 1,
22    "F8_E4M3": 1, "F8_E5M2": 1,
23}
24
25SIZE_UNITS = ['B', 'KB', 'MB', 'GB', 'TB']
26
27
28def get_weight_map(model_path: Path) -> Optional[dict[str, str]]:
29    index_file = model_path / MODEL_SAFETENSORS_INDEX
30
31    if index_file.exists():
32        with open(index_file, 'r') as f:
33            index = json.load(f)
34            return index.get("weight_map", {})
35
36    return None
37
38
39def get_all_tensor_names(model_path: Path) -> list[str]:
40    weight_map = get_weight_map(model_path)
41
42    if weight_map is not None:
43        return list(weight_map.keys())
44
45    single_file = model_path / MODEL_SAFETENSORS_FILE
46    if single_file.exists():
47        try:
48            with safe_open(single_file, framework="pt", device="cpu") as f:
49                return list(f.keys())
50        except Exception as e:
51            print(f"Error reading {single_file}: {e}")
52            sys.exit(1)
53
54    print(f"Error: No safetensors files found in {model_path}")
55    sys.exit(1)
56
57
58def find_tensor_file(model_path: Path, tensor_name: str) -> Optional[str]:
59    weight_map = get_weight_map(model_path)
60
61    if weight_map is not None:
62        return weight_map.get(tensor_name)
63
64    single_file = model_path / MODEL_SAFETENSORS_FILE
65    if single_file.exists():
66        return single_file.name
67
68    return None
69
70
71def read_safetensors_header(file_path: Path) -> dict:
72    with open(file_path, 'rb') as f:
73        header_size = struct.unpack('<Q', f.read(8))[0]
74        return json.loads(f.read(header_size))
75
76
77def get_tensor_size_bytes(tensor_meta: dict) -> int:
78    offsets = tensor_meta.get("data_offsets")
79    if offsets and len(offsets) == 2:
80        return offsets[1] - offsets[0]
81    n_elements = 1
82    for d in tensor_meta.get("shape", []):
83        n_elements *= d
84    return n_elements * DTYPE_SIZES.get(tensor_meta.get("dtype", "F32"), 4)
85
86
87def format_size(size_bytes: int) -> str:
88    val = float(size_bytes)
89    for unit in SIZE_UNITS[:-1]:
90        if val < 1024.0:
91            return f"{val:.2f} {unit}"
92        val /= 1024.0
93    return f"{val:.2f} {SIZE_UNITS[-1]}"
94
95
96def get_all_tensor_metadata(model_path: Path) -> dict[str, dict]:
97    weight_map = get_weight_map(model_path)
98
99    if weight_map is not None:
100        file_to_tensors: dict[str, list[str]] = {}
101        for tensor_name, file_name in weight_map.items():
102            file_to_tensors.setdefault(file_name, []).append(tensor_name)
103
104        all_metadata: dict[str, dict] = {}
105        for file_name, tensor_names in file_to_tensors.items():
106            try:
107                header = read_safetensors_header(model_path / file_name)
108                for tensor_name in tensor_names:
109                    if tensor_name in header:
110                        all_metadata[tensor_name] = header[tensor_name]
111            except Exception as e:
112                print(f"Warning: Could not read header from {file_name}: {e}", file=sys.stderr)
113        return all_metadata
114
115    single_file = model_path / MODEL_SAFETENSORS_FILE
116    if single_file.exists():
117        try:
118            header = read_safetensors_header(single_file)
119            return {k: v for k, v in header.items() if k != "__metadata__"}
120        except Exception as e:
121            print(f"Error reading {single_file}: {e}")
122            sys.exit(1)
123
124    print(f"Error: No safetensors files found in {model_path}")
125    sys.exit(1)
126
127
128def normalize_tensor_name(tensor_name: str) -> str:
129    normalized = re.sub(r'\.\d+\.', '.#.', tensor_name)
130    normalized = re.sub(r'\.\d+$', '.#', normalized)
131    return normalized
132
133
134def list_all_tensors(
135    model_path: Path,
136    short: bool = False,
137    show_sizes: bool = False,
138):
139    tensor_names = get_all_tensor_names(model_path)
140
141    metadata: Optional[dict[str, dict]] = None
142    if show_sizes:
143        metadata = get_all_tensor_metadata(model_path)
144
145    total_bytes = 0
146
147    if short:
148        seen: dict[str, str] = {}
149        for tensor_name in sorted(tensor_names):
150            normalized = normalize_tensor_name(tensor_name)
151            if normalized not in seen:
152                seen[normalized] = tensor_name
153        display_pairs = list(sorted(seen.items()))
154        name_width = max((len(n) for n, _ in display_pairs), default=0)
155        for normalized, first_name in display_pairs:
156            if metadata and first_name in metadata:
157                m = metadata[first_name]
158                size = get_tensor_size_bytes(m)
159                total_bytes += size
160                print(f"{normalized:{name_width}}  {m.get('dtype', '?'):6s}  {str(m.get('shape', '')):30s}  {format_size(size)}")
161            else:
162                print(normalized)
163    else:
164        name_width = max((len(n) for n in tensor_names), default=0)
165        for tensor_name in sorted(tensor_names):
166            if metadata and tensor_name in metadata:
167                m = metadata[tensor_name]
168                size = get_tensor_size_bytes(m)
169                total_bytes += size
170                print(f"{tensor_name:{name_width}}  {m.get('dtype', '?'):6s}  {str(m.get('shape', '')):30s}  {format_size(size)}")
171            else:
172                print(tensor_name)
173
174    if show_sizes:
175        print(f"\nTotal: {format_size(total_bytes)}")
176
177
178def print_tensor_info(model_path: Path, tensor_name: str, num_values: Optional[int] = None):
179    tensor_file = find_tensor_file(model_path, tensor_name)
180
181    if tensor_file is None:
182        print(f"Error: Could not find tensor '{tensor_name}' in model index")
183        print(f"Model path: {model_path}")
184        sys.exit(1)
185
186    file_path = model_path / tensor_file
187
188    try:
189        header = read_safetensors_header(file_path)
190        tensor_meta = header.get(tensor_name, {})
191        dtype_str = tensor_meta.get("dtype")
192
193        with safe_open(file_path, framework="pt", device="cpu") as f:
194            if tensor_name in f.keys():
195                tensor_slice = f.get_slice(tensor_name)
196                shape = tensor_slice.get_shape()
197                print(f"Tensor: {tensor_name}")
198                print(f"File:   {tensor_file}")
199                print(f"Shape:  {shape}")
200                if dtype_str:
201                    print(f"Dtype:  {dtype_str}")
202                if tensor_meta:
203                    print(f"Size:   {format_size(get_tensor_size_bytes(tensor_meta))}")
204                if num_values is not None:
205                    tensor = f.get_tensor(tensor_name)
206                    if not dtype_str:
207                        print(f"Dtype:  {tensor.dtype}")
208                    flat = tensor.flatten()
209                    n = min(num_values, flat.numel())
210                    print(f"Values: {flat[:n].tolist()}")
211            else:
212                print(f"Error: Tensor '{tensor_name}' not found in {tensor_file}")
213                sys.exit(1)
214
215    except FileNotFoundError:
216        print(f"Error: The file '{file_path}' was not found.")
217        sys.exit(1)
218    except Exception as e:
219        print(f"An error occurred: {e}")
220        sys.exit(1)
221
222
223def main():
224    parser = argparse.ArgumentParser(
225        description="Print tensor information from a safetensors model"
226    )
227    parser.add_argument(
228        "tensor_name",
229        nargs="?",
230        help="Name of the tensor to inspect"
231    )
232    parser.add_argument(
233        "-m", "--model-path",
234        type=Path,
235        help="Path to the model directory (default: MODEL_PATH environment variable)"
236    )
237    parser.add_argument(
238        "-l", "--list-all-short",
239        action="store_true",
240        help="List unique tensor patterns (layer numbers replaced with #)"
241    )
242    parser.add_argument(
243        "-la", "--list-all",
244        action="store_true",
245        help="List all tensor names with actual layer numbers"
246    )
247    parser.add_argument(
248        "-n", "--num-values",
249        nargs="?",
250        const=10,
251        default=None,
252        type=int,
253        metavar="N",
254        help="Print the first N values of the tensor flattened (default: 10 if flag is given without a number)"
255    )
256    parser.add_argument(
257        "-s", "--sizes",
258        action="store_true",
259        help="Show dtype, shape, and size for each tensor when listing"
260    )
261
262    args = parser.parse_args()
263
264    model_path = args.model_path
265    if model_path is None:
266        model_path_str = os.environ.get("MODEL_PATH")
267        if model_path_str is None:
268            print("Error: --model-path not provided and MODEL_PATH environment variable not set")
269            sys.exit(1)
270        model_path = Path(model_path_str)
271
272    if not model_path.exists():
273        print(f"Error: Model path does not exist: {model_path}")
274        sys.exit(1)
275
276    if not model_path.is_dir():
277        print(f"Error: Model path is not a directory: {model_path}")
278        sys.exit(1)
279
280    if args.list_all_short or args.list_all:
281        list_all_tensors(model_path, short=args.list_all_short, show_sizes=args.sizes)
282    else:
283        if args.tensor_name is None:
284            print("Error: tensor_name is required when not using --list-all-short or --list-all")
285            sys.exit(1)
286        print_tensor_info(model_path, args.tensor_name, args.num_values)
287
288
289if __name__ == "__main__":
290    main()