data_path = "gs://beans-lowres/tfrecords"
local_data_path = "data"

model_file = "modules/model.py"
model_fn = "modules.model.run_fn"import tfx
tfx.__version__from tfx import v1 as tfx
from tfx.orchestration.experimental.interactive.interactive_context import InteractiveContext
from tfx.components import ImportExampleGen
from tfx.components import Trainer
from tfx.components import Pusher
from tfx.components import StatisticsGen
from tfx.proto import example_gen_pb2context = InteractiveContext()input_config = example_gen_pb2.Input(
    splits=[
        example_gen_pb2.Input.Split(name="train", pattern="train-*.tfrec"),
        example_gen_pb2.Input.Split(name="eval", pattern="val-*.tfrec"),
    ]
)
example_gen = ImportExampleGen(
    input_base=local_data_path, 
    input_config=input_config
)context.run(example_gen)%%writefile {model_file}

from typing import List, Dict, Tuple

import absl
import tensorflow as tf
from transformers import ViTFeatureExtractor, TFViTForImageClassification
from tfx.components.trainer.fn_args_utils import FnArgs

feature_extractor = ViTFeatureExtractor()

_TRAIN_LENGTH = 128
_EVAL_LENGTH = 128
_TRAIN_BATCH_SIZE = 8
_EVAL_BATCH_SIZE = 8

_EPOCHS = 1

_LABELS = ['angular_leaf_spot', 'bean_rust', 'healthy']

_CONCRETE_INPUT = "pixel_values"

def INFO(text: str):
    absl.logging.info(text)

###

def _normalize_img(
    img, mean=feature_extractor.image_mean, std=feature_extractor.image_std
):
    img = img / 255
    mean = tf.constant(mean)
    std = tf.constant(std)
    return (img - mean) / std


def _preprocess_serving(string_input):
    decoded_input = tf.io.decode_base64(string_input)
    decoded = tf.io.decode_jpeg(decoded_input, channels=3)
    resized = tf.image.resize(decoded, size=(224, 224))
    normalized = _normalize_img(resized)
    normalized = tf.transpose(
        normalized, (2, 0, 1)
    )  # Since HF models are channel-first.
    return normalized


@tf.function(input_signature=[tf.TensorSpec([None], tf.string)])
def _preprocess_fn(string_input):
    decoded_images = tf.map_fn(
        _preprocess_serving, string_input, dtype=tf.float32, back_prop=False
    )
    return {_CONCRETE_INPUT: decoded_images}


def _model_exporter(model: tf.keras.Model):
    m_call = tf.function(model.call).get_concrete_function(
        tf.TensorSpec(
            shape=[None, 3, 224, 224], dtype=tf.float32, name=_CONCRETE_INPUT
        )
    )

    @tf.function(input_signature=[tf.TensorSpec([None], tf.string)])
    def serving_fn(string_input):
        labels = tf.constant(list(model.config.id2label.values()), dtype=tf.string)
        images = _preprocess_fn(string_input)

        predictions = m_call(**images)
        indices = tf.argmax(predictions.logits, axis=1)
        pred_source = tf.gather(params=labels, indices=indices)
        probs = tf.nn.softmax(predictions.logits, axis=1)
        pred_confidence = tf.reduce_max(probs, axis=1)
        return {"label": pred_source, "confidence": pred_confidence}

    return serving_fn

###

def _parse_tfr(proto):
    feature_description = {
        "image": tf.io.VarLenFeature(tf.float32),
        "image_shape": tf.io.VarLenFeature(tf.int64),
        "label": tf.io.VarLenFeature(tf.int64),
    }
    rec = tf.io.parse_single_example(proto, feature_description)
    image_shape = tf.sparse.to_dense(rec["image_shape"])
    image = tf.reshape(tf.sparse.to_dense(rec["image"]), image_shape)
    label = tf.sparse.to_dense(rec["label"])
    return {"pixel_values": image, "labels": label}

def _preprocess(example_batch):
    images = example_batch["pixel_values"]
    images = tf.transpose(images, perm=[0, 1, 2, 3]) # (batch_size, height, width, num_channels)
    images = tf.image.resize(images, (224, 224))
    images = tf.transpose(images, perm=[0, 3, 1, 2])

    labels = example_batch["labels"]
    labels = tf.transpose(labels, perm=[0, 1]) # So, that TF can evaluation the shapes.
    
    return {"pixel_values": images, "labels": labels}

def _input_fn(
    file_pattern: List[str],
    batch_size: int = 32,
    is_train: bool = False,
) -> tf.data.Dataset:
    INFO(f"Reading data from: {file_pattern}")

    dataset = tf.data.TFRecordDataset(
        tf.io.gfile.glob(file_pattern[0] + ".gz"),
        num_parallel_reads=tf.data.AUTOTUNE,
        compression_type="GZIP",
    ).map(_parse_tfr, num_parallel_calls=tf.data.AUTOTUNE)

    if is_train:
        dataset = dataset.shuffle(batch_size * 2)

    dataset = dataset.batch(batch_size)
    dataset = dataset.prefetch(tf.data.AUTOTUNE)
    dataset = dataset.map(_preprocess)
    return dataset

def _build_model():
  id2label={str(i): c for i, c in enumerate(_LABELS)}
  label2id={c: str(i) for i, c in enumerate(_LABELS)}

  model = TFViTForImageClassification.from_pretrained(
    "google/vit-base-patch16-224-in21k",
    num_labels=len(_LABELS),
    label2id=label2id,
    id2label=id2label,
  )

  loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
  model.compile(optimizer='adam', loss=loss, metrics=["accuracy"])

  return model  

def run_fn(fn_args: FnArgs):
    train_dataset = _input_fn(
        fn_args.train_files,
        is_train=True,
        batch_size=_TRAIN_BATCH_SIZE,
    )

    eval_dataset = _input_fn(
        fn_args.eval_files,
        is_train=False,
        batch_size=_EVAL_BATCH_SIZE,
    )

    model = _build_model()

    model.fit(
        train_dataset,
        steps_per_epoch=_TRAIN_LENGTH // _TRAIN_BATCH_SIZE,
        validation_data=eval_dataset,
        validation_steps=_EVAL_LENGTH // _TRAIN_BATCH_SIZE,
        epochs=_EPOCHS,
    )

    model.save(
        fn_args.serving_model_dir, save_format="tf", signatures=_model_exporter(model)
    )trainer = Trainer(
    run_fn=model_fn,
    examples=example_gen.outputs["examples"],
)context.run(trainer)