Package org.springframework.ai.transformers

Class TransformersEmbeddingModel

java.lang.Object
org.springframework.ai.embedding.AbstractEmbeddingModel
org.springframework.ai.transformers.TransformersEmbeddingModel
All Implemented Interfaces:
AutoCloseable, EmbeddingModel, Model<EmbeddingRequest,EmbeddingResponse>, org.springframework.beans.factory.InitializingBean
public class TransformersEmbeddingModel extends AbstractEmbeddingModel implements org.springframework.beans.factory.InitializingBean, AutoCloseable
An implementation of the AbstractEmbeddingModel that uses ONNX-based Transformer models for text embeddings.

By default, it uses the all-MiniLM-L6-v2 model, but can be configured to use other ONNX-compatible models. The class supports both CPU and GPU inference, caching of model resources, and various tokenization options.

For more information on the underlying SBERT framework, see: SBERT Documentation SBERT Pre-trained Models

Since:
1.0.0
Author:
Christian Tzolov, Soby Chacko

  • Field Details

  • Constructor Details

    • TransformersEmbeddingModel

      public TransformersEmbeddingModel()

    • TransformersEmbeddingModel

      public TransformersEmbeddingModel(MetadataMode metadataMode)

    • TransformersEmbeddingModel

      public TransformersEmbeddingModel(MetadataMode metadataMode, io.micrometer.observation.ObservationRegistry observationRegistry)

  • Method Details

    • setTokenizerOptions

      public void setTokenizerOptions(Map<String,String> tokenizerOptions)

    • setDisableCaching

      public void setDisableCaching(boolean disableCaching)

    • setResourceCacheDirectory

      public void setResourceCacheDirectory(@Nullable String resourceCacheDir)

    • setGpuDeviceId

      public void setGpuDeviceId(int gpuDeviceId)

    • setTokenizerResource

      public void setTokenizerResource(org.springframework.core.io.Resource tokenizerResource)

    • setModelResource

      public void setModelResource(org.springframework.core.io.Resource modelResource)

    • setTokenizerResource

      public void setTokenizerResource(String tokenizerResourceUri)

    • setModelResource

      public void setModelResource(String modelResourceUri)

    • setModelOutputName

      public void setModelOutputName(String modelOutputName)

    • afterPropertiesSet

      public void afterPropertiesSet() throws Exception
      Specified by:
      afterPropertiesSet in interface org.springframework.beans.factory.InitializingBean
      Throws:
      Exception

    • close

      public void close() throws ai.onnxruntime.OrtException
      Release the native ONNX runtime session and tokenizer acquired in afterPropertiesSet(). Spring registers this as the bean destroy method automatically (inferred close() method), and the model can also be used with try-with-resources.
      Specified by:
      close in interface AutoCloseable
      Throws:
      ai.onnxruntime.OrtException

    • embed

      public float[] embed(String text)
      Description copied from interface: EmbeddingModel
      Embeds the given text into a vector.
      Specified by:
      embed in interface EmbeddingModel
      Parameters:
      text - the text to embed.
      Returns:
      the embedded vector.

    • getEmbeddingContent

      public @Nullable String getEmbeddingContent(Document document)
      Description copied from interface: EmbeddingModel
      Extracts the text content from a Document to be used for embedding. By default, returns Document.getText(). Implementations that support MetadataMode should override this method to return Document.getFormattedContent(org.springframework.ai.document.MetadataMode) with the appropriate metadata mode, so that metadata is included in the text sent to the embedding API.
      Specified by:
      getEmbeddingContent in interface EmbeddingModel
      Parameters:
      document - the document to extract embedding content from.
      Returns:
      the text content to embed.

    • embed

      public float[] embed(Document document)
      Description copied from interface: EmbeddingModel
      Embeds the given document's content into a vector.
      Specified by:
      embed in interface EmbeddingModel
      Parameters:
      document - the document to embed.
      Returns:
      the embedded vector.

    • embedForResponse

      public EmbeddingResponse embedForResponse(List<String> texts)
      Description copied from interface: EmbeddingModel
      Embeds a batch of texts into vectors and returns the EmbeddingResponse.
      Specified by:
      embedForResponse in interface EmbeddingModel
      Parameters:
      texts - list of texts to embed.
      Returns:
      the embedding response.

    • embed

      public List<float[]> embed(List<String> texts)
      Description copied from interface: EmbeddingModel
      Embeds a batch of texts into vectors.
      Specified by:
      embed in interface EmbeddingModel
      Parameters:
      texts - list of texts to embed.
      Returns:
      list of embedded vectors.

    • call

      public EmbeddingResponse call(EmbeddingRequest request)
      Description copied from interface: Model
      Executes a method call to the AI model.
      Specified by:
      call in interface EmbeddingModel
      Specified by:
      call in interface Model<EmbeddingRequest,EmbeddingResponse>
      Parameters:
      request - the request object to be sent to the AI model
      Returns:
      the response from the AI model

    • setObservationConvention

      public void setObservationConvention(EmbeddingModelObservationConvention observationConvention)
      Use the provided convention for reporting observation data
      Parameters:
      observationConvention - The provided convention