@@ -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