Ensure entity types to be str in prompt tune (#1015)

microsoft · Aug 24, 2024 · e15df44 · e15df44
1 parent 13e17d2
commit e15df44
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 2 deletions.
diff --git a/.semversioner/next-release/patch-20240823225051898406.json b/.semversioner/next-release/patch-20240823225051898406.json
@@ -0,0 +1,4 @@
+{
+  "type": "patch",
+  "description": "Ensure entity types to be str when running prompt tuning"
+}
diff --git a/graphrag/prompt_tune/generator/entity_extraction_prompt.py b/graphrag/prompt_tune/generator/entity_extraction_prompt.py
@@ -54,7 +54,7 @@ def create_entity_extraction_prompt(
         else UNTYPED_GRAPH_EXTRACTION_PROMPT
     )
     if isinstance(entity_types, list):
-        entity_types = ", ".join(entity_types)
+        entity_types = ", ".join(map(str, entity_types))
 
     tokens_left = (
         max_token_count

diff --git a/graphrag/prompt_tune/generator/entity_relationship.py b/graphrag/prompt_tune/generator/entity_relationship.py
@@ -34,7 +34,9 @@ async def generate_entity_relationship_examples(
 
     if entity_types:
         entity_types_str = (
-            entity_types if isinstance(entity_types, str) else ", ".join(entity_types)
+            entity_types
+            if isinstance(entity_types, str)
+            else ", ".join(map(str, entity_types))
         )
 
         messages = [