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
importos
importtorch
fromtransformersimport AutoModelForSeq2SeqLM
defrun_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!")
defrun_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.
 """
 fromgoogle.cloudimport 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__":
 importargparse
 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
importlogging
importapache_beamasbeam
fromapache_beam.ml.inference.baseimport PredictionResult
fromapache_beam.ml.inference.baseimport RunInference
fromapache_beam.ml.inference.pytorch_inferenceimport make_tensor_model_fn
fromapache_beam.ml.inference.pytorch_inferenceimport PytorchModelHandlerTensor
fromapache_beam.options.pipeline_optionsimport PipelineOptions
importtorch
fromtransformersimport AutoConfig
fromtransformersimport AutoModelForSeq2SeqLM
fromtransformersimport AutoTokenizer
fromtransformers.tokenization_utilsimport PreTrainedTokenizer
MAX_RESPONSE_TOKENS = 256
defto_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]
defdecode_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)
classAskModel(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
 defexpand(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__":
 importargparse
 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()