Generative TF

import tensorflow as tf
from tensorflow import keras
import numpy as np

# course-edition deterministic TensorFlow news fixture v1
_course_titles = tf.constant([
    "Regional leaders agree on a public health plan",
    "The home team wins the basketball final",
    "Markets rise after a strong business report",
    "Researchers release an efficient intelligence model",
    "Diplomats begin a new round of peace talks",
    "The tennis champion returns after an injury",
    "Retail sales improve with consumer confidence",
    "A software update improves privacy",
])
_course_descriptions = tf.constant([
    "World policy and international cooperation",
    "Sports results and player development",
    "Business revenue and economic indicators",
    "Science and technology research",
    "World diplomacy and humanitarian aid",
    "Sports training and competition",
    "Business investment and employment",
    "Science security and technology",
])
_course_labels = tf.constant([0, 1, 2, 3, 0, 1, 2, 3], dtype=tf.int64)
_course_records = {"title": _course_titles, "description": _course_descriptions, "label": _course_labels}
ds_train = tf.data.Dataset.from_tensor_slices(_course_records).take(6)
ds_test = tf.data.Dataset.from_tensor_slices(_course_records).skip(6)

titles = ds_train.map(lambda item: item["title"])
vectorizer = keras.layers.TextVectorization(
    standardize=None,
    split="character",
    output_mode="int",
    ragged=True,
)
vectorizer.adapt(titles.batch(256))

vocabulary = vectorizer.get_vocabulary()
eos_token = len(vocabulary)
vocab_size = eos_token + 1
reverse_map = {index: token for index, token in enumerate(vocabulary)}
reverse_map[eos_token] = "<eos>"

def encode(text):
    return vectorizer(tf.constant([text])).values.numpy().tolist()

encode("Hello, world!")

def title_batch_fn(batch):
    encoded = vectorizer(batch["title"]).to_tensor(default_value=0)
    eos_column = tf.fill([tf.shape(encoded)[0], 1], tf.cast(eos_token, encoded.dtype))
    targets = tf.concat([encoded[:, 1:], eos_column], axis=1)
    return tf.one_hot(encoded, vocab_size), tf.one_hot(targets, vocab_size)

training_ds = (
    ds_train.shuffle(10_000, seed=2026)
    .batch(8)
    .map(title_batch_fn, num_parallel_calls=tf.data.AUTOTUNE)
    .prefetch(tf.data.AUTOTUNE)
)

model = keras.models.Sequential([
    keras.Input(shape=(None, vocab_size)),
    keras.layers.Masking(mask_value=0.0),
    keras.layers.LSTM(128, return_sequences=True),
    keras.layers.Dense(vocab_size, activation="softmax"),
])

model.summary()
model.compile(optimizer=keras.optimizers.Adam(), loss="categorical_crossentropy")
model.fit(training_ds, epochs=1)

def decode(token_ids):
    return "".join(reverse_map.get(int(token_id), "�") for token_id in token_ids if int(token_id) != 0)

def generate(model, size=100, start="Today "):
    tokens = encode(start)
    generated = list(tokens)
    for _ in range(size):
        inputs = tf.one_hot(tokens, vocab_size)[None, ...]
        probabilities = model(inputs, training=False)[0, -1]
        next_token = int(tf.argmax(probabilities))
        if next_token == eos_token:
            break
        generated.append(next_token)
        tokens.append(next_token)
    return decode(generated)

generate(model)

sampling_callback = keras.callbacks.LambdaCallback(
    on_epoch_end=lambda _epoch, _logs: print(generate(model))
)

model.fit(training_ds, callbacks=[sampling_callback], epochs=3)

def generate_soft(model, size=100, start="Today ", temperature=1.0):
    tokens = encode(start)
    generated = list(tokens)
    for _ in range(size):
        inputs = tf.one_hot(tokens, vocab_size)[None, ...]
        logits = tf.math.log(tf.clip_by_value(model(inputs, training=False)[0, -1], 1e-8, 1.0))
        next_token = int(tf.random.categorical(logits[None, :] / temperature, 1, seed=2026)[0, 0])
        if next_token == eos_token:
            break
        generated.append(next_token)
        tokens.append(next_token)
    return decode(generated)

prompts = ["Today ", "On Sunday ", "Riyadh, ", "President ", "Once upon a time "]
for temperature in [0.3, 0.8, 1.0, 1.3]:
    print(f"\n--- Temperature = {temperature}")
    for prompt in prompts:
        print(generate_soft(model, size=160, start=prompt, temperature=temperature))