jrmccluskey commented on code in PR #31514:
URL: https://github.com/apache/beam/pull/31514#discussion_r1629759836
##########
sdks/python/apache_beam/ml/transforms/embeddings/vertex_ai.py:
##########
@@ -157,3 +160,95 @@ def get_ptransform_for_processing(self, **kwargs) ->
beam.PTransform:
return RunInference(
model_handler=_TextEmbeddingHandler(self),
inference_args=self.inference_args)
+
+
+class _VertexAIImageEmbeddingHandler(ModelHandler):
+ def __init__(
+ self,
+ model_name: str,
+ dimension: Optional[int] = None,
+ project: Optional[str] = None,
+ location: Optional[str] = None,
+ credentials: Optional[Credentials] = None,
+ ):
+ vertexai.init(project=project, location=location, credentials=credentials)
+ self.model_name = model_name
+ self.dimension = dimension
+
+ def run_inference(
+ self,
+ batch: Sequence[Image],
+ model: MultiModalEmbeddingModel,
+ inference_args: Optional[Dict[str, Any]] = None,
+ ) -> Iterable:
+ embeddings = []
+ # Maximum request size for muli-model embedding models is 1.
+ for img in batch:
+ try:
+ embedding_response = model.get_embeddings(
+ image=img, dimension=self.dimension)
+ embeddings.append(embedding_response.image_embedding)
+ except Exception as e:
+ print(e)
Review Comment:
ah this was a holdover from some testing. I was getting a weird python error
that was apparently a python-on-M1 architecture thing, not a code issue.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscr...@beam.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
