Felipe97/llama-cpp-compiled

1#!/usr/bin/env python32 3import argparse4import json5import os6import re7import struct8import sys9from pathlib import Path10from typing import Optional11from safetensors import safe_open12 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_INDEX30 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 None37 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_FILE46    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_FILE65    if single_file.exists():66        return single_file.name67 68    return None69 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 = 182    for d in tensor_meta.get("shape", []):83        n_elements *= d84    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.093    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_metadata114 115    single_file = model_path / MODEL_SAFETENSORS_FILE116    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 normalized132 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]] = None142    if show_sizes:143        metadata = get_all_tensor_metadata(model_path)144 145    total_bytes = 0146 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_name153        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 += size160                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 += size170                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_file187 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_path265    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()291