diff --git a/litellm/llms/huggingface_restapi.py b/litellm/llms/huggingface_restapi.py index da5f7bebea..ec35532af5 100644 --- a/litellm/llms/huggingface_restapi.py +++ b/litellm/llms/huggingface_restapi.py @@ -375,9 +375,19 @@ def embedding( else: embed_url = f"https://api-inference.huggingface.co/models/{model}" - data = { - "inputs": input - } + if "sentence-transformers" in model: + if len(input) == 0: + raise HuggingfaceError(status_code=400, message="sentence transformers requires 2+ sentences") + data = { + "inputs": { + "source_sentence": input[0], + "sentences": [ "That is a happy dog", "That is a very happy person", "Today is a sunny day" ] + } + } + else: + data = { + "inputs": input + } ## LOGGING logging_obj.pre_call( @@ -402,15 +412,38 @@ def embedding( embeddings = response.json() + if "error" in embeddings: + raise HuggingfaceError(status_code=500, message=embeddings['error']) + output_data = [] - for idx, embedding in enumerate(embeddings): - output_data.append( + print(f"embeddings: {embeddings}") + if "similarities" in embeddings: + for idx, embedding in embeddings["similarities"]: + output_data.append( { "object": "embedding", "index": idx, - "embedding": embedding[0][0] # flatten list returned from hf + "embedding": embedding # flatten list returned from hf } ) + else: + for idx, embedding in enumerate(embeddings): + if isinstance(embedding, float): + output_data.append( + { + "object": "embedding", + "index": idx, + "embedding": embedding # flatten list returned from hf + } + ) + else: + output_data.append( + { + "object": "embedding", + "index": idx, + "embedding": embedding[0][0] # flatten list returned from hf + } + ) model_response["object"] = "list" model_response["data"] = output_data model_response["model"] = model diff --git a/litellm/tests/test_embedding.py b/litellm/tests/test_embedding.py index d7bf867ab4..10c2ce8b20 100644 --- a/litellm/tests/test_embedding.py +++ b/litellm/tests/test_embedding.py @@ -89,17 +89,17 @@ def test_bedrock_embedding(): # test_bedrock_embedding() # comment out hf tests - since hf endpoints are unstable -# def test_hf_embedding(): -# try: -# # huggingface/microsoft/codebert-base -# # huggingface/facebook/bart-large -# response = embedding( -# model="huggingface/BAAI/bge-large-zh", input=["good morning from litellm", "this is another item"] -# ) -# print(f"response:", response) -# except Exception as e: -# pytest.fail(f"Error occurred: {e}") -# test_hf_embedding() +def test_hf_embedding(): + try: + # huggingface/microsoft/codebert-base + # huggingface/facebook/bart-large + response = embedding( + model="huggingface/sentence-transformers/all-MiniLM-L6-v2", input=["good morning from litellm", "this is another item"] + ) + print(f"response:", response) + except Exception as e: + pytest.fail(f"Error occurred: {e}") +test_hf_embedding() # test async embeddings def test_aembedding():