Resolve comments

Author: Amar3tto

sdks/python/apache_beam/examples/inference/pytorch_image_captioning.py

@@ -84,7 +84,7 @@ def process(self, kv: Tuple[str, str]):
       with FileSystems.open(uri) as f:
         image_bytes = f.read()
       yield uri, {"image_bytes": image_bytes}
-    except Exception as e:
+    except OSError as e:
       logging.warning("Failed to read image %s: %s", uri, e)
       return
 
@@ -164,6 +164,8 @@ def load_model(self):
     from transformers import BlipForConditionalGeneration, BlipProcessor
     processor = BlipProcessor.from_pretrained(self.model_name)
     model = BlipForConditionalGeneration.from_pretrained(self.model_name)
+    model.to(self.device)
+    model.eval()
     return (model, processor)
 
   def batch_elements_kwargs(self):
@@ -173,8 +175,6 @@ def run_inference(
       self, batch: List[Dict[str, Any]], model_bundle, inference_args=None):
 
     model, processor = model_bundle
-    model.to(self.device)
-    model.eval()
     start = now_millis()
 
     images = [x["image"] for x in batch]
@@ -236,6 +236,8 @@ def load_model(self):
     from transformers import CLIPModel, CLIPProcessor
     processor = CLIPProcessor.from_pretrained(self.model_name)
     model = CLIPModel.from_pretrained(self.model_name)
+    model.to(self.device)
+    model.eval()
     return (model, processor)
 
   def batch_elements_kwargs(self):
@@ -245,8 +247,6 @@ def run_inference(
       self, batch: List[Dict[str, Any]], model_bundle, inference_args=None):
 
     model, processor = model_bundle
-    model.to(self.device)
-    model.eval()
     start_batch = now_millis()
 
     # Flat lists for a single batched CLIP forward pass
@@ -464,15 +464,15 @@ def cleanup_pubsub_resources(
   try:
     subscriber.delete_subscription(
         request={"subscription": full_subscription_path})
-    print(f"Deleted subscription: {subscription_name}")
+    logging.info(f"Deleted subscription: {subscription_name}")
   except NotFound:
-    print(f"Subscription already deleted: {subscription_name}")
+    logging.info(f"Subscription already deleted: {subscription_name}")
 
   try:
     publisher.delete_topic(request={"topic": full_topic_path})
-    print(f"Deleted topic: {topic_name}")
+    logging.info(f"Deleted topic: {topic_name}")
   except NotFound:
-    print(f"Topic already deleted: {topic_name}")
+    logging.info(f"Topic already deleted: {topic_name}")
 
 
 def override_or_add(args, flag, value):

sdks/python/apache_beam/examples/inference/pytorch_image_object_detection.py

@@ -342,15 +342,15 @@ def cleanup_pubsub_resources(
   try:
     subscriber.delete_subscription(
         request={"subscription": full_subscription_path})
-    print(f"Deleted subscription: {subscription_name}")
+    logging.info(f"Deleted subscription: {subscription_name}")
   except NotFound:
-    print(f"Subscription already deleted: {subscription_name}")
+    logging.info(f"Subscription already deleted: {subscription_name}")
 
   try:
     publisher.delete_topic(request={"topic": full_topic_path})
-    print(f"Deleted topic: {topic_name}")
+    logging.info(f"Deleted topic: {topic_name}")
   except NotFound:
-    print(f"Topic already deleted: {topic_name}")
+    logging.info(f"Topic already deleted: {topic_name}")
 
 
 def override_or_add(args, flag, value):

sdks/python/apache_beam/examples/inference/pytorch_imagenet_rightfit.py

@@ -157,14 +157,10 @@ def process(self, kv: Tuple[str, PredictionResult]):
     if isinstance(inference_obj, dict):
       logits = inference_obj.get("logits", None)
       if logits is None:
-        # fallback: try first value if dict shape differs
-        try:
-          logits = next(iter(inference_obj.values()))
-          logging.warning(
-              'Could not find <logits> key in model output.'
-              'Falling back to first value in dict.')
-        except Exception:
-          logging.warning('Could not find <logits> key in dict.')
+        raise ValueError(
+            f"Unable to find 'logits' in model output. "
+            f"Available keys: {list(inference_obj.keys())}"
+        )
     else:
       logits = inference_obj
 
@@ -297,15 +293,15 @@ def cleanup_pubsub_resources(
   try:
     subscriber.delete_subscription(
         request={"subscription": full_subscription_path})
-    print(f"Deleted subscription: {subscription_name}")
+    logging.info(f"Deleted subscription: {subscription_name}")
   except NotFound:
-    print(f"Subscription already deleted: {subscription_name}")
+    logging.info(f"Subscription already deleted: {subscription_name}")
 
   try:
     publisher.delete_topic(request={"topic": full_topic_path})
-    print(f"Deleted topic: {topic_name}")
+    logging.info(f"Deleted topic: {topic_name}")
   except NotFound:
-    print(f"Topic already deleted: {topic_name}")
+    logging.info(f"Topic already deleted: {topic_name}")
 
 
 def override_or_add(args, flag, value):
@@ -475,6 +471,7 @@ def run(
 
   predictions = (
       to_infer
+      | 'Reshuffle' >> beam.Reshuffle()
       | 'RunInference' >> RunInference(
           KeyedModelHandler(model_handler)).with_resource_hints(
               accelerator="type:nvidia-tesla-t4;count:1;install-nvidia-driver"))