From a2e03a5ddad5949581768e17fd834099f8d94ff2 Mon Sep 17 00:00:00 2001
From: Jerry Liu <jerryjliu98@gmail.com>
Date: Mon, 21 Nov 2022 22:05:08 -0800
Subject: [PATCH] add keyword lookup to GPTListIndex (#30)
---
gpt_index/indices/list/base.py | 3 ++-
gpt_index/indices/list/query.py | 6 ++++++
requirements.txt | 4 ++--
3 files changed, 10 insertions(+), 3 deletions(-)
diff --git a/gpt_index/indices/list/base.py b/gpt_index/indices/list/base.py
index a994ab56c2..1f763cf1f6 100644
--- a/gpt_index/indices/list/base.py
+++ b/gpt_index/indices/list/base.py
@@ -14,6 +14,7 @@ from gpt_index.indices.data_structs import IndexList
from gpt_index.indices.list.query import GPTListIndexQuery
from gpt_index.indices.utils import get_chunk_size_given_prompt, truncate_text
from gpt_index.langchain_helpers.text_splitter import TokenTextSplitter
+from gpt_index.prompts.base import Prompt
from gpt_index.prompts.default_prompts import DEFAULT_TEXT_QA_PROMPT
from gpt_index.schema import Document
@@ -25,7 +26,7 @@ class GPTListIndex(BaseGPTIndex[IndexList]):
self,
documents: Optional[List[Document]] = None,
index_struct: Optional[IndexList] = None,
- text_qa_template: str = DEFAULT_TEXT_QA_PROMPT,
+ text_qa_template: Prompt = DEFAULT_TEXT_QA_PROMPT,
) -> None:
"""Initialize params."""
self.text_qa_template = text_qa_template
diff --git a/gpt_index/indices/list/query.py b/gpt_index/indices/list/query.py
index 5bb7bbc5d9..b38a47ae33 100644
--- a/gpt_index/indices/list/query.py
+++ b/gpt_index/indices/list/query.py
@@ -1,4 +1,6 @@
"""Default query for GPTListIndex."""
+from typing import Optional
+
from gpt_index.indices.base import BaseGPTIndexQuery
from gpt_index.indices.data_structs import IndexList
from gpt_index.indices.response_utils import give_response, refine_response
@@ -18,17 +20,21 @@ class GPTListIndexQuery(BaseGPTIndexQuery[IndexList]):
index_struct: IndexList,
text_qa_template: Prompt = DEFAULT_TEXT_QA_PROMPT,
refine_template: Prompt = DEFAULT_REFINE_PROMPT,
+ keyword: Optional[str] = None,
) -> None:
"""Initialize params."""
super().__init__(index_struct=index_struct)
self.text_qa_template = text_qa_template
self.refine_template = refine_template
+ self.keyword = keyword
def query(self, query_str: str, verbose: bool = False) -> str:
"""Answer a query."""
print(f"> Starting query: {query_str}")
response = None
for node in self.index_struct.nodes:
+ if self.keyword is not None and self.keyword not in node.text:
+ continue
fmt_text_chunk = truncate_text(node.text, 50)
if verbose:
print(f"> Searching in chunk: {fmt_text_chunk}")
diff --git a/requirements.txt b/requirements.txt
index a211e87c81..2c9695e47d 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -3,6 +3,6 @@
black
isort
mypy
-# flake8
-# flake8-docstrings
+flake8
+flake8-docstrings
pylint
--
GitLab