def load_file()

in bindings/python/py_src/safetensors/flax.py [0:0]

• 6 lines of code
• 2 McCabe index (conditional complexity)

def load_file(filename: Union[str, os.PathLike]) -> Dict[str, Array]:
    """
    Loads a safetensors file into flax format.

    Args:
        filename (`str`, or `os.PathLike`)):
            The name of the file which contains the tensors

    Returns:
        `Dict[str, Array]`: dictionary that contains name as key, value as `Array`

    Example:

    ```python
    from safetensors.flax import load_file

    file_path = "./my_folder/bert.safetensors"
    loaded = load_file(file_path)
    ```
    """
    result = {}
    with safe_open(filename, framework="flax") as f:
        for k in f.offset_keys():
            result[k] = f.get_tensor(k)
    return result