Is my code for implementing a custom model correct?

[Mahran-xo]

Hello,

I am in the process of implementing an audio classifier intended for production use. However, I am encountering challenges with the logic, as I am relatively new to AutoML and MLOps. I would greatly appreciate any tips or guidance on the following code. Please excuse any inconsistencies in the logic, as I am still learning and refining the implementation.

import autokeras as ak
import train_utils as cb
from train_utils import (
    stereo_to_mono_converter,
    squeeze,
    get_spectrogram
)


class SoundClf():
    def build(self,hp,inputs):
        # Apply your data processing pipeline to the inputs
        x, y = inputs
        x = stereo_to_mono_converter(x, y)
        x = squeeze(x, y)
        x = get_spectrogram(x)
        input = ak.ImageInput()
        resize = cb.ResizingBlock()(input)
        norm_layer = ak.Normalization()(resize)
        conv1 = ak.ConvBlock()(norm_layer)
        conv2 = ak.ConvBlock()(conv1)
        conv3 = ak.ConvBlock()(conv2)
        res = ak.ResNetBlock(version="v2")(norm_layer)
        merge = ak.Merge()[conv3,res]
        conv4 = ak.ConvBlock()(merge)
        output = ak.ClassificationHead()(conv4)

        auto_model = ak.AutoModel(
            inputs=input, outputs=output, overwrite=True, max_trials=1
        )
        return auto_model(x)

the functions from my custom module train_utils are as follows:

Resizing Block

class ResizingBlock(ak.Block):
    """
    Resizing Block implemented from keras.layers.Resizing
    """
    def build(self, hp, inputs=None):
        # Get the input_node from inputs.
        input_node = inputs[0]
        layer = keras.layers.Resizing(
            target_height=hp.Int("height", min_value=32, max_value=1024, step=16),
            target_width=hp.Int("width", min_value=32, max_value=1024, step=16)
        )
        output_node = layer(input_node)
        return output_node

The rest of the functions

seed = 42
tf.random.set_seed(seed)
np.random.seed(seed)

def stereo_to_mono_converter(example,labels):
    audio = example
    # If it has multiple channels, take the mean to convert to mono
    audio = tf.reduce_mean(audio, axis=-1, keepdims=True)
    # Add any additional preprocessing steps here
    return audio,labels

def squeeze(audio, labels):
  audio = tf.squeeze(audio, axis=-1)
  return audio, labels


def get_spectrogram(waveform):
  spectrogram = tf.signal.stft(
      waveform, frame_length=255, frame_step=128)

  spectrogram = tf.abs(spectrogram)
 
  spectrogram = spectrogram[..., tf.newaxis]
  return spectrogram

def make_spec_ds(ds):
  return ds.map(
      map_func=lambda audio,label: (get_spectrogram(audio), label),
      num_parallel_calls=tf.data.AUTOTUNE)