feedback implemented

microsoft · Oct 25, 2024 · a646e2c · a646e2c
1 parent 8c07704
commit a646e2c
Showing 1 changed file with 10 additions and 2 deletions.
diff --git a/graphrag/index/flows/generate_text_embeddings.py b/graphrag/index/flows/generate_text_embeddings.py
@@ -37,7 +37,7 @@ async def generate_text_embeddings(
     cache: PipelineCache,
     storage: PipelineStorage,
     embedded_fields: set[str],
-    base_text_embed: dict | None = None,
+    base_text_embed: dict,
 ) -> pd.DataFrame:
     """All the steps to generate all embeddings."""
     documents_embeddings = final_documents.loc[:, ["id", "raw_content"]]
@@ -118,10 +118,18 @@ async def _run_and_snapshot_embeddings(
     callbacks: VerbCallbacks,
     cache: PipelineCache,
     storage: PipelineStorage,
-    base_text_embed: dict | None = None,
+    base_text_embed: dict,
 ) -> None:
     """All the steps to generate single embedding."""
+    log.info(filename)
     if base_text_embed:
+        new_vector_store = {
+            "title_column": column_to_embed,
+            "collection_name": filename,
+        }
+
+        base_text_embed["strategy"]["vector_store"].update(new_vector_store)
+
         data["embedding"] = await embed_text(
             data,
             callbacks,