SO question 59697878

script.py

import argparse
import logging

import apache_beam as beam

from apache_beam.options.pipeline_options import SetupOptions
from apache_beam.options.pipeline_options import PipelineOptions


def prediction_handler(element):
    logging.info("Hi")
    yield element


def main(argv=None):

    parser = argparse.ArgumentParser()
    known_args, pipeline_args = parser.parse_known_args(argv)

    pipeline_options = PipelineOptions(pipeline_args)
    pipeline_options.view_as(SetupOptions).save_main_session = True
    
    p = beam.Pipeline(options=pipeline_options)

    output_tweets = (p
                     | 'start' >> beam.Create(['Start'])
                     | 'predict sentiment' >> beam.FlatMap(
            lambda messages: prediction_handler(messages))
                     )

    result = p.run()
    result.wait_until_finish()


if __name__ == '__main__':
    logging.getLogger().setLevel(logging.INFO)
    main()