Run an LLM in a streaming pipeline

  # Copyright 2023 Google LLC 
 # 
 # Licensed under the Apache License, Version 2.0 (the "License"); 
 # you may not use this file except in compliance with the License. 
 # You may obtain a copy of the License at 
 # 
 #      http://www.apache.org/licenses/LICENSE-2.0 
 # 
 # Unless required by applicable law or agreed to in writing, software 
 # distributed under the License is distributed on an "AS IS" BASIS, 
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 
 # See the License for the specific language governing permissions and 
 # limitations under the License. 
 """Loads the state_dict for an LLM model into Cloud Storage.""" 
 from 
  
 __future__ 
  
 import 
 annotations 
 import 
  
 os 
 import 
  
 torch 
 from 
  
 transformers 
  
 import 
 AutoModelForSeq2SeqLM 
 def 
  
 run_local 
 ( 
 model_name 
 : 
 str 
 , 
 state_dict_path 
 : 
 str 
 ) 
 - 
> None 
 : 
  
 """Loads the state dict and saves it into the desired path. 
 If the `state_dict_path` is a Cloud Storage location starting 
 with "gs://", this assumes Cloud Storage is mounted with 
 Cloud Storage FUSE in `/gcs`. Vertex AI is set up like this. 
 Args: 
 model_name: HuggingFace model name compatible with AutoModelForSeq2SeqLM. 
 state_dict_path: File path to the model's state_dict, can be in Cloud Storage. 
 """ 
 print 
 ( 
 f 
 "Loading model: 
 { 
 model_name 
 } 
 " 
 ) 
 model 
 = 
 AutoModelForSeq2SeqLM 
 . 
 from_pretrained 
 ( 
 model_name 
 , 
 torch_dtype 
 = 
 torch 
 . 
 bfloat16 
 ) 
 print 
 ( 
 f 
 "Model loaded, saving state dict to: 
 { 
 state_dict_path 
 } 
 " 
 ) 
 # Assume Cloud Storage FUSE is mounted in `/gcs`. 
 state_dict_path 
 = 
 state_dict_path 
 . 
 replace 
 ( 
 "gs://" 
 , 
 "/gcs/" 
 ) 
 directory 
 = 
 os 
 . 
 path 
 . 
 dirname 
 ( 
 state_dict_path 
 ) 
 if 
 directory 
 and 
 not 
 os 
 . 
 path 
 . 
 exists 
 ( 
 directory 
 ): 
 os 
 . 
 makedirs 
 ( 
 os 
 . 
 path 
 . 
 dirname 
 ( 
 state_dict_path 
 ), 
 exist_ok 
 = 
 True 
 ) 
 torch 
 . 
 save 
 ( 
 model 
 . 
 state_dict 
 (), 
 state_dict_path 
 ) 
 print 
 ( 
 "State dict saved successfully!" 
 ) 
 def 
  
 run_vertex_job 
 ( 
 model_name 
 : 
 str 
 , 
 state_dict_path 
 : 
 str 
 , 
 job_name 
 : 
 str 
 , 
 project 
 : 
 str 
 , 
 bucket 
 : 
 str 
 , 
 location 
 : 
 str 
 = 
 "us-central1" 
 , 
 machine_type 
 : 
 str 
 = 
 "e2-highmem-2" 
 , 
 disk_size_gb 
 : 
 int 
 = 
 100 
 , 
 ) 
 - 
> None 
 : 
  
 """Launches a Vertex AI custom job to load the state dict. 
 If the model is too large to fit into memory or disk, we can launch 
 a Vertex AI custom job with a large enough VM for this to work. 
 Depending on the model's size, it might require a different VM 
 configuration. The model MUST fit into the VM's memory, and there 
 must be enough disk space to stage the entire model while it gets 
 copied to Cloud Storage. 
 Args: 
 model_name: HuggingFace model name compatible with AutoModelForSeq2SeqLM. 
 state_dict_path: File path to the model's state_dict, can be in Cloud Storage. 
 job_name: Job display name in the Vertex AI console. 
 project: Google Cloud Project ID. 
 bucket: Cloud Storage bucket name, without the "gs://" prefix. 
 location: Google Cloud regional location. 
 machine_type: Machine type for the VM to run the job. 
 disk_size_gb: Disk size in GB for the VM to run the job. 
 """ 
 from 
  
 google.cloud 
  
 import 
 aiplatform 
 aiplatform 
 . 
  init 
 
 ( 
 project 
 = 
 project 
 , 
 staging_bucket 
 = 
 bucket 
 , 
 location 
 = 
 location 
 ) 
 job 
 = 
 aiplatform 
 . 
  CustomJob 
 
 . 
 from_local_script 
 ( 
 display_name 
 = 
 job_name 
 , 
 container_uri 
 = 
 "us-docker.pkg.dev/vertex-ai/training/pytorch-gpu.1-13:latest" 
 , 
 script_path 
 = 
 "download_model.py" 
 , 
 args 
 = 
 [ 
 "local" 
 , 
 f 
 "--model-name= 
 { 
 model_name 
 } 
 " 
 , 
 f 
 "--state-dict-path= 
 { 
 state_dict_path 
 } 
 " 
 , 
 ], 
 machine_type 
 = 
 machine_type 
 , 
 boot_disk_size_gb 
 = 
 disk_size_gb 
 , 
 requirements 
 = 
 [ 
 "transformers" 
 ], 
 ) 
 job 
 . 
 run 
 () 
 if 
 __name__ 
 == 
 "__main__" 
 : 
 import 
  
 argparse 
 parser 
 = 
 argparse 
 . 
 ArgumentParser 
 () 
 subparsers 
 = 
 parser 
 . 
 add_subparsers 
 ( 
 required 
 = 
 True 
 ) 
 parser_local 
 = 
 subparsers 
 . 
 add_parser 
 ( 
 "local" 
 ) 
 parser_local 
 . 
 add_argument 
 ( 
 "--model-name" 
 , 
 required 
 = 
 True 
 , 
 help 
 = 
 "HuggingFace model name compatible with AutoModelForSeq2SeqLM" 
 , 
 ) 
 parser_local 
 . 
 add_argument 
 ( 
 "--state-dict-path" 
 , 
 required 
 = 
 True 
 , 
 help 
 = 
 "File path to the model's state_dict, can be in Cloud Storage" 
 , 
 ) 
 parser_local 
 . 
 set_defaults 
 ( 
 run 
 = 
 run_local 
 ) 
 parser_vertex 
 = 
 subparsers 
 . 
 add_parser 
 ( 
 "vertex" 
 ) 
 parser_vertex 
 . 
 add_argument 
 ( 
 "--model-name" 
 , 
 required 
 = 
 True 
 , 
 help 
 = 
 "HuggingFace model name compatible with AutoModelForSeq2SeqLM" 
 , 
 ) 
 parser_vertex 
 . 
 add_argument 
 ( 
 "--state-dict-path" 
 , 
 required 
 = 
 True 
 , 
 help 
 = 
 "File path to the model's state_dict, can be in Cloud Storage" 
 , 
 ) 
 parser_vertex 
 . 
 add_argument 
 ( 
 "--job-name" 
 , 
 required 
 = 
 True 
 , 
 help 
 = 
 "Job display name in the Vertex AI console" 
 ) 
 parser_vertex 
 . 
 add_argument 
 ( 
 "--project" 
 , 
 required 
 = 
 True 
 , 
 help 
 = 
 "Google Cloud Project ID" 
 ) 
 parser_vertex 
 . 
 add_argument 
 ( 
 "--bucket" 
 , 
 required 
 = 
 True 
 , 
 help 
 = 
 'Cloud Storage bucket name, without the "gs://" prefix' 
 , 
 ) 
 parser_vertex 
 . 
 add_argument 
 ( 
 "--location" 
 , 
 default 
 = 
 "us-central1" 
 , 
 help 
 = 
 "Google Cloud regional location" 
 ) 
 parser_vertex 
 . 
 add_argument 
 ( 
 "--machine-type" 
 , 
 default 
 = 
 "e2-highmem-2" 
 , 
 help 
 = 
 "Machine type for the VM to run the job" 
 , 
 ) 
 parser_vertex 
 . 
 add_argument 
 ( 
 "--disk-size-gb" 
 , 
 type 
 = 
 int 
 , 
 default 
 = 
 100 
 , 
 help 
 = 
 "Disk size in GB for the VM to run the job" 
 , 
 ) 
 parser_vertex 
 . 
 set_defaults 
 ( 
 run 
 = 
 run_vertex_job 
 ) 
 args 
 = 
 parser 
 . 
 parse_args 
 () 
 kwargs 
 = 
 args 
 . 
 __dict__ 
 . 
 copy 
 () 
 kwargs 
 . 
 pop 
 ( 
 "run" 
 ) 
 args 
 . 
 run 
 ( 
 ** 
 kwargs 
 ) 
 

  # Copyright 2023 Google LLC 
 # 
 # Licensed under the Apache License, Version 2.0 (the "License"); 
 # you may not use this file except in compliance with the License. 
 # You may obtain a copy of the License at 
 # 
 #      http://www.apache.org/licenses/LICENSE-2.0 
 # 
 # Unless required by applicable law or agreed to in writing, software 
 # distributed under the License is distributed on an "AS IS" BASIS, 
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 
 # See the License for the specific language governing permissions and 
 # limitations under the License. 
 """Runs a streaming RunInference Language Model pipeline.""" 
 from 
  
 __future__ 
  
 import 
 annotations 
 import 
  
 logging 
 import 
  
 apache_beam 
  
 as 
  
 beam 
 from 
  
 apache_beam.ml.inference.base 
  
 import 
 PredictionResult 
 from 
  
 apache_beam.ml.inference.base 
  
 import 
 RunInference 
 from 
  
 apache_beam.ml.inference.pytorch_inference 
  
 import 
 make_tensor_model_fn 
 from 
  
 apache_beam.ml.inference.pytorch_inference 
  
 import 
 PytorchModelHandlerTensor 
 from 
  
 apache_beam.options.pipeline_options 
  
 import 
 PipelineOptions 
 import 
  
 torch 
 from 
  
 transformers 
  
 import 
 AutoConfig 
 from 
  
 transformers 
  
 import 
 AutoModelForSeq2SeqLM 
 from 
  
 transformers 
  
 import 
 AutoTokenizer 
 from 
  
 transformers.tokenization_utils 
  
 import 
 PreTrainedTokenizer 
 MAX_RESPONSE_TOKENS 
 = 
 256 
 def 
  
 to_tensors 
 ( 
 input_text 
 : 
 str 
 , 
 tokenizer 
 : 
 PreTrainedTokenizer 
 ) 
 - 
> torch 
 . 
 Tensor 
 : 
  
 """Encodes input text into token tensors. 
 Args: 
 input_text: Input text for the language model. 
 tokenizer: Tokenizer for the language model. 
 Returns: Tokenized input tokens. 
 """ 
 return 
 tokenizer 
 ( 
 input_text 
 , 
 return_tensors 
 = 
 "pt" 
 ) 
 . 
 input_ids 
 [ 
 0 
 ] 
 def 
  
 decode_response 
 ( 
 result 
 : 
 PredictionResult 
 , 
 tokenizer 
 : 
 PreTrainedTokenizer 
 ) 
 - 
> str 
 : 
  
 """Decodes output token tensors into text. 
 Args: 
 result: Prediction results from the RunInference transform. 
 tokenizer: Tokenizer for the language model. 
 Returns: The model's response as text. 
 """ 
 output_tokens 
 = 
 result 
 . 
 inference 
 return 
 tokenizer 
 . 
 decode 
 ( 
 output_tokens 
 , 
 skip_special_tokens 
 = 
 True 
 ) 
 class 
  
 AskModel 
 ( 
 beam 
 . 
 PTransform 
 ): 
  
 """Asks an language model a prompt message and gets its responses. 
 Attributes: 
 model_name: HuggingFace model name compatible with AutoModelForSeq2SeqLM. 
 state_dict_path: File path to the model's state_dict, can be in Cloud Storage. 
 max_response_tokens: Maximum number of tokens for the model to generate. 
 """ 
 def 
  
 __init__ 
 ( 
 self 
 , 
 model_name 
 : 
 str 
 , 
 state_dict_path 
 : 
 str 
 , 
 max_response_tokens 
 : 
 int 
 = 
 MAX_RESPONSE_TOKENS 
 , 
 ) 
 - 
> None 
 : 
 self 
 . 
 model_handler 
 = 
 PytorchModelHandlerTensor 
 ( 
 state_dict_path 
 = 
 state_dict_path 
 , 
 model_class 
 = 
 AutoModelForSeq2SeqLM 
 . 
 from_config 
 , 
 model_params 
 = 
 { 
 "config" 
 : 
 AutoConfig 
 . 
 from_pretrained 
 ( 
 model_name 
 )}, 
 inference_fn 
 = 
 make_tensor_model_fn 
 ( 
 "generate" 
 ), 
 ) 
 self 
 . 
 tokenizer 
 = 
 AutoTokenizer 
 . 
 from_pretrained 
 ( 
 model_name 
 ) 
 self 
 . 
 max_response_tokens 
 = 
 max_response_tokens 
 def 
  
 expand 
 ( 
 self 
 , 
 pcollection 
 : 
 beam 
 . 
 PCollection 
 [ 
 str 
 ]) 
 - 
> beam 
 . 
 PCollection 
 [ 
 str 
 ]: 
 return 
 ( 
 pcollection 
 | 
 "To tensors" 
>> beam 
 . 
 Map 
 ( 
 to_tensors 
 , 
 self 
 . 
 tokenizer 
 ) 
 | 
 "RunInference" 
>> RunInference 
 ( 
 self 
 . 
 model_handler 
 , 
 inference_args 
 = 
 { 
 "max_new_tokens" 
 : 
 self 
 . 
 max_response_tokens 
 }, 
 ) 
 | 
 "Get response" 
>> beam 
 . 
 Map 
 ( 
 decode_response 
 , 
 self 
 . 
 tokenizer 
 ) 
 ) 
 if 
 __name__ 
 == 
 "__main__" 
 : 
 import 
  
 argparse 
 parser 
 = 
 argparse 
 . 
 ArgumentParser 
 () 
 parser 
 . 
 add_argument 
 ( 
 "--messages-topic" 
 , 
 required 
 = 
 True 
 , 
 help 
 = 
 "Pub/Sub topic for input text messages" 
 , 
 ) 
 parser 
 . 
 add_argument 
 ( 
 "--responses-topic" 
 , 
 required 
 = 
 True 
 , 
 help 
 = 
 "Pub/Sub topic for output text responses" 
 , 
 ) 
 parser 
 . 
 add_argument 
 ( 
 "--model-name" 
 , 
 required 
 = 
 True 
 , 
 help 
 = 
 "HuggingFace model name compatible with AutoModelForSeq2SeqLM" 
 , 
 ) 
 parser 
 . 
 add_argument 
 ( 
 "--state-dict-path" 
 , 
 required 
 = 
 True 
 , 
 help 
 = 
 "File path to the model's state_dict, can be in Cloud Storage" 
 , 
 ) 
 args 
 , 
 beam_args 
 = 
 parser 
 . 
 parse_known_args 
 () 
 logging 
 . 
 getLogger 
 () 
 . 
 setLevel 
 ( 
 logging 
 . 
 INFO 
 ) 
 beam_options 
 = 
 PipelineOptions 
 ( 
 beam_args 
 , 
 pickle_library 
 = 
 "cloudpickle" 
 , 
 streaming 
 = 
 True 
 , 
 ) 
 simple_name 
 = 
 args 
 . 
 model_name 
 . 
 split 
 ( 
 "/" 
 )[ 
 - 
 1 
 ] 
 pipeline 
 = 
 beam 
 . 
 Pipeline 
 ( 
 options 
 = 
 beam_options 
 ) 
 _ 
 = 
 ( 
 pipeline 
 | 
 "Read from Pub/Sub" 
>> beam 
 . 
 io 
 . 
 ReadFromPubSub 
 ( 
 args 
 . 
 messages_topic 
 ) 
 | 
 "Decode bytes" 
>> beam 
 . 
 Map 
 ( 
 lambda 
 msg 
 : 
 msg 
 . 
 decode 
 ( 
 "utf-8" 
 )) 
 | 
 f 
 "Ask 
 { 
 simple_name 
 } 
 " 
>> AskModel 
 ( 
 args 
 . 
 model_name 
 , 
 args 
 . 
 state_dict_path 
 ) 
 | 
 "Encode bytes" 
>> beam 
 . 
 Map 
 ( 
 lambda 
 msg 
 : 
 msg 
 . 
 encode 
 ( 
 "utf-8" 
 )) 
 | 
 "Write to Pub/Sub" 
>> beam 
 . 
 io 
 . 
 WriteToPubSub 
 ( 
 args 
 . 
 responses_topic 
 ) 
 ) 
 pipeline 
 . 
 run 
 ()