fn_registry/python/functions/ml/safetensors_inspect.py

"""Lee solo el header de un archivo safetensors sin cargar los tensores en RAM."""

from __future__ import annotations

import json
import struct


def safetensors_inspect(path: str) -> dict:
    """Lee el header de un archivo safetensors sin cargar los tensores.

    El formato safetensors almacena al inicio del archivo:
      - 8 bytes: uint64 little-endian con la longitud del header JSON (N).
      - N bytes: JSON con metadata y descriptores de tensores.

    Este enfoque evita cargar gigabytes de pesos en RAM para inspeccionar
    un checkpoint: solo se leen los primeros 8 + N bytes.

    Spec: https://github.com/huggingface/safetensors#format

    Args:
        path: ruta al archivo .safetensors (absoluta o relativa).

    Returns:
        dict con claves:
            path (str): ruta absoluta del archivo.
            metadata (dict): metadatos del modelo (campo "__metadata__" del header).
            tensors (list[dict]): lista de tensores, cada uno con:
                name (str): nombre del tensor.
                dtype (str): tipo de dato (F32, BF16, F16, I32, etc.).
                shape (list[int]): dimensiones del tensor.
                data_offsets (list[int]): [inicio, fin] en bytes dentro del
                    bloque de datos (para acceso lazy si se necesita).
            total_size_bytes (int): tamano total del archivo en bytes.
            n_tensors (int): numero de tensores en el archivo.

    Raises:
        FileNotFoundError: si el archivo no existe.
        ValueError: si el archivo no es un safetensors valido.
        OSError: si no se puede leer el archivo.
    """
    import os

    abs_path = os.path.abspath(path)

    if not os.path.isfile(abs_path):
        raise FileNotFoundError(f"safetensors_inspect: archivo no encontrado: {abs_path}")

    total_size = os.path.getsize(abs_path)

    with open(abs_path, "rb") as f:
        # Leer los 8 bytes del tamano del header
        raw_len = f.read(8)
        if len(raw_len) < 8:
            raise ValueError(
                f"safetensors_inspect: archivo demasiado corto para ser safetensors: {abs_path}"
            )

        header_len = struct.unpack("<Q", raw_len)[0]  # uint64 little-endian

        if header_len == 0 or header_len > total_size:
            raise ValueError(
                f"safetensors_inspect: header_len invalido ({header_len}) en {abs_path}"
            )

        raw_header = f.read(header_len)
        if len(raw_header) < header_len:
            raise ValueError(
                f"safetensors_inspect: no se pudo leer el header completo en {abs_path}"
            )

    try:
        header = json.loads(raw_header.decode("utf-8"))
    except (json.JSONDecodeError, UnicodeDecodeError) as exc:
        raise ValueError(
            f"safetensors_inspect: header JSON invalido en {abs_path}: {exc}"
        ) from exc

    metadata = header.get("__metadata__", {})

    tensors = []
    for name, desc in header.items():
        if name == "__metadata__":
            continue
        tensors.append(
            {
                "name": name,
                "dtype": desc.get("dtype", ""),
                "shape": desc.get("shape", []),
                "data_offsets": desc.get("data_offsets", []),
            }
        )

    return {
        "path": abs_path,
        "metadata": metadata,
        "tensors": tensors,
        "total_size_bytes": total_size,
        "n_tensors": len(tensors),
    }