tensor-man

Crates.iotensor-man
lib.rstensor-man
version0.4.2
sourcesrc
created_at2024-10-24 17:25:10.021953
updated_at2024-11-01 13:12:31.475075
descriptionA small utility to inspect and validate safetensors and ONNX files.
homepagehttps://github.com/dreadnode/tman
repositoryhttps://github.com/dreadnode/tman
max_upload_size
id1421574
size191,707
Simone Margaritelli (evilsocket)

documentation

README

Release Crate Docker Hub Rust Report GitHub Actions Workflow Status Software License

tensor-man is a utility to inspect, validate, sign and verify machine learning model files.

Supported Formats

[!IMPORTANT] PyTorch models are loaded and inspected in a networkless Docker container in order to prevent unintended code execution on the host machine.

Install with Cargo

This is the recommended way to install and use the tool:

cargo install tensor-man

Pull from Docker Hub

docker pull dreadnode/tensor-man:latest

Build Docker image

To build your own Docker image for the tool, run:

docker build . -t tman  

Note about Docker

If you want to inspect PyTorch models and you are using tensor-man inside a container, make sure to share the docker socket from the host machine with the container:

docker run -it \
  # these paths must match
  -v/path/to/pytorch_model.bin:/path/to/pytorch_model.bin \
  # allow the container itself to instrument docker on the host
  -v/var/run/docker.sock:/var/run/docker.sock \
  # the rest of the command line
  tman inspect /path/to/pytorch_model.bin

Build from source

Alternatively you can build the project from source, in which case you'll need to have Rust and Cargo installed on your system.

Once you have those set up, clone the repository and build the project:

cargo build --release

The compiled binary will be available in the target/release directory. You can run it directly or add it to your system's PATH:

# Run directly
./target/release/tman

# Or, copy to a directory in your PATH (e.g., /usr/local/bin)
sudo cp target/release/tman /usr/local/bin/

Usage

Inspect

Inspect a file and print a brief summary:

tman inspect /path/to/whatever/llama-3.1-8b-instruct.safetensors

Print detailed information about each tensor:

tman inspect /path/to/whatever/llama-3.1-8b-instruct.safetensors --detail full

Filter by tensor name:

tman inspect /path/to/whatever/llama-3.1-8b-instruct.onnx -D full --filter "q_proj"

Save the output as JSON:

tman inspect /path/to/whatever/llama-3.1-8b-instruct.gguf -D full --to-json output.json

Sign and Verify

The tool allows you to generate an Ed25519 key pair to sign your models:

tman create-key --private-key private.key --public-key public.key

Then you can use the private key to sign a model (this will automatically include and sign external data files if referenced by the format):

# this will generate the tinyyolov2-8.signature file
tman sign /path/to/whatever/tinyyolov2-8.onnx -K /path/to/private.key

# you can provide a safetensors index file and all files referenced by it will be signed as well
tman sign /path/to/whatever/Meta-Llama-3-8B/model.safetensors.index.json -K /path/to/private.key

# this will sign the entire model folder with every file in it
tman sign /path/to/whatever/Meta-Llama-3-8B/ -K /path/to/private.key

And the public one to verify the signature:

# will verify the signature in tinyyolov2-8.signature
tman verify /path/to/whatever/tinyyolov2-8.onnx -K /path/to/public.key

# will verify with an alternative signature file 
tman verify /path/to/whatever/tinyyolov2-8.onnx -K /path/to/public.key --signature /path/to/your.signature

# this will verify every file in the model folder
tman sign /path/to/whatever/Meta-Llama-3-8B/ -K /path/to/public.key

Inference Graph

Generate a .dot file for the execution graph of an ONNX model:

tman graph /path/to/whatever/tinyyolov2-8.onnx --output tinyyolov2-8.dot

More

For the full list of commands and options, run:

tman --help

# get command specific help
tman inspect --help

License

This tool is released under the GPL 3 license. To see the licenses of the project dependencies, install cargo license with cargo install cargo-license and then run cargo license.

Commit count: 67

cargo fmt