[Question]: GPT and Simple KeywordTableRetriever won't respond to keyword template

[kylemassimilian]

Question Validation

• I have searched both the documentation and discord for an answer.

Question

I am trying to use a KeywordTableRetriever to combine vector and keyword search. My keyword retriever will not extract important keywords, specifically dates, even when I use a keyword prompt template to do so. Is there a way to give more direct guidance to the KeywordTableGPTRetriever? Has anyone also come across this problem?

`INFO:llama_index.indices.keyword_table.retrievers:> Starting query: Tell me about the February 2023 Alpha Monthly Update.
> Starting query: Tell me about the February 2023 Alpha Monthly Update.
INFO:llama_index.indices.keyword_table.retrievers:query keywords: ['2023', 'alpha', 'february', 'update', 'monthly']
query keywords: ['2023', 'alpha', 'february', 'update', 'monthly']
INFO:llama_index.indices.keyword_table.retrievers:> Extracted keywords: ['2023', 'alpha', 'update']
> Extracted keywords: ['2023', 'alpha', 'update']
INFO:llama_index.indices.keyword_table.retrievers:> Starting query: Tell me about the February 2023 Alpha Monthly Update.
> Starting query: Tell me about the February 2023 Alpha Monthly Update.
INFO:llama_index.indices.keyword_table.retrievers:query keywords: ['2023', 'alpha', 'february', 'update', 'monthly']
query keywords: ['2023', 'alpha', 'february', 'update', 'monthly']
INFO:llama_index.indices.keyword_table.retrievers:> Extracted keywords: ['2023', 'alpha', 'update']
> Extracted keywords: ['2023', 'alpha', 'update']`

`class CustomRetriever(BaseRetriever):
    """Custom retriever that performs both semantic search and hybrid search"""

    def __init__(
            self,
            vector_retriever: VectorIndexRetriever,
            keyword_retriever: KeywordTableGPTRetriever,
            mode: str = "OR"

    ) -> None:
        """Init params."""

        self._vector_retriever = vector_retriever
        self._keyword_retriever = keyword_retriever
        if mode not in ("AND", "OR"):
            raise ValueError("Invalid mode.")
        self._mode = mode
    
    def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
        """Retrieve nodes given query"""

        vector_nodes = self._vector_retriever.retrieve(query_bundle)
        keyword_nodes = self._keyword_retriever.retrieve(query_bundle)

        # extract just the node_id from each NodeWithScore object in vector_nodes and keyword_nodes and store in set
        vector_ids = {n.node.node_id for n in vector_nodes}
        keyword_ids = {n.node.node_id for n in keyword_nodes}

        # combined_dict stores all the NodeWithScore objects as values, with the node IDs as keys.
        combined_dict = {n.node.node_id: n for n in vector_nodes}
        combined_dict.update({n.node.node_id: n for n in keyword_nodes})

        # if mode is AND, take the intersection of the vector and keyword node IDs
        if self._mode == "AND":
            retrieve_ids = vector_ids.intersection(keyword_ids)
        else:
            retrieve_ids = vector_ids.union(keyword_ids)

        # extract the NodeWithScore objects for the retrieved IDs into a list and return that list
        retrieve_nodes = [combined_dict[rid] for rid in retrieve_ids]
        return retrieve_nodes


def kw_vector_search(vector_index, keyword_index, query):

    logging.basicConfig(stream=sys.stdout, level=logging.INFO)
    logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))


    QUERY_KEYWORD_EXTRACT_TEMPLATE_TMPL = (
        "A question is provided below. Given the question, extract as many keywords from the question as possible up to {max_keywords} "
        "keywords from the text.  "
        "to best lookup answers to the question\n"
        "---------------------\n"
        "{question}\n"
        "---------------------\n"
        "Extract keywords from the question that denote months (ex: February, March, April)"
        "Provide keywords in the following comma-separated format: 'KEYWORDS: <keywords>'\n"
    )
    QUERY_KEYWORD_EXTRACT_TEMPLATE = Prompt(
        QUERY_KEYWORD_EXTRACT_TEMPLATE_TMPL,
        prompt_type=PromptType.QUERY_KEYWORD_EXTRACT,
    )

    #define custom retriever
    vector_retriever = VectorIndexRetriever(index=vector_index, similarity_top_k=3)
    keyword_retriever = KeywordTableGPTRetriever(index=keyword_index, max_keywords=10,
                                                 query_keyword_extract_template=QUERY_KEYWORD_EXTRACT_TEMPLATE) #######
    custom_retriever = CustomRetriever(vector_retriever, keyword_retriever)

    # Define response synthesizer - callback manager?
    response_synthesizer = get_response_synthesizer()

    # Assemble query engine
    custom_query_engine = RetrieverQueryEngine(
        retriever=custom_retriever, 
        response_synthesizer=response_synthesizer,
    )
    #vector query engine
    vector_query_engine = RetrieverQueryEngine(
        retriever=vector_retriever, 
        response_synthesizer=response_synthesizer,
    )
    # keyword query engine
    keyword_query_engine = RetrieverQueryEngine(
        retriever=keyword_retriever, 
        response_synthesizer=response_synthesizer,
    )

    custom_response = custom_query_engine.query(query)
    vector_response = vector_query_engine.query(query)
    keyword_response = keyword_query_engine.query(query)
    
    return custom_response, vector_response, keyword_response

When I ask a question that does not have a month (or date) in the query, the keywords are months. Seems like the KeywordTableGPTRetriever is interpreting my prompt in a strange way.

`INFO:llama_index.indices.keyword_table.retrievers:> Starting query: What is CRC?
> Starting query: What is CRC?
> Starting query: What is CRC?
INFO:llama_index.indices.keyword_table.retrievers:query keywords: ['february', 'march', 'crc', 'april']
query keywords: ['february', 'march', 'crc', 'april']
query keywords: ['february', 'march', 'crc', 'april']
INFO:llama_index.indices.keyword_table.retrievers:> Extracted keywords: ['march', 'crc']
> Extracted keywords: ['march', 'crc']
> Extracted keywords: ['march', 'crc']`

When I just say the phrase "February 2023," there are no keywords fetched.

`INFO:llama_index.indices.keyword_table.retrievers:> Starting query: February 2023
> Starting query: February 2023
> Starting query: February 2023
INFO:llama_index.indices.keyword_table.retrievers:query keywords: ['february']
query keywords: ['february']
query keywords: ['february']
INFO:llama_index.indices.keyword_table.retrievers:> Extracted keywords: []
> Extracted keywords: []
> Extracted keywords: []`

Interestingly, when I ask "What date did we discuss x?" it's able to answer pretty well.

[dosubot[bot]]

Answer generated by a 🤖

Answer

Hello,

Thank you for your detailed question. It seems like you're having trouble with the KeywordTableGPTRetriever not extracting the keywords you expect, especially when it comes to dates.

The KeywordTableGPTRetriever uses a GPT model to predict the response to the query and then extracts keywords from this response. The query_keyword_extract_template is used as the prompt for the GPT model, and the max_keywords_per_query parameter is used to limit the number of keywords extracted. The question parameter is the query string from which keywords are to be extracted.

The response from the GPT model is then passed to the extract_keywords_given_response function, which extracts the keywords from the response. The start token for keyword extraction is "KEYWORDS:". The function returns a list of keywords.

If you're finding that the KeywordTableGPTRetriever is not extracting dates as expected, you might want to modify the _extract_keywords method in the KeywordTableIndex class. You can create a custom function that is specifically designed to extract dates from the response. This function can use regular expressions or a date parsing library like dateutil to identify and extract dates.

Here is an example of how you can modify the _extract_keywords method:

from dateutil.parser import parse
import re

def _extract_keywords(self, text: str) -> Set[str]:
    """Extract keywords from text."""
    response = self._service_context.llm_predictor.predict(
        self.keyword_extract_template,
        text=text,
    )
    keywords = extract_keywords_given_response(response, start_token="KEYWORDS:")

    # Custom function to extract dates
    def extract_dates(text):
        matches = re.findall(r'\b\d{1,2}[/-]\d{1,2}[/-]\d{4}\b', text)
        dates = [parse(match) for match in matches]
        return dates

    dates = extract_dates(response)
    keywords.update(dates)

    return keywords

In this example, the extract_dates function uses a regular expression to find dates in the format dd/mm/yyyy or dd-mm-yyyy in the response. It then uses the parse function from the dateutil.parser module to convert these matches into date objects. These date objects are then added to the set of keywords.

Please note that this is just an example and the actual implementation might need to be adjusted based on the specific requirements of your application. For instance, you might need to adjust the regular expression to match the date format used in your data, or use a different method to extract dates if they are not in a standard format.

I hope this helps! Let me know if you have any other questions.

Sources

• llama_index/indices/keyword_table/retrievers.py
• llama_index/indices/keyword_table/base.py

---

This response is meant to be useful, save you time, and share context. It is not meant to be a precise solution, but rather a starting point for your own research.

Help me be more useful! Please leave a 👍 if this is helpful and 👎 if it is irrelevant.

[kylemassimilian]

Could someone please tell me where these keywords are dropping off and how to fix this??

INFO:llama_index.indices.keyword_table.retrievers:> Starting query: Please summarize the February 2023 Alpha Monthly Update
> Starting query: Please summarize the February 2023 Alpha Monthly Update
Response from _get_keywords method: 

February, 2023, Alpha, Monthly, Update, Summarize
Results before expanding tokens with subtokens: ['february', '2023', 'alpha', 'monthly', 'update', 'summarize']
Results after expanding tokens with subtokens: ['february', '2023', 'alpha', 'monthly', 'update', 'summarize']
Keywords extracted from Custom Retriever
INFO:llama_index.indices.keyword_table.retrievers:query keywords: ['summarize', '2023', 'update', 'monthly', 'alpha', 'february']
query keywords: ['summarize', '2023', 'update', 'monthly', 'alpha', 'february']
INFO:llama_index.indices.keyword_table.retrievers:> Extracted keywords: ['2023', 'alpha']
> Extracted keywords: ['2023', 'alpha']

def build_index():

    ## Build Vector store index from Chroma
    chroma_db = chromadb.HttpClient(host="3.208.86.57", port=8000)
    
    collections = [col.name for col in chroma_db.list_collections()]
    
    collection_name = "cr-alpha-monthly-1"

    if collection_name in collections:
        print("Building index..")
        chroma_collection = chroma_db.get_collection(collection_name)
        #Set chunk size and LLM for querying phase
        service_context = ServiceContext.from_defaults(chunk_size=1024, llm=llm) #chunk_overlap=, chunk_size_limit=
        vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
        vector_index = VectorStoreIndex.from_vector_store(vector_store=vector_store, service_context=service_context,  
                                                   show_progress=True) #show_progress to display a progress bar during index construction
    else:
        print("Check name of index")
        sys.exit()


    ## Build List and Keyword indices from documents
    try:
        print("Loading data...") 
        documents = loader.load_data()
        print("Data loaded.")
    except Exception as e:
        print("Error loading data: {e}")
        sys.exit()
    

    # Initialize service context (set chunk size, LLM)
    service_context = ServiceContext.from_defaults(chunk_size=1024, llm=llm)
    
    #Compose nodes from documents (splits docs into smaller chunks of text)
    nodes = service_context.node_parser.get_nodes_from_documents(documents)

    # Initialize storage context - a utility container for storing nodes, indices, and vectors (by default it's in-memory - RAM)
    storage_context = StorageContext.from_defaults() #Check this in docs######
    
    # Add the nodes to the storage context's document store
    storage_context.docstore.add_documents(nodes)

    # Create the List and Keyword indices using the nodes and storage context
        #ListIndex is a simple data structure where nodes are stored in a sequence
    list_index = ListIndex(nodes, storage_context=storage_context)
        #KeywordIndex extracts keywords from each node and builds a mapping from each keyword to the corresponding nodes of that keyword
    keyword_index = CustomKeywordTableIndex(nodes, storage_context=storage_context, show_progress=True)
    print("Keyword index built")
    return vector_index, keyword_index



def extract_keywords_given_response(
    response: str, lowercase: bool = True, start_token: str = ""
) -> Set[str]:
    """Extract keywords given the GPT-generated response.

    Used by keyword table indices.
    Parses <start_token>: <word1>, <word2>, ... into [word1, word2, ...]
    Raises exception if response doesn't start with <start_token>
    """
    results = []
    response = response.strip()  # Strip newlines from responses.

    if response.startswith(start_token):
        response = response[len(start_token) :]

    keywords = response.split(",")
    for k in keywords:
        rk = k
        if lowercase:
            rk = rk.lower()
        results.append(rk.strip())
    
    # if keyword consists of multiple words, split into subwords
    # (removing stopwords)
    
    print("Results before expanding tokens with subtokens:", results)
    expanded_results = expand_tokens_with_subtokens(set(results))
    print("Results after expanding tokens with subtokens:", results)
    
    return expanded_results




class CustomKeywordTableGPTRetriever(BaseKeywordTableRetriever):
    """Custom Keyword Table GPT Retriever.

    Extracts keywords using GPT directly from the query string.

    See BaseGPTKeywordTableQuery for arguments.

    """

    def _get_keywords(self, query_str: str) -> List[str]:
        """Extract keywords."""
        # Modify the prompt to ask the GPT model to extract keywords from the query string
        prompt = f"Extract all keywords from the following query: '{query_str}'. Keywords:"

        response = self._service_context.llm_predictor.predict(
            prompt,
            max_keywords=self.max_keywords_per_query,
            question=query_str,
        )
        print( "Response from _get_keywords method:",response)
        keywords = extract_keywords_given_response(response, start_token="Keywords:")
        print("Keywords extracted from Custom Retriever")
        return list(keywords)


class CustomKeywordTableIndex(BaseKeywordTableIndex):
    """Custom Keyword Table Index.

    This index uses the CustomKeywordTableGPTRetriever to extract keywords from the text.

    """

    def _extract_keywords(self, text: str) -> Set[str]:
        """Extract keywords from text."""
        retriever = CustomKeywordTableGPTRetriever(self)
        keywords = retriever._get_keywords(text)
        print("Keywords extracted from Custom Index")
        return set(keywords)
        


class CustomRetriever(BaseRetriever):
    """Custom retriever that performs both semantic search and hybrid search"""

    def __init__(
            self,
            vector_retriever: VectorIndexRetriever,
            keyword_retriever: CustomKeywordTableGPTRetriever,
            mode: str = "OR"

    ) -> None:
        """Init params."""

        self._vector_retriever = vector_retriever
        self._keyword_retriever = keyword_retriever
        if mode not in ("AND", "OR"):
            raise ValueError("Invalid mode.")
        self._mode = mode
    
    def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
        """Retrieve nodes given query"""

        vector_nodes = self._vector_retriever.retrieve(query_bundle)
        keyword_nodes = self._keyword_retriever.retrieve(query_bundle)

        # extract just the node_id from each NodeWithScore object in vector_nodes and keyword_nodes and store in set
        vector_ids = {n.node.node_id for n in vector_nodes}
        keyword_ids = {n.node.node_id for n in keyword_nodes}

        # combined_dict stores all the NodeWithScore objects as values, with the node IDs as keys.
        combined_dict = {n.node.node_id: n for n in vector_nodes}
        combined_dict.update({n.node.node_id: n for n in keyword_nodes})

        # if mode is AND, take the intersection of the vector and keyword node IDs
        if self._mode == "AND":
            retrieve_ids = vector_ids.intersection(keyword_ids)
        else:
            retrieve_ids = vector_ids.union(keyword_ids)

        # extract the NodeWithScore objects for the retrieved IDs into a list and return that list
        retrieve_nodes = [combined_dict[rid] for rid in retrieve_ids]
        return retrieve_nodes


def kw_vector_search(vector_index, keyword_index, query):

    logging.basicConfig(stream=sys.stdout, level=logging.INFO)
    logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))


    QUERY_KEYWORD_EXTRACT_TEMPLATE_TMPL = (
        "A question is provided below. Given the question, extract as many keywords from the question as possible up to {max_keywords} "
        "keywords from the text.  "
        "to best lookup answers to the question\n"
        "If the question mentions a specific month, use this in your response. "
        "---------------------\n"
        "{question}\n"
        "---------------------\n"
        "Provide keywords in the following comma-separated format: 'KEYWORDS: <keywords>'\n"
    )
    QUERY_KEYWORD_EXTRACT_TEMPLATE = Prompt(
        QUERY_KEYWORD_EXTRACT_TEMPLATE_TMPL,
        prompt_type=PromptType.QUERY_KEYWORD_EXTRACT,
    )

    #define custom retriever
    vector_retriever = VectorIndexRetriever(index=vector_index, similarity_top_k=3)
    keyword_retriever = CustomKeywordTableGPTRetriever(index=keyword_index) #######
    custom_retriever = CustomRetriever(vector_retriever, keyword_retriever)

    # Define response synthesizer - callback manager?
    response_synthesizer = get_response_synthesizer()

    # Assemble query engine
    custom_query_engine = RetrieverQueryEngine(
        retriever=custom_retriever, 
        response_synthesizer=response_synthesizer,
    )
    #vector query engine
    vector_query_engine = RetrieverQueryEngine(
        retriever=vector_retriever, 
        response_synthesizer=response_synthesizer,
    )
    # keyword query engine
    keyword_query_engine = RetrieverQueryEngine(
        retriever=keyword_retriever, 
        response_synthesizer=response_synthesizer,
    )

    custom_response = custom_query_engine.query(query)
    vector_response = vector_query_engine.query(query)
    keyword_response = keyword_query_engine.query(query)
    
    return custom_response, vector_response, keyword_response

[dosubot[bot]]

Answer generated by a 🤖

Answer

Hello,

Based on the logs and code snippet you've provided, it seems like all the keywords are being correctly identified and extracted in the _get_keywords method of your CustomKeywordTableGPTRetriever. However, not all of these keywords are being returned by the _extract_keywords method in your CustomKeywordTableIndex.

This discrepancy could be due to the filtering process that happens in the _retrieve method of the BaseKeywordTableRetriever class. In this method, the extracted keywords are filtered to include only those present in the index structure (keywords = [k for k in keywords if k in self._index_struct.keywords]). If a keyword is not in the index structure, it will not be considered.

To resolve this issue, you can ensure that all the necessary keywords are present in your index structure. If a keyword is important but not present in the index structure, you might need to update your index.

Additionally, you might want to check the max_keywords_per_query parameter in your CustomKeywordTableGPTRetriever. This parameter limits the number of keywords extracted from the query. If it's set to a low value, some keywords might be missed. You can increase this value to allow more keywords to be extracted from the query.

Here's how you can adjust the max_keywords_per_query parameter:

class CustomKeywordTableGPTRetriever(BaseKeywordTableRetriever):
    """Custom Keyword Table GPT Retriever."""

    def __init__(self, index, max_keywords_per_query=20):  # Increase this value
        super().__init__(index)
        self.max_keywords_per_query = max_keywords_per_query

Remember to re-index your data if you make changes to the keyword extraction method or the max_keywords_per_query parameter.

I hope this helps! Let me know if you have any other questions.

Sources

• llama_index/indices/keyword_table/retrievers.py
• llama_index/indices/keyword_table/base.py
• llama_index/indices/keyword_table/simple_base.py

---

This response is meant to be useful, save you time, and share context. It is not meant to be a precise solution, but rather a starting point for your own research.

Help me be more useful! Please leave a 👍 if this is helpful and 👎 if it is irrelevant.