human
I am working on the file "09_deploying/09a_inmemory.ipynb".
The first task is:
In memory model (Beam pipeline)
Make sure to share the model so that it is not loaded for each element.
There will be a different model on different worker machines.
Can you write Python program for it?
assistant
#Here is the Python program you requested.

import apache_beam as beam
from apache_beam.utils.shared import Shared
import tensorflow as tf

class ModelPredict:
    def __init__(self, shared_handle, model_location):
        self._shared_handle = shared_handle
        self._model_location = model_location 

    def __call__(self, filenames):
        def initialize_model():
            print('Loading Keras model from ' + self._model_location)
            return (tf.keras.models.load_model(self._model_location)
                    .signatures['serving_default'])
        
        serving_fn = self._shared_handle.acquire(initialize_model)
        if isinstance(filenames, str):
            # only one element, put it into a batch of 1
            result = serving_fn(tf.convert_to_tensor([filenames]))
        else:
            # a list
            result = serving_fn(tf.convert_to_tensor(filenames))
        return {
            'filenames': filenames,
            'probability': result['probability'].numpy(),
            'pred_label': result['flower_type_str'].numpy()
        }

with beam.Pipeline() as p:
    
    shared_handle = Shared()
    
    (p 
     | 'input' >> beam.Create([
        'gs://practical-ml-vision-book/flowers_5_jpeg/flower_photos/dandelion/9818247_e2eac18894.jpg',
        'gs://practical-ml-vision-book/flowers_5_jpeg/flower_photos/dandelion/9853885425_4a82356f1d_m.jpg',
        'gs://practical-ml-vision-book/flowers_5_jpeg/flower_photos/daisy/9299302012_958c70564c_n.jpg',
        'gs://practical-ml-vision-book/flowers_5_jpeg/flower_photos/tulips/8733586143_3139db6e9e_n.jpg',
        'gs://practical-ml-vision-book/flowers_5_jpeg/flower_photos/tulips/8713397358_0505cc0176_n.jpg'
     ]) 
     | 'pred' >> beam.Map(ModelPredict(shared_handle, MODEL_LOCATION))
     | 'write' >> beam.io.textio.WriteToText('/tmp/out.txt', num_shards=1)
    )

!cat /tmp/out.txt*