Sagemaker embeddings

[kirtishrinkhala]

I am following the example here

https://python.langchain.com/en/harrison-docs-refactor-3-24/modules/models/text_embedding/examples/sagemaker-endpoint.html

to integrate with an embedding model deployed using : https://www.philschmid.de/custom-inference-huggingface-sagemaker

However, it is failing with :

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
Cell In[3], line 37
     26 content_handler = ContentHandler()
     29 embeddings = SagemakerEndpointEmbeddings(
     30     # endpoint_name="endpoint-name", 
     31     # credentials_profile_name="credentials-profile-name", 
   (...)
     34     content_handler=content_handler
     35 )
---> 37 query_result = embeddings.embed_query("foo")

File ~/.pyenv/versions/3.8.11/lib/python3.8/site-packages/langchain/embeddings/sagemaker_endpoint.py:194, in SagemakerEndpointEmbeddings.embed_query(self, text)
    185 def embed_query(self, text: str) -> List[float]:
    186     """Compute query embeddings using a SageMaker inference endpoint.
    187 
    188     Args:
   (...)
    192         Embeddings for the text.
    193     """
--> 194     return self._embedding_func([text])

File ~/.pyenv/versions/3.8.11/lib/python3.8/site-packages/langchain/embeddings/sagemaker_endpoint.py:161, in SagemakerEndpointEmbeddings._embedding_func(self, texts)
    158 except Exception as e:
    159     raise ValueError(f"Error raised by inference endpoint: {e}")
--> 161 return self.content_handler.transform_output(response["Body"])

Cell In[3], line 24, in ContentHandler.transform_output(self, output)
     22 def transform_output(self, output: bytes) -> str:
     23     response_json = json.loads(output.read().decode("utf-8"))
---> 24     return response_json["embeddings"]

KeyError: 'embeddings'

What might be going wrong here?

[kirtishrinkhala]

It turns out that different models have different JSON structures for the embedding that was causing the issue. In my case the embeddings were in vector field