From acb49c4ce11dd1224d4cad79268b8dd001c6fca5 Mon Sep 17 00:00:00 2001
From: Razvan Dinu <rdinu@nvidia.com>
Date: Fri, 1 Sep 2023 15:33:10 +0300
Subject: [PATCH] Add documentation for custom embeddings model and small
 tweaks.

---
 CHANGELOG.md                             |  1 +
 docs/user_guide/configuration-guide.md   | 22 ++++++++++++++++++++++
 nemoguardrails/actions/llm/generation.py |  7 +++++--
 nemoguardrails/rails/llm/llmrails.py     |  2 +-
 4 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8578b0d0f..28e0d9df1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Support for [custom configuration data](./docs/user_guide/configuration-guide.md#custom-data).
 - Example for using [custom LLM and multiple KBs](./examples/multi_kb/README.md)
 - Support for [`PROMPTS_DIR`](./docs/user_guide/advanced/prompt-customization.md#prompt-configuration).
+- [#101](https://github.com/NVIDIA/NeMo-Guardrails/pull/101) Support for [using OpenAI embeddings](./docs/user_guide/configuration-guide.md#the-embeddings-model) models in addition to SentenceTransformers.
 
 ### Fixed
 
diff --git a/docs/user_guide/configuration-guide.md b/docs/user_guide/configuration-guide.md
index 39ad365a1..120888192 100644
--- a/docs/user_guide/configuration-guide.md
+++ b/docs/user_guide/configuration-guide.md
@@ -111,6 +111,28 @@ models:
 ```
 
 
+### The Embeddings Model
+
+To configure the embeddings model that is used for the various steps in the [guardrails process](../architecture/README.md) (e.g., canonical form generation, next step generation) you can add a model configuration in the `models` key as shown below:
+
+```yaml
+models:
+  - ...
+  - type: embedding
+    engine: SentenceTransformers
+    model: all-MiniLM-L6-v2
+```
+
+The `SentenceTransformers` engine is the default one and uses the `all-MiniLM-L6-v2` model. NeMo Guardrails also supports using OpenAI models for computing the embeddings, e.g.:
+
+```yaml
+models:
+  - ...
+  - type: embedding
+    engine: openai
+    model: text-embedding-ada-002
+```
+
 ### General Instruction
 
 The general instruction (similar to a system prompt) gets appended at the beginning of every prompt, and you can configure it as shown below:
diff --git a/nemoguardrails/actions/llm/generation.py b/nemoguardrails/actions/llm/generation.py
index 57569604a..4b8c2d0bb 100644
--- a/nemoguardrails/actions/llm/generation.py
+++ b/nemoguardrails/actions/llm/generation.py
@@ -67,10 +67,13 @@ def __init__(
         self.llm = llm
         self.verbose = verbose
 
-        # If we have a customized embedding model, we'll use it.
+        # The default embeddings model is using SentenceTransformers
         self.embedding_model = "all-MiniLM-L6-v2"
+        self.embedding_engine = "SentenceTransformers"
+
+        # If we have a customized embedding model, we'll use it.
         for model in self.config.models:
-            if model.type == "embedding":
+            if model.type == "embeddings":
                 self.embedding_model = model.model
                 assert model.engine in get_embedding_provider_names()
                 self.embedding_engine = model.engine
diff --git a/nemoguardrails/rails/llm/llmrails.py b/nemoguardrails/rails/llm/llmrails.py
index 7139abec4..5e875d435 100644
--- a/nemoguardrails/rails/llm/llmrails.py
+++ b/nemoguardrails/rails/llm/llmrails.py
@@ -124,7 +124,7 @@ def _init_llms(self):
         #  to search for the main model config.
 
         for llm_config in self.config.models:
-            if llm_config.type == "embedding":
+            if llm_config.type == "embeddings":
                 if llm_config.engine not in get_embedding_provider_names():
                     raise Exception(f"Unknown embedding engine: {llm_config.engine}")
             else: