File tree Expand file tree Collapse file tree
sdks/python/apache_beam/ml/rag/embeddings Expand file tree Collapse file tree Original file line number Diff line number Diff line change 1616
1717"""Tests for apache_beam.ml.rag.embeddings.vertex_ai."""
1818
19- import os
2019import shutil
2120import struct
2221import tempfile
@@ -186,21 +185,17 @@ def test_image_embedding_pipeline(self):
186185
187186 def test_image_embedding_pipeline_from_path (self ):
188187 png_bytes = _create_png_bytes ()
189- img_path = os .path .join (self .artifact_location , 'test_img.png' )
190- with open (img_path , 'wb' ) as f :
191- f .write (png_bytes )
192-
193188 test_items = [
194189 EmbeddableItem .from_image (
195- img_path , id = "img1" , metadata = {"source" : "test" }),
190+ png_bytes , id = "img1" , metadata = {"source" : "test" }),
196191 ]
197192
198193 expected = [
199194 EmbeddableItem (
200195 id = "img1" ,
201196 embedding = Embedding (dense_embedding = [0.0 ] * 1408 ),
202197 metadata = {"source" : "test" },
203- content = Content (image = img_path )),
198+ content = Content (image = png_bytes )),
204199 ]
205200
206201 artifact_location = tempfile .mkdtemp (prefix = 'vertex_ai_img_path_' )
You can’t perform that action at this time.
0 commit comments