Merge pull request #2 from telekom/ms/fix-embedding-dim

fix: use EMBEDDING_DIM for milvus and set it to 768 as default
telekom · Apr 23, 2024 · 692fbe1 · 692fbe1
2 parents 2852cbc + b154e84
commit 692fbe1
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 2 deletions.
diff --git a/helpers/env-variables.ts b/helpers/env-variables.ts
@@ -150,6 +150,8 @@ export const createBackendEnvFile = async (
   ];
   let envVars: EnvVar[] = [];
   if (opts.framework === "fastapi") {
+    // TODO: select right dimension depending on the model
+    const defaultDimension = opts.model === "large-model" ? "1536" : "768";
     envVars = [
       ...defaultEnvs,
       ...[
@@ -171,6 +173,7 @@ export const createBackendEnvFile = async (
         {
           name: "EMBEDDING_DIM",
           description: "Dimension of the embedding model to use.",
+          value: defaultDimension,
         },
         {
           name: "LLM_TEMPERATURE",

diff --git a/templates/components/vectordbs/python/milvus/generate.py b/templates/components/vectordbs/python/milvus/generate.py
@@ -27,7 +27,7 @@ def generate_datasource():
         user=os.getenv("MILVUS_USERNAME"),
         password=os.getenv("MILVUS_PASSWORD"),
         collection_name=os.getenv("MILVUS_COLLECTION"),
-        dim=int(os.getenv("MILVUS_DIMENSION", "1536")),
+        dim=int(os.getenv("EMBEDDING_DIM")),
     )
     storage_context = StorageContext.from_defaults(vector_store=store)
     VectorStoreIndex.from_documents(

diff --git a/templates/components/vectordbs/python/milvus/index.py b/templates/components/vectordbs/python/milvus/index.py
@@ -19,7 +19,7 @@ def get_index():
         user=os.getenv("MILVUS_USERNAME"),
         password=os.getenv("MILVUS_PASSWORD"),
         collection_name=os.getenv("MILVUS_COLLECTION"),
-        dim=int(os.getenv("EMBEDDING_DIM", "1536")),
+        dim=int(os.getenv("EMBEDDING_DIM")),
     )
     index = VectorStoreIndex.from_vector_store(store)
     logger.info("Finished connecting to index from Milvus.")