Skip to content

Commit bd94408

Browse files
edknvcharlesbluca
andauthored
fix embed stage ignoring pipeline endpoint (#1459)
Co-authored-by: Charles Blackmon-Luca <20627856+charlesbluca@users.noreply.github.com>
1 parent ac5518f commit bd94408

2 files changed

Lines changed: 19 additions & 27 deletions

File tree

api/api_tests/internal/transform/test_embed_text.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -201,7 +201,6 @@ def _local_embedder(texts):
201201
return [[0.9, 0.8]] * len(list(texts))
202202

203203
task_cfg = {
204-
"endpoint_url": None, # Explicitly disable remote endpoint
205204
"embedder": _local_embedder,
206205
"local_batch_size": 2,
207206
}

api/src/nv_ingest_api/internal/transform/embed_text.py

Lines changed: 19 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -717,18 +717,11 @@ def transform_create_text_embeddings_internal(
717717
- The updated DataFrame with embeddings applied.
718718
- A dictionary with trace information.
719719
"""
720-
# Allow task_config to explicitly override values with None by checking key presence.
721-
api_key = task_config["api_key"] if "api_key" in task_config else transform_config.api_key
722-
endpoint_url = (
723-
task_config["endpoint_url"] if "endpoint_url" in task_config else transform_config.embedding_nim_endpoint
724-
)
725-
model_name = task_config["model_name"] if "model_name" in task_config else transform_config.embedding_model
726-
custom_content_field = (
727-
task_config["custom_content_field"]
728-
if "custom_content_field" in task_config
729-
else transform_config.custom_content_field
730-
)
731-
dimensions = task_config["dimensions"] if "dimensions" in task_config else transform_config.dimensions
720+
api_key = task_config.get("api_key") or transform_config.api_key
721+
endpoint_url = task_config.get("endpoint_url") or transform_config.embedding_nim_endpoint
722+
model_name = task_config.get("model_name") or transform_config.embedding_model
723+
custom_content_field = task_config.get("custom_content_field") or transform_config.custom_content_field
724+
dimensions = task_config.get("dimensions") or transform_config.dimensions
732725

733726
endpoint_url = endpoint_url.strip() if isinstance(endpoint_url, str) else endpoint_url
734727
if isinstance(endpoint_url, str) and not endpoint_url:
@@ -868,7 +861,13 @@ def _content_type_getter(row):
868861
else:
869862
modality_batches = None
870863

871-
if endpoint_url:
864+
if callable(embedder):
865+
content_embeddings = _callable_runner(
866+
filtered_content_batches,
867+
embedder=embedder,
868+
batch_size=local_batch_size,
869+
)
870+
elif endpoint_url:
872871
content_embeddings = _async_runner(
873872
filtered_content_batches,
874873
api_key,
@@ -881,12 +880,6 @@ def _content_type_getter(row):
881880
modalities=modality_batches,
882881
dimensions=dimensions,
883882
)
884-
elif callable(embedder):
885-
content_embeddings = _callable_runner(
886-
filtered_content_batches,
887-
embedder=embedder,
888-
batch_size=local_batch_size,
889-
)
890883
else:
891884
raise ValueError(
892885
"No embedding endpoint configured (endpoint_url/embedding_nim_endpoint are empty) "
@@ -926,7 +919,13 @@ def _content_type_getter(row):
926919
custom_content_list = extracted_custom_content[valid_custom_content_mask].to_list()
927920
custom_content_batches = _generate_batches(custom_content_list, batch_size=transform_config.batch_size)
928921

929-
if endpoint_url:
922+
if callable(embedder):
923+
custom_content_embeddings = _callable_runner(
924+
custom_content_batches,
925+
embedder=embedder,
926+
batch_size=local_batch_size,
927+
)
928+
elif endpoint_url:
930929
custom_content_embeddings = _async_runner(
931930
custom_content_batches,
932931
api_key,
@@ -938,12 +937,6 @@ def _content_type_getter(row):
938937
False,
939938
dimensions=dimensions,
940939
)
941-
elif callable(embedder):
942-
custom_content_embeddings = _callable_runner(
943-
custom_content_batches,
944-
embedder=embedder,
945-
batch_size=local_batch_size,
946-
)
947940
else:
948941
raise ValueError(
949942
"No embedding endpoint configured (endpoint_url/embedding_nim_endpoint are empty) "

0 commit comments

Comments
 (0)