vajra-streamer 0.0.22

A fast zero-copy PyTorch tensor streamer

Vajra

A fast zero-copy PyTorch tensor streamer powered by Dlang.

Usage

Once installed, you can use Vajra in any Python script or Jupyter Notebook:

import torch
from vajra import VajraStreamer, StreamConfig

# Configure the streamer (optional, uses defaults otherwise)
config = StreamConfig(
    auth_token="hf_YOUR_TOKEN", # Required for gated models like Llama 3
    chunk_size_mb=64,
    chunk_workers=16,
    gpu_workers=3,
    disable_cache=False
)

# Context manager ensures VRAM is freed when done
with VajraStreamer(config) as streamer:
    # Pass the URL or a Hugging Face repo ID
    # (e.g. "meta-llama/Meta-Llama-3-8B")
    tensors = streamer.load("meta-llama/Meta-Llama-3-8B")

    # 'tensors' is a dictionary mapping tensor names to zero-copy PyTorch tensors
    # that are backed directly by the downloaded GPU memory.
    for name, tensor in tensors.items():
        print(f"Tensor: {name}, Shape: {tensor.shape}, Dtype: {tensor.dtype}")