New workflow to generate embeddings in a single workflow

.semversioner/next-release/minor-20241018204541069382.json

@@ -0,0 +1,4 @@
+{
+  "type": "minor",
+  "description": "embeddings moved to a different workflow"
+}

graphrag/config/enums.py

@@ -86,6 +86,7 @@ class TextEmbeddingTarget(str, Enum):
 
     all = "all"
     required = "required"
+    none = "none"
 
     def __repr__(self):
         """Get a string representation."""

graphrag/index/config/__init__.py

@@ -11,6 +11,18 @@
     PipelineMemoryCacheConfig,
     PipelineNoneCacheConfig,
 )
+from .embeddings import (
+    all_embeddings,
+    community_full_content_embedding,
+    community_summary_embedding,
+    community_title_embedding,
+    document_raw_content_embedding,
+    entity_description_embedding,
+    entity_name_embedding,
+    relationship_description_embedding,
+    required_embeddings,
+    text_unit_text_embedding,
+)
 from .input import (
     PipelineCSVInputConfig,
     PipelineInputConfig,
@@ -66,4 +78,14 @@
     "PipelineWorkflowConfig",
     "PipelineWorkflowReference",
     "PipelineWorkflowStep",
+    "all_embeddings",
+    "community_full_content_embedding",
+    "community_summary_embedding",
+    "community_title_embedding",
+    "document_raw_content_embedding",
+    "entity_description_embedding",
+    "entity_name_embedding",
+    "relationship_description_embedding",
+    "required_embeddings",
+    "text_unit_text_embedding",
 ]

graphrag/index/config/embeddings.py

@@ -0,0 +1,25 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""A module containing embeddings values."""
+
+entity_name_embedding = "entity.name"
+entity_description_embedding = "entity.description"
+relationship_description_embedding = "relationship.description"
+document_raw_content_embedding = "document.raw_content"
+community_title_embedding = "community.title"
+community_summary_embedding = "community.summary"
+community_full_content_embedding = "community.full_content"
+text_unit_text_embedding = "text_unit.text"
+
+all_embeddings: set[str] = {
+    entity_name_embedding,
+    entity_description_embedding,
+    relationship_description_embedding,
+    document_raw_content_embedding,
+    community_title_embedding,
+    community_summary_embedding,
+    community_full_content_embedding,
+    text_unit_text_embedding,
+}
+required_embeddings: set[str] = {entity_description_embedding}

graphrag/index/create_pipeline_config.py

@@ -25,6 +25,10 @@
     PipelineMemoryCacheConfig,
     PipelineNoneCacheConfig,
 )
+from graphrag.index.config.embeddings import (
+    all_embeddings,
+    required_embeddings,
+)
 from graphrag.index.config.input import (
     PipelineCSVInputConfig,
     PipelineInputConfigTypes,
@@ -55,6 +59,7 @@
     create_final_community_reports,
     create_final_covariates,
     create_final_documents,
+    create_final_embeddings,
     create_final_entities,
     create_final_nodes,
     create_final_relationships,
@@ -63,29 +68,6 @@
 
 log = logging.getLogger(__name__)
 
-
-entity_name_embedding = "entity.name"
-entity_description_embedding = "entity.description"
-relationship_description_embedding = "relationship.description"
-document_raw_content_embedding = "document.raw_content"
-community_title_embedding = "community.title"
-community_summary_embedding = "community.summary"
-community_full_content_embedding = "community.full_content"
-text_unit_text_embedding = "text_unit.text"
-
-all_embeddings: set[str] = {
-    entity_name_embedding,
-    entity_description_embedding,
-    relationship_description_embedding,
-    document_raw_content_embedding,
-    community_title_embedding,
-    community_summary_embedding,
-    community_full_content_embedding,
-    text_unit_text_embedding,
-}
-required_embeddings: set[str] = {entity_description_embedding}
-
-
 builtin_document_attributes: set[str] = {
     "id",
     "source",
@@ -121,11 +103,12 @@ def create_pipeline_config(settings: GraphRagConfig, verbose=False) -> PipelineC
         storage=_get_storage_config(settings),
         cache=_get_cache_config(settings),
         workflows=[
-            *_document_workflows(settings, embedded_fields),
-            *_text_unit_workflows(settings, covariates_enabled, embedded_fields),
-            *_graph_workflows(settings, embedded_fields),
-            *_community_workflows(settings, covariates_enabled, embedded_fields),
+            *_document_workflows(settings),
+            *_text_unit_workflows(settings, covariates_enabled),
+            *_graph_workflows(settings),
+            *_community_workflows(settings, covariates_enabled),
             *(_covariate_workflows(settings) if covariates_enabled else []),
+            *(_embeddings_workflows(settings, embedded_fields)),
         ],
     )
 
@@ -138,9 +121,11 @@ def create_pipeline_config(settings: GraphRagConfig, verbose=False) -> PipelineC
 def _get_embedded_fields(settings: GraphRagConfig) -> set[str]:
     match settings.embeddings.target:
         case TextEmbeddingTarget.all:
-            return all_embeddings - {*settings.embeddings.skip}
+            return all_embeddings.difference(settings.embeddings.skip)
         case TextEmbeddingTarget.required:
             return required_embeddings
+        case TextEmbeddingTarget.none:
+            return set()
         case _:
             msg = f"Unknown embeddings target: {settings.embeddings.target}"
             raise ValueError(msg)
@@ -163,11 +148,8 @@ def _log_llm_settings(settings: GraphRagConfig) -> None:
 
 
 def _document_workflows(
-    settings: GraphRagConfig, embedded_fields: set[str]
+    settings: GraphRagConfig,
 ) -> list[PipelineWorkflowReference]:
-    skip_document_raw_content_embedding = (
-        document_raw_content_embedding not in embedded_fields
-    )
     return [
         PipelineWorkflowReference(
             name=create_final_documents,
@@ -176,15 +158,6 @@ def _document_workflows(
                     {*(settings.input.document_attribute_columns)}
                     - builtin_document_attributes
                 ),
-                "document_raw_content_embed": _get_embedding_settings(
-                    settings.embeddings,
-                    "document_raw_content",
-                    {
-                        "title_column": "raw_content",
-                        "collection_name": "final_documents_raw_content_embedding",
-                    },
-                ),
-                "skip_raw_content_embedding": skip_document_raw_content_embedding,
             },
         ),
     ]
@@ -193,9 +166,7 @@ def _document_workflows(
 def _text_unit_workflows(
     settings: GraphRagConfig,
     covariates_enabled: bool,
-    embedded_fields: set[str],
 ) -> list[PipelineWorkflowReference]:
-    skip_text_unit_embedding = text_unit_text_embedding not in embedded_fields
     return [
         PipelineWorkflowReference(
             name=create_base_text_units,
@@ -211,13 +182,7 @@ def _text_unit_workflows(
         PipelineWorkflowReference(
             name=create_final_text_units,
             config={
-                "text_unit_text_embed": _get_embedding_settings(
-                    settings.embeddings,
-                    "text_unit_text",
-                    {"title_column": "text", "collection_name": "text_units_embedding"},
-                ),
                 "covariates_enabled": covariates_enabled,
-                "skip_text_unit_embedding": skip_text_unit_embedding,
             },
         ),
     ]
@@ -247,16 +212,7 @@ def _get_embedding_settings(
     }
 
 
-def _graph_workflows(
-    settings: GraphRagConfig, embedded_fields: set[str]
-) -> list[PipelineWorkflowReference]:
-    skip_entity_name_embedding = entity_name_embedding not in embedded_fields
-    skip_entity_description_embedding = (
-        entity_description_embedding not in embedded_fields
-    )
-    skip_relationship_description_embedding = (
-        relationship_description_embedding not in embedded_fields
-    )
+def _graph_workflows(settings: GraphRagConfig) -> list[PipelineWorkflowReference]:
     return [
         PipelineWorkflowReference(
             name=create_base_entity_graph,
@@ -286,40 +242,11 @@ def _graph_workflows(
         ),
         PipelineWorkflowReference(
             name=create_final_entities,
-            config={
-                "entity_name_embed": _get_embedding_settings(
-                    settings.embeddings,
-                    "entity_name",
-                    {
-                        "title_column": "name",
-                        "collection_name": "entity_name_embeddings",
-                    },
-                ),
-                "entity_name_description_embed": _get_embedding_settings(
-                    settings.embeddings,
-                    "entity_name_description",
-                    {
-                        "title_column": "description",
-                        "collection_name": "entity_description_embeddings",
-                    },
-                ),
-                "skip_name_embedding": skip_entity_name_embedding,
-                "skip_description_embedding": skip_entity_description_embedding,
-            },
+            config={},
         ),
         PipelineWorkflowReference(
             name=create_final_relationships,
-            config={
-                "relationship_description_embed": _get_embedding_settings(
-                    settings.embeddings,
-                    "relationship_description",
-                    {
-                        "title_column": "description",
-                        "collection_name": "relationships_description_embeddings",
-                    },
-                ),
-                "skip_description_embedding": skip_relationship_description_embedding,
-            },
+            config={},
         ),
         PipelineWorkflowReference(
             name=create_final_nodes,
@@ -332,52 +259,21 @@ def _graph_workflows(
 
 
 def _community_workflows(
-    settings: GraphRagConfig, covariates_enabled: bool, embedded_fields: set[str]
+    settings: GraphRagConfig, covariates_enabled: bool
 ) -> list[PipelineWorkflowReference]:
-    skip_community_title_embedding = community_title_embedding not in embedded_fields
-    skip_community_summary_embedding = (
-        community_summary_embedding not in embedded_fields
-    )
-    skip_community_full_content_embedding = (
-        community_full_content_embedding not in embedded_fields
-    )
     return [
         PipelineWorkflowReference(name=create_final_communities),
         PipelineWorkflowReference(
             name=create_final_community_reports,
             config={
                 "covariates_enabled": covariates_enabled,
-                "skip_title_embedding": skip_community_title_embedding,
-                "skip_summary_embedding": skip_community_summary_embedding,
-                "skip_full_content_embedding": skip_community_full_content_embedding,
                 "create_community_reports": {
                     **settings.community_reports.parallelization.model_dump(),
                     "async_mode": settings.community_reports.async_mode,
                     "strategy": settings.community_reports.resolved_strategy(
                         settings.root_dir
                     ),
                 },
-                "community_report_full_content_embed": _get_embedding_settings(
-                    settings.embeddings,
-                    "community_report_full_content",
-                    {
-                        "title_column": "full_content",
-                        "collection_name": "final_community_reports_full_content_embedding",
-                    },
-                ),
-                "community_report_summary_embed": _get_embedding_settings(
-                    settings.embeddings,
-                    "community_report_summary",
-                    {
-                        "title_column": "summary",
-                        "collection_name": "final_community_reports_summary_embedding",
-                    },
-                ),
-                "community_report_title_embed": _get_embedding_settings(
-                    settings.embeddings,
-                    "community_report_title",
-                    {"title_column": "title"},
-                ),
             },
         ),
     ]
@@ -401,6 +297,27 @@ def _covariate_workflows(
     ]
 
 
+def _embeddings_workflows(
+    settings: GraphRagConfig, embedded_fields: set[str]
+) -> list[PipelineWorkflowReference]:
+    return [
+        PipelineWorkflowReference(
+            name=create_final_embeddings,
+            config={
+                "text_embed": _get_embedding_settings(
+                    settings.embeddings,
+                    "column_to_embed",
+                    {
+                        "title_column": "column_to_embed",
+                        "collection_name": "embedding",
+                    },
+                ),
+                "embedded_fields": embedded_fields,
+            },
+        ),
+    ]
+
+
 def _get_pipeline_input_config(
     settings: GraphRagConfig,
 ) -> PipelineInputConfigTypes: