runai_model_streamer_gcs_bucket_to_transformers_quickstart.py

1from runai.model_streamer.gcs import GCSStreamer
2from transformers import AutoModelForCausalLM
3
4# Define the GCS bucket and the path to the model
5bucket_name = "my-model-bucket"
6model_path = "path/to/my/model"
7
8# Initialize the GCSStreamer
9streamer = GCSStreamer(bucket_name=bucket_name)
10
11# Stream the model directly into the transformers library
12# The streamer acts as a file-like object or provider for the model files
13with streamer.stream(model_path) as local_model_dir:
14    model = AutoModelForCausalLM.from_pretrained(local_model_dir)
15
16print("Model loaded successfully from GCS!")