This is an automated email from the ASF dual-hosted git repository. vterentev pushed a commit to branch oss-image-cpu in repository https://gitbox.apache.org/repos/asf/beam.git
commit 008ccbe1576c084db3e2dd7889bca028c98096d6 Author: Vitaly Terentyev <[email protected]> AuthorDate: Tue Dec 23 20:26:12 2025 +0400 Fix lint --- .../examples/inference/pytorch_imagenet_rightfit.py | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/sdks/python/apache_beam/examples/inference/pytorch_imagenet_rightfit.py b/sdks/python/apache_beam/examples/inference/pytorch_imagenet_rightfit.py index 0d16f786581..fa380ed8ac7 100644 --- a/sdks/python/apache_beam/examples/inference/pytorch_imagenet_rightfit.py +++ b/sdks/python/apache_beam/examples/inference/pytorch_imagenet_rightfit.py @@ -192,7 +192,7 @@ class PostProcessDoFn(beam.DoFn): "image_id": image_id, "model_name": self.model_name, "topk": json.dumps(topk), - "infer_ts_ms": now_millis(), + "infer_ms": now_millis(), } @@ -319,7 +319,7 @@ def override_or_add(args, flag, value): # ============ Model factory (timm) ============ -def create_timm_model(model_name: str, num_classes: int = 1000): +def create_timm_m(model_name: str, num_classes: int = 1000): import timm model = timm.create_model( model_name, pretrained=True, num_classes=num_classes) @@ -367,8 +367,7 @@ def run_load_pipeline(known_args, pipeline_args): _ = ( lines | 'ToBytes' >> beam.Map(lambda line: line.encode('utf-8')) - | - 'WriteToPubSub' >> beam.io.WriteToPubSub(topic=known_args.pubsub_topic)) + | 'ToPubSub' >> beam.io.WriteToPubSub(topic=known_args.pubsub_topic)) return pipeline.run() @@ -409,7 +408,7 @@ def run( for bs in tried: try: model_handler = PytorchModelHandlerTensor( - model_class=lambda: create_timm_model(known_args.pretrained_model_name), + model_class=lambda: create_timm_m(known_args.pretrained_model_name), model_params={}, state_dict_path=known_args.model_state_dict_path, device=device, @@ -435,9 +434,7 @@ def run( "Falling back to batch_size=8 due to previous errors: %s", last_err) bs_ok = 8 model_handler = PytorchModelHandlerTensor( - model_class=lambda: create_timm_model( - known_args.pretrained_model_name - ), + model_class=lambda: create_timm_m(known_args.pretrained_model_name), model_params={}, state_dict_path=known_args.model_state_dict_path, device=device, @@ -500,7 +497,7 @@ def run( | 'WriteToBigQuery' >> beam.io.WriteToBigQuery( known_args.output_table, schema= - 'image_id:STRING, model_name:STRING, topk:STRING, infer_ts_ms:INT64', + 'image_id:STRING, model_name:STRING, topk:STRING, infer_ms:INT64', write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND, create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED, method=beam.io.WriteToBigQuery.Method.STREAMING_INSERTS))
