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

model_file = "modules/model.py"
model_fn = "modules.model.run_fn"
tuner_fn = "modules.model.tuner_fn"

proprocessing_file = "modules/preprocessing.py"
preprocessing_fn = "modules.preprocessing.preprocessing_fn"

schema_file = "schema.pbtxt"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 StatisticsGen
from tfx.components import SchemaGen
from tfx.components import ExampleValidator
from tfx.components import Transform
from tfx.components import Trainer
from tfx.components import Tuner
from tfx.components import Evaluator
from tfx.components import Pusher
from tfx.proto import example_gen_pb2

from tfx.types import Channel
from tfx.types.standard_artifacts import Model
from tfx.types.standard_artifacts import ModelBlessing
from tfx.dsl.components.common import resolver
from tfx.dsl.experimental.latest_blessed_model_resolver import LatestBlessedModelResolver

import tensorflow_model_analysis as tfmacontext = 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)statistics_gen = StatisticsGen(examples=example_gen.outputs['examples'])
context.run(statistics_gen)context.show(statistics_gen.outputs['statistics'])%%writefile {schema_file}

feature {
  name: "image"
  type: FLOAT

  presence {
    min_fraction: 1.0
  }

  float_domain {
    min: 0
    max: 255
  }

  shape {
    dim { size: 256 }
    dim { size: 256 }
    dim { size: 3 }
  }
}

feature {
  name: "image_shape"
  type: INT

  presence {
    min_fraction: 1.0
  }

  shape {
    dim { size: 3 }
  }
}

feature {
  name: "label"
  type: INT

  presence {
    min_fraction: 1.0
  }

  int_domain {
    min: 0
    max: 2
  }

  shape {
    dim { size: 1 }
  }  
}schema_gen = tfx.components.ImportSchemaGen(
    schema_file=schema_file)

context.run(schema_gen)example_validator = ExampleValidator(
  statistics=statistics_gen.outputs['statistics'],
  schema=schema_gen.outputs['schema']   
)

context.run(example_validator)context.show(example_validator.outputs['anomalies'])%%writefile {proprocessing_file}

import tensorflow as tf

IMAGE_KEY = "image"
LABEL_KEY = "label"

MODEL_INPUT_IMAGE_KEY = "pixel_values"
MODEL_INPUT_LABEL_KEY = "labels"

INPUT_IMG_SIZE = 224

def preprocessing_fn(inputs):
    """tf.transform's callback function for preprocessing inputs.
    Args:
      inputs: map from feature keys to raw not-yet-transformed features.
    Returns:
      Map from string feature key to transformed feature operations.
    """
    # print(inputs)
    outputs = {}

    inputs[IMAGE_KEY] = tf.image.resize(
        inputs[IMAGE_KEY], [INPUT_IMG_SIZE, INPUT_IMG_SIZE]
    )

    inputs[IMAGE_KEY] = inputs[IMAGE_KEY] / 255.0
    inputs[IMAGE_KEY] = tf.transpose(inputs[IMAGE_KEY], [0, 3, 1, 2])

    outputs[MODEL_INPUT_IMAGE_KEY] = inputs[IMAGE_KEY]
    outputs[MODEL_INPUT_LABEL_KEY] = inputs[LABEL_KEY]

    return outputstransform = Transform(
    examples=example_gen.outputs['examples'],
    schema=schema_gen.outputs['schema'],
    preprocessing_fn=preprocessing_fn)context.run(transform)%%writefile {model_file}

from typing import List, Dict, Tuple

import absl
import tensorflow as tf
import keras_tuner
import tensorflow_transform as tft
from transformers import ViTFeatureExtractor, TFViTForImageClassification
from tfx.components.trainer.fn_args_utils import FnArgs
from tfx.v1.components import TunerFnResult

from tfx_bsl.tfxio import dataset_options
from tfx.components.trainer.fn_args_utils import DataAccessor

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"

_MODEL_INPUT_LABEL_KEY = "labels"

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 _transform_features_signature(
    model: tf.keras.Model, tf_transform_output: tft.TFTransformOutput
):
    """
    transform_features_signature simply returns a function that transforms
    any data of the type of tf.Example which is denoted as the type of sta
    ndard_artifacts.Examples in TFX. The purpose of this function is to ap
    ply Transform Graph obtained from Transform component to the data prod
    uced by ImportExampleGen. This function will be used in the Evaluator
    component, so the raw evaluation inputs from ImportExampleGen can be a
    pporiately transformed that the model could understand.
    """

    # basically, what Transform component emits is a SavedModel that knows
    # how to transform data. transform_features_layer() simply returns the
    # layer from the Transform.
    model.tft_layer = tf_transform_output.transform_features_layer()

    @tf.function(
        input_signature=[tf.TensorSpec(shape=[None], dtype=tf.string, name="examples")]
    )
    def serve_tf_examples_fn(serialized_tf_examples):
        """
        raw_feature_spec returns a set of feature maps(dict) for the input
        TFRecords based on the knowledge that Transform component has lear
        ned(learn doesn't mean training here). By using this information,
        the raw data from ImportExampleGen could be parsed with tf.io.parse
        _example utility function.
        Then, it is passed to the model.tft_layer, so the final output we
        get is the transformed data of the raw input.
        """
        feature_spec = tf_transform_output.raw_feature_spec()
        parsed_features = tf.io.parse_example(serialized_tf_examples, feature_spec)

        transformed_features = model.tft_layer(parsed_features)

        return transformed_features

    return serve_tf_examples_fn


def _tf_examples_serving_signature(model, tf_transform_output):
    """
    tf_examples_serving_signature simply returns a function that performs
    data transformation(preprocessing) and model prediction in a sequential
    manner. How data transformation is done is idential to the process of
    transform_features_signature function.
    """

    model.tft_layer = tf_transform_output.transform_features_layer()

    @tf.function(
        input_signature=[tf.TensorSpec(shape=[None], dtype=tf.string, name="examples")]
    )
    def serve_tf_examples_fn(
        serialized_tf_example: tf.Tensor,
    ) -> Dict[str, tf.Tensor]:
        raw_feature_spec = tf_transform_output.raw_feature_spec()
        raw_features = tf.io.parse_example(serialized_tf_example, raw_feature_spec)

        transformed_features = model.tft_layer(raw_features)
        logits = model(transformed_features).logits

        return {_MODEL_INPUT_LABEL_KEY: logits}

    return serve_tf_examples_fn

###

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

    dataset = data_accessor.tf_dataset_factory(
        file_pattern,
        dataset_options.TensorFlowDatasetOptions(
            batch_size=batch_size, label_key=_MODEL_INPUT_LABEL_KEY
        ),
        tf_transform_output.transformed_metadata.schema,
    )

    return dataset

def _get_hyperparameters() -> keras_tuner.HyperParameters:
    hp = keras_tuner.HyperParameters()
    hp.Choice("learning_rate", [1e-3, 1e-2], default=1e-3)
    return hp

def _build_model(hparams: keras_tuner.HyperParameters):
  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,
  )

  model.layers[0].trainable=False

  loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
  optimizer = tf.keras.optimizers.Adam(learning_rate=hparams.get("learning_rate"))
  model.compile(optimizer=optimizer, loss=loss, metrics=["accuracy"])

  return model  

def run_fn(fn_args: FnArgs):
    tf_transform_output = tft.TFTransformOutput(fn_args.transform_output)

    train_dataset = _input_fn(
        fn_args.train_files,
        fn_args.data_accessor,
        tf_transform_output,
        is_train=True,
        batch_size=_TRAIN_BATCH_SIZE,
    )

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

    hparams = keras_tuner.HyperParameters.from_config(fn_args.hyperparameters)
    INFO(f"HyperParameters for training: {hparams.get_config()}")
    model = _build_model(hparams)

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

    model.save(
        fn_args.serving_model_dir, 
        save_format="tf", 
        signatures={
            "serving_default": _model_exporter(model),
            "transform_features": _transform_features_signature(
                model, tf_transform_output
            ),
            "from_examples": _tf_examples_serving_signature(model, tf_transform_output),
        },
    )

def tuner_fn(fn_args: FnArgs) -> TunerFnResult:
    tuner = keras_tuner.RandomSearch(
        _build_model,
        max_trials=6,
        hyperparameters=_get_hyperparameters(),
        allow_new_entries=False,
        objective=keras_tuner.Objective("val_accuracy", "max"),
        directory=fn_args.working_dir,
        project_name="img_classification_tuning",
    )

    tf_transform_output = tft.TFTransformOutput(fn_args.transform_graph_path)

    train_dataset = _input_fn(
        fn_args.train_files,
        fn_args.data_accessor,
        tf_transform_output,
        is_train=True,
        batch_size=_TRAIN_BATCH_SIZE,
    )

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

    return TunerFnResult(
        tuner=tuner,
        fit_kwargs={
            "x": train_dataset,
            "validation_data": eval_dataset,
            "steps_per_epoch": _TRAIN_LENGTH // _TRAIN_BATCH_SIZE,
            "validation_steps": _EVAL_LENGTH // _EVAL_BATCH_SIZE,
        },
    )
tuner = Tuner(
    tuner_fn=tuner_fn,
    examples=transform.outputs["transformed_examples"],
    schema=schema_gen.outputs["schema"],
    transform_graph=transform.outputs["transform_graph"],
)context.run(tuner)trainer = Trainer(
    run_fn=model_fn,
    transformed_examples=transform.outputs["transformed_examples"],
    transform_graph=transform.outputs["transform_graph"],
    schema=schema_gen.outputs["schema"],
    hyperparameters=tuner.outputs["best_hyperparameters"],
)context.run(trainer)model_resolver = resolver.Resolver(
    strategy_class=LatestBlessedModelResolver,
    model=Channel(type=Model),
    model_blessing=Channel(type=ModelBlessing),
).with_id("latest_blessed_model_resolver")context.run(model_resolver)eval_configs = tfma.EvalConfig(
    model_specs=[
        tfma.ModelSpec(
            signature_name="from_examples",
            preprocessing_function_names=["transform_features"],
            label_key="labels",
            prediction_key="labels",
        )
    ],
    slicing_specs=[tfma.SlicingSpec()],
    metrics_specs=[
        tfma.MetricsSpec(
            metrics=[
                tfma.MetricConfig(
                    class_name="SparseCategoricalAccuracy",
                    threshold=tfma.MetricThreshold(
                        value_threshold=tfma.GenericValueThreshold(
                            lower_bound={"value": 0.55}
                        ),
                        # Change threshold will be ignored if there is no
                        # baseline model resolved from MLMD (first run).
                        change_threshold=tfma.GenericChangeThreshold(
                            direction=tfma.MetricDirection.HIGHER_IS_BETTER,
                            absolute={"value": -1e-3},
                        ),
                    ),
                )
            ]
        )
    ],
)


evaluator = Evaluator(
    examples=example_gen.outputs["examples"],
    model=trainer.outputs["model"],
    baseline_model=model_resolver.outputs["model"],
    eval_config=eval_configs,
)context.run(evaluator)