feat: add new EntityExtractor as option

note: extractors are currently off by default.
docqai · Sep 27, 2023 · 42790ae · 42790ae
1 parent 6b2169b
commit 42790ae
Show file tree

Hide file tree

Showing 3 changed files with 254 additions and 1 deletion.
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -22,6 +22,7 @@ opendal = "^0.38.1"
 llama-index = "0.8.34"
 transformers = "^4.33.2"
 optimum = {extras = ["exporters"], version = "^1.13.2"}
+span-marker = "^1.3.0"
 
 [tool.poetry.group.dev.dependencies]
 pre-commit = "^2.18.1"

diff --git a/source/docq/support/llm.py b/source/docq/support/llm.py
@@ -22,6 +22,7 @@
 from llama_index.indices.composability import ComposableGraph
 from llama_index.node_parser import SimpleNodeParser
 from llama_index.node_parser.extractors import (
+    EntityExtractor,
     KeywordExtractor,
     MetadataExtractor,
     QuestionsAnsweredExtractor,
@@ -179,6 +180,7 @@ def _get_node_parser() -> SimpleNodeParser:
             QuestionsAnsweredExtractor(questions=3),
             SummaryExtractor(summaries=["prev", "self"]),
             KeywordExtractor(keywords=10),
+            EntityExtractor(entities=["prev", "self"]),
             # CustomExtractor()
         ],
     )