Fix weight casting during graph extraction (#1016)

* Fix weight casting during graph extraction * Format * Format
microsoft · Aug 24, 2024 · 55e74a0 · 55e74a0
1 parent e15df44
commit 55e74a0
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 6 deletions.
diff --git a/.semversioner/next-release/patch-20240823233325895089.json b/.semversioner/next-release/patch-20240823233325895089.json
@@ -0,0 +1,4 @@
+{
+  "type": "patch",
+  "description": "Fix weight casting during graph extraction"
+}
diff --git a/graphrag/index/graph/extractors/graph/graph_extractor.py b/graphrag/index/graph/extractors/graph/graph_extractor.py
@@ -4,7 +4,6 @@
 """A module containing 'GraphExtractionResult' and 'GraphExtractor' models."""
 
 import logging
-import numbers
 import re
 import traceback
 from collections.abc import Mapping
@@ -248,11 +247,11 @@ async def _process_results(
                     target = clean_str(record_attributes[2].upper())
                     edge_description = clean_str(record_attributes[3])
                     edge_source_id = clean_str(str(source_doc_id))
-                    weight = (
-                        float(record_attributes[-1])
-                        if isinstance(record_attributes[-1], numbers.Number)
-                        else 1.0
-                    )
+                    try:
+                        weight = float(record_attributes[-1])
+                    except ValueError:
+                        weight = 1.0
+
                     if source not in graph.nodes():
                         graph.add_node(
                             source,