gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add data visualization for Anthropic

Browse source
This commit is contained in:
vasilije 2025-01-14 00:40:14 +01:00
parent cf4737b0a2
commit 55e9d64d87
1 changed files with 5 additions and 25 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

30
cognee/shared/utils.py
View file

@ -11,7 +11,7 @@ import networkx as nx
import pandas as pd import pandas as pd
import matplotlib.pyplot as plt import matplotlib.pyplot as plt
import tiktoken import tiktoken
import nltk
import base64 import base64
import time import time
@ -243,31 +243,9 @@ async def render_graph(
# return df.replace([np.inf, -np.inf, np.nan], None) # return df.replace([np.inf, -np.inf, np.nan], None)
def get_entities(tagged_tokens):
nltk.download("maxent_ne_chunker", quiet=True)
from nltk.chunk import ne_chunk
return ne_chunk(tagged_tokens)
def extract_pos_tags(sentence):
"""Extract Part-of-Speech (POS) tags for words in a sentence."""
# Ensure that the necessary NLTK resources are downloaded
nltk.download("words", quiet=True)
nltk.download("punkt", quiet=True)
nltk.download("averaged_perceptron_tagger", quiet=True)
from nltk.tag import pos_tag
from nltk.tokenize import word_tokenize
# Tokenize the sentence into words
tokens = word_tokenize(sentence)
# Tag each word with its corresponding POS tag
pos_tags = pos_tag(tokens)
return pos_tags
logging.basicConfig(level=logging.INFO) logging.basicConfig(level=logging.INFO)
@ -450,13 +428,15 @@ async def create_cognee_style_network_with_logo(
# Construct your filename # Construct your filename
filename = f"{timestamp}.png" filename = f"{timestamp}.png"
export_png(p, filename=filename)
logging.info(f"Saving visualization to {output_filename}...") logging.info(f"Saving visualization to {output_filename}...")
html_content = file_html(p, CDN, title) html_content = file_html(p, CDN, title)
with open(output_filename, "w") as f: with open(output_filename, "w") as f:
f.write(html_content) f.write(html_content)
logging.info("Visualization complete.") logging.info("Visualization complete.")
if bokeh_object: if bokeh_object:
@ -517,7 +497,7 @@ if __name__ == "__main__":
G, G,
output_filename="example_network.html", output_filename="example_network.html",
title="Example Cognee Network", title="Example Cognee Network",
node_attribute="group", # Attribute to use for coloring nodes label="group", # Attribute to use for coloring nodes
layout_func=nx.spring_layout, # Layout function layout_func=nx.spring_layout, # Layout function
layout_scale=3.0, # Scale for the layout layout_scale=3.0, # Scale for the layout
logo_alpha=0.2, logo_alpha=0.2,