Add Option for Knowledge Graph Similarities #487

logan-markewich · 2023-02-20T01:56:01Z

This PR adds support for using embedding similarities within the knowledge graph index.

You may have noticed that relying purely on keyword mapping for finding relation tuples may result in no relation tuples being found! If you didn't turn on the include_text=True option, then chances are the model will hallucinate an answer.

This PR allows you to use the existing keyword mapping, the top_k related knowledge graph tuples, or a combination of both!

The following is some example usage using both keywords and embeddings:

from gpt_index.data_structs.data_structs import  MODE_KEYWORDS, MODE_EMBEDDINGS
from gpt_index.indices.knowledge_graph.base import GPTKnowledgeGraphIndex

index = GPTKnowledgeGraphIndex(documents, llm_predictor=llm_predictor, mode=MODE_EMBEDDINGS+MODE_KEYWORDS)
response = index.query("What did the author do growing up?", include_text=False, similarity_cutoff=0.8, top_k_embeddings=1, response_mode="tree_summarize")

In my testing, the similarity cutoff is pretty important. Unrelated nodes are added to the prompt, it ends up hurting the response accuracy (maybe the prompt itself needs some investigation too 🤔)

logan-markewich · 2023-02-20T02:49:12Z

Ah maybe don't look at this quite yet, just realized _should_use_node filters out everything when similarity_cutoff is set

logan-markewich · 2023-02-20T04:37:05Z

Ok bug is fixed! (Also I need to figure out my flake8 setup, it wasn't catching the same rules as the automated checks)

jerryjliu

This is amazing! I have a high-level comment about how to structure this a bit more nicely, let me know if you agree!

Once this PR is ready to land, let me know if you have a notebook sample demonstrating this capability - this would be a great feature to publicize :)

gpt_index/data_structs/data_structs.py

gpt_index/indices/knowledge_graph/base.py

gpt_index/indices/query/knowledge_graph/query.py

gpt_index/indices/knowledge_graph/base.py

logan-markewich · 2023-02-21T14:37:35Z

Thanks for the comments! I'll work to de-couple the query from the index, and add some examples of the usage to a notebook 👍🏻

logan-markewich · 2023-02-21T21:00:26Z

Alright, I think I've addressed your comments 🚀

I also added a mechanism to use the "include_text" option for triplets that are found using the embeddings

Just working on getting some examples, but I'm struggling a bit with the Paul Graham essay. Just trying to find a few queries where it can really shine. I might try with a different example dataset 🤔

I think a lot of it also depends on the quality of triplets (which could be alleviated by inserting 3rd party triplets directly into the index struct 😉)

jerryjliu

this looks pretty good - some minor comments. let me know if you need any help on the notebook as well!

gpt_index/indices/knowledge_graph/base.py

gpt_index/indices/query/knowledge_graph/query.py

jerryjliu · 2023-02-22T06:50:53Z

gpt_index/indices/query/knowledge_graph/query.py

@@ -93,6 +154,9 @@ def _get_nodes_for_response(
        # filter sorted nodes
        sorted_nodes = [node for node in sorted_nodes if self._should_use_node(node)]
        for chunk_idx, node in zip(sorted_chunk_indices, sorted_nodes):
+            # nodes are found with keyword mapping, give high conf to avoid cutoff
+            if similarity_tracker is not None:
+                similarity_tracker.add(node, 1000.0)


this is a bit hacky (the num 1000 is a bit arbitrary), but it's ok, i'll take a stab at this later

Definitely hacky, without the comment this line will look pretty strange.

Maybe to be less hacky, we can override the get_nodes_and_similarities_for_response or _should_use_node? to avoid the original problem of nodes being removed

tests/indices/knowledge_graph/test_base.py

logan-markewich · 2023-02-23T05:00:09Z

Ok, addressed basically everything. Not 100% sure where to put the new enum, so its in the query file for now.

Still working on the example, but could maybe use a hand if you have time.

I was using the NYC wiki data example, I'll attach the index to save you the tokens (constructed using the same index options from the knowledge graph example notebook)

If you don't have time that's totally fine! I'll explore this more after work tomorrow otherwise. 💪🏻
index_kg_emb_nyc.zip

jerryjliu · 2023-02-23T09:03:24Z

If you don't have time that's totally fine! I'll explore this more after work tomorrow otherwise. 💪🏻 index_kg_emb_nyc.zip

thanks for sending over @logan-markewich! I might not have time until friday or this weekend or so to take a deeper look, but can get to it then if you'd like me to. just lmk

jerryjliu

oops sorry i accidentally meant to approve another PR! putting a hold on this until i get back to it

logan-markewich · 2023-02-25T04:17:44Z

Settled on an example inside the existing KG notebook. I'm pretty happy with the current state of this PR 💪🏻

jerryjliu · 2023-02-25T05:44:53Z

Settled on an example inside the existing KG notebook. I'm pretty happy with the current state of this PR 💪🏻

sweet! will take a look in a bit

jerryjliu

looks good! and amazing work on the demo notebook! I'll take a look at fiing some of these nits before merging

jerryjliu · 2023-02-25T06:57:09Z

gpt_index/indices/query/knowledge_graph/query.py

+        KEYWORD ("keyword"): Default query mode, using keywords to find triplets.
+        EMBEDDING ("embedding"): Embedding mode, using embeddings to find
+            similar triplets.
+        HYBRID ("hyrbid"): Hyrbid mode, combining both keywords and embeddings


nit: hybrid

jerryjliu · 2023-02-25T06:57:59Z

gpt_index/indices/query/knowledge_graph/query.py

+        embedding_mode (KGQueryMode): Specifies whether to use keyowrds,
+            embeddings, or both to find relevent triplets. Should be one of "keyword",
+            "embedding", or "hybrid".
+        top_k_embeddings (int): The number of top embeddings to use


nit: let's rename similarity_top_k to keep consistent with vector index queries

Logan Markewich and others added 7 commits February 17, 2023 21:05

support for keeping track of KG triplet sources

d8202ba

cr

cf8a4ec

cr

03e2453

Merge branch 'jerryjliu:main' into main

27d2cdc

Merge branch 'jerryjliu:main' into main

a4894f1

Add support for KG similarity indexes+queries

de97619

Merge branch 'jerryjliu:main' into main

7aab095

logan-markewich marked this pull request as draft February 20, 2023 01:56

logan-markewich changed the title ~~Add Option for Knowledge Graph Similarities~~ [WIP] Add Option for Knowledge Graph Similarities Feb 20, 2023

logan-markewich changed the title ~~[WIP] Add Option for Knowledge Graph Similarities~~ Add Option for Knowledge Graph Similarities Feb 20, 2023

logan-markewich marked this pull request as ready for review February 20, 2023 02:01

Logan Markewich added 2 commits February 19, 2023 20:37

Fix mypi tpye issues

7668017

Fix various linter issues

02afca1

Logan Markewich added 5 commits February 19, 2023 22:02

Fix kg similarity bug with tracker

474357f

Fix flake8 errors

c125af5

Shorten comment linewidth

a6a311d

Add docstring to new KG struct methods

c6198b8

Add period to docstring

ba8d9e5

Merge branch 'main' into main

6d0a9c7

jerryjliu reviewed Feb 21, 2023

View reviewed changes

logan-markewich and others added 5 commits February 21, 2023 12:07

Merge branch 'jerryjliu:main' into main

926999e

De-couple KG index and query for embeddings

5465d39

Use queued embeddings, linting fixes

29c0dd1

use include_text option for KG embeddings

bdf2628

Fix flake8 error

fb4c8a5

jerryjliu reviewed Feb 22, 2023

View reviewed changes

Logan Markewich added 2 commits February 22, 2023 22:50

remove print inside test

5a447ff

Use kg mode enum, bug/formatting fixes

ff628d1

jerryjliu approved these changes Feb 23, 2023

View reviewed changes

jerryjliu requested changes Feb 23, 2023

View reviewed changes

logan-markewich and others added 3 commits February 24, 2023 21:32

Merge branch 'jerryjliu:main' into main

d5d2eb0

Add small KG embeddings example

a3d584f

Fix mypy errors

0013ed6

jerryjliu approved these changes Feb 25, 2023

View reviewed changes

Jerry Liu added 2 commits February 24, 2023 23:27

cr

e1ce619

cr

4c92cbe

jerryjliu merged commit 7281dea into run-llama:main Feb 25, 2023

wey-gu mentioned this pull request Apr 25, 2023

RFC: Bring graph store to llama index #1318

Closed

3 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add Option for Knowledge Graph Similarities #487

Add Option for Knowledge Graph Similarities #487

logan-markewich commented Feb 20, 2023 •

edited

Loading

logan-markewich commented Feb 20, 2023 •

edited

Loading

logan-markewich commented Feb 20, 2023

jerryjliu left a comment

logan-markewich commented Feb 21, 2023

logan-markewich commented Feb 21, 2023 •

edited

Loading

jerryjliu left a comment

jerryjliu Feb 22, 2023

logan-markewich Feb 22, 2023 •

edited

Loading

logan-markewich commented Feb 23, 2023 •

edited

Loading

jerryjliu commented Feb 23, 2023

jerryjliu left a comment

logan-markewich commented Feb 25, 2023

jerryjliu commented Feb 25, 2023

jerryjliu left a comment

jerryjliu Feb 25, 2023

jerryjliu Feb 25, 2023

Add Option for Knowledge Graph Similarities #487

Add Option for Knowledge Graph Similarities #487

Conversation

logan-markewich commented Feb 20, 2023 • edited Loading

logan-markewich commented Feb 20, 2023 • edited Loading

logan-markewich commented Feb 20, 2023

jerryjliu left a comment

Choose a reason for hiding this comment

logan-markewich commented Feb 21, 2023

logan-markewich commented Feb 21, 2023 • edited Loading

jerryjliu left a comment

Choose a reason for hiding this comment

jerryjliu Feb 22, 2023

Choose a reason for hiding this comment

logan-markewich Feb 22, 2023 • edited Loading

Choose a reason for hiding this comment

logan-markewich commented Feb 23, 2023 • edited Loading

jerryjliu commented Feb 23, 2023

jerryjliu left a comment

Choose a reason for hiding this comment

logan-markewich commented Feb 25, 2023

jerryjliu commented Feb 25, 2023

jerryjliu left a comment

Choose a reason for hiding this comment

jerryjliu Feb 25, 2023

Choose a reason for hiding this comment

jerryjliu Feb 25, 2023

Choose a reason for hiding this comment

logan-markewich commented Feb 20, 2023 •

edited

Loading

logan-markewich commented Feb 20, 2023 •

edited

Loading

logan-markewich commented Feb 21, 2023 •

edited

Loading

logan-markewich Feb 22, 2023 •

edited

Loading

logan-markewich commented Feb 23, 2023 •

edited

Loading