Skip to content

Sagemaker huggingface embedding model driver

SageMakerHuggingFaceEmbeddingModelDriver

Bases: BaseEmbeddingModelDriver

Source code in griptape/drivers/embedding_model/sagemaker_huggingface_embedding_model_driver.py 
 5
 6
 7
 8
 9
10
11
@define
class SageMakerHuggingFaceEmbeddingModelDriver(BaseEmbeddingModelDriver):
    def chunk_to_model_params(self, chunk: str) -> dict:
        return {"text_inputs": chunk}

    def process_output(self, output: dict) -> list[float]:
        return output["embedding"][0]

chunk_to_model_params(chunk)

Source code in griptape/drivers/embedding_model/sagemaker_huggingface_embedding_model_driver.py 
7
8
def chunk_to_model_params(self, chunk: str) -> dict:
    return {"text_inputs": chunk}

process_output(output)

Source code in griptape/drivers/embedding_model/sagemaker_huggingface_embedding_model_driver.py 
10
11
def process_output(self, output: dict) -> list[float]:
    return output["embedding"][0]