import os

os.environ["KERAS_BACKEND"] = "tensorflow"

import keras
import keras_hub


PRESET = "gpt2_base_en"
PROMPT = "KerasHub lets developers"


preprocessor = keras_hub.models.GPT2CausalLMPreprocessor.from_preset(
    PRESET,
    sequence_length=64,
)
model = keras_hub.models.GPT2CausalLM.from_preset(
    PRESET,
    preprocessor=preprocessor,
)
model.compile(sampler="greedy")

completion = model.generate(PROMPT, max_length=40, strip_prompt=True)

print(f"backend: {keras.config.backend()}")
print(f"preset: {PRESET}")
print(f"prompt: {PROMPT}")
print(f"completion: {completion.strip()}")