File tree Expand file tree Collapse file tree 2 files changed +5
-5
lines changed
sdks/python/apache_beam/ml Expand file tree Collapse file tree 2 files changed +5
-5
lines changed Original file line number Diff line number Diff line change 37
37
try :
38
38
import vertexai
39
39
except ImportError :
40
- vertexai = None
40
+ vertexai = None # type: ignore[assignment]
41
41
42
42
43
43
class VertexAITextEmbeddings (EmbeddingsManager ):
Original file line number Diff line number Diff line change @@ -115,7 +115,7 @@ def __init__(
115
115
def get_request (
116
116
self ,
117
117
text_batch : Sequence [TextEmbeddingInput ],
118
- model : MultiModalEmbeddingModel ,
118
+ model : TextEmbeddingModel ,
119
119
throttle_delay_secs : int ):
120
120
while self .throttler .throttle_request (time .time () * _MSEC_TO_SEC ):
121
121
LOGGER .info (
@@ -126,7 +126,7 @@ def get_request(
126
126
127
127
try :
128
128
req_time = time .time ()
129
- prediction = model .get_embeddings (text_batch )
129
+ prediction = model .get_embeddings (list ( text_batch ) )
130
130
self .throttler .successful_request (req_time * _MSEC_TO_SEC )
131
131
return prediction
132
132
except TooManyRequests as e :
@@ -145,11 +145,11 @@ def run_inference(
145
145
embeddings = []
146
146
batch_size = _BATCH_SIZE
147
147
for i in range (0 , len (batch ), batch_size ):
148
- text_batch = batch [i :i + batch_size ]
148
+ text_batch_strs = batch [i :i + batch_size ]
149
149
text_batch = [
150
150
TextEmbeddingInput (
151
151
text = text , title = self .title , task_type = self .task_type )
152
- for text in text_batch
152
+ for text in text_batch_strs
153
153
]
154
154
embeddings_batch = self .get_request (
155
155
text_batch = text_batch , model = model , throttle_delay_secs = 5 )
You can’t perform that action at this time.
0 commit comments