generate.py

#
# -*- coding: utf-8 -*-
#
# Copyright (c) 2024 Intel Corporation
#
# 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.
#

# -*- coding: utf-8 -*-
"""get_started.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/github/google/generative-ai-docs/blob/main/site/en/gemma/docs/get_started.ipynb

##### Copyright 2024 Google LLC.
"""

#@title 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
#
# https://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.


import argparse
import keras
import keras_nlp
import os
import time

parser = argparse.ArgumentParser("Sampling task with GEMMA using FLAX")
parser.add_argument(
    "-p", "--precision", required=True, type=str,
    choices=['fp32', 'fp16', 'bfloat16'],
    help="precision for the model",
)
parser.add_argument(
    "-m", "--model_path", required=True, type=str,
    help="path to the Keras model weights & tokenizer folder"
)
parser.add_argument(
    "-b", "--keras_backend", required=True, type=str,
    choices=['tensorflow', 'jax'],
    help="keras backend to use for the model",
)
parser.add_argument(
    "-l", "--max_length", default=64, type=int,
    help="maximum length of the generated text",
)
parser.add_argument(
    "-s", "--batch_size", default=128, type=int,
    help="input batch size",
)
args = parser.parse_args()

os.environ["KERAS_BACKEND"] = args.keras_backend
print("Using Keras backend: " + args.keras_backend)
print("Configured max_length: " + str(args.max_length))

if args.precision == "fp16":
    print("Running with " + args.precision + " precision. Setting keras mixed precision policy.")
    keras.mixed_precision.set_global_policy("mixed_float16")
elif args.precision == "bfloat16":
    print("Running with " + args.precision + " precision. Setting keras mixed precision policy.")
    keras.mixed_precision.set_global_policy("mixed_bfloat16")

gemma_lm = keras_nlp.models.GemmaCausalLM.from_preset(args.model_path)

print(gemma_lm.summary())

start = time.time()
gemma_lm.generate("Warming up the model", max_length=args.max_length)
print("Time taken for first generate (warmup):" + str(time.time() - start) + " seconds")

input_query = ["What is the meaning of life?"]
input_batch = input_query * args.batch_size

start = time.time()
out_data = gemma_lm.generate(input_batch, max_length=args.max_length)
end = time.time()
print("Time taken for second generate:" + str(end - start) + " seconds")

latency = (end - start)
print("Latency: {} sec".format(latency))
print("Throughput: {} inputs/sec".format(
    len(input_batch) / latency
))

# print("\n\nText generated by GEMMA for the input prompts:")
# for out_string in out_data:
#     print(100*'#')
#     print(out_string)