Register and Configure Custom Embedding Model

[sivakumarl]

Hi All,
I am tyring to integrate and configure Huggingface TEI , able to implement the custom embedding model using EmbeddingModel , but not sure how to register this custom embedding class in conflig.yml for models of type embeddings . Please assist, appreciate your help. Thanks

from nemoguardrails.embeddings.embedding_providers import EmbeddingModel
from typing import List
import requests

class CustomEmbedLLM(EmbeddingModel):
    def encode( self, documents: List[str]) -> List[List[float]]:
        API_URL = "https://tsta.com/api/v1/models/all-MiniLM-L6-v2/embed"
        API_TOKEN = "abc1234"
        payload = {
            "inputs": documents
        }
        headers = {"Authorization": f"Basic {API_TOKEN}"}
        res = requests.post(API_URL, json=payload, headers=headers, verify=False)
        return res.text

Regards,
Siva Kumar.

[drazvan]

Hi @sivakumarl !

Currently, there isn't an API for registering an embedding model, similar to LLM providers/actions/embedding search providers.
If you're working with a cloned version of NeMo Guardrails, you'll need to edit here: https://github.com/NVIDIA/NeMo-Guardrails/blob/develop/nemoguardrails/embeddings/embedding_providers/__init__.py#L71.

If not, a quick workaround would be to add it manually into the _embedding_model_cache, as a sort of monkey patching. Let me know if you need more guidance on that.

If you have a bit of time to contribute the addition of a register_embedding_model, that would be great.

[sivakumarl]

Hi @drazvan ,

Thank you for the response , Could you please provide guidance on utilizing the _embedding_model_cache? I'm eager to contribute to adding the register_embedding_model feature. Thanks!

[drazvan]

Hi @sivakumarl !

In config.py or in the Python code that uses the LLMRails instance, you should do something along these lines:

from nemoguardrails.embeddings import embedding_providers

embedding_providers._embedding_model_cache["custom_embed"] = CustomEmbedLLM()

Then, in your config.yml you should be able to use it like this:

models:
  ...
  - type: embeddings
    engine: custom_embed

I haven't tested this, so let me know if you encounter any issues.

[sivakumarl]

Hi @drazvan ,
Able to make it work after some tweaks in nemoguardrails/embeddings/embedding_providers/init.py file with additional condition for custom_embed, appreciated your help.

    if model_key not in _embedding_model_cache:
        if embedding_engine == "SentenceTransformers":
            from .sentence_transformers import SentenceTransformerEmbeddingModel

            model = SentenceTransformerEmbeddingModel(embedding_model)

        elif embedding_engine == "FastEmbed":
            from .fastembed import FastEmbedEmbeddingModel

            model = FastEmbedEmbeddingModel(embedding_model)

        elif embedding_engine == "openai":
            from .openai import OpenAIEmbeddingModel

            model = OpenAIEmbeddingModel(embedding_model)
          elif embedding_engine == "custom_embed":
            
            model = _embedding_model_cache["custom_embed"]


        else:
            raise ValueError(f"Invalid embedding engine: {embedding_engine}")

Regards,
Siva Kumar.