runai_model_streamer_safetensors_pytorch_loading_quickstart.py

1import torch
2from runai_model_streamer import load_model
3
4# Define the path to your model (directory containing safetensors files)
5model_path = "path/to/your/model"
6
7# Load the model using Run:ai Model Streamer
8# This replaces the standard torch.load or transformer's from_pretrained
9# by streaming the weights directly into the device memory
10model = load_model(model_path)
11
12# Verify the model is loaded on the desired device (e.g., CUDA)
13device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
14model.to(device)
15
16print(f"Model successfully loaded from {model_path} using Run:ai Model Streamer")