mirror of
https://github.com/microsoft/graphrag.git
synced 2025-07-23 08:53:04 +00:00
57 lines
1.6 KiB
Python
57 lines
1.6 KiB
Python
|
# Copyright (c) 2024 Microsoft Corporation.
|
||
|
|
# Licensed under the MIT License
|
||
|
|
|
||
|
|
"""Full graph module."""
|
||
|
|
|
||
|
|
import altair as alt
|
||
|
|
import pandas as pd
|
||
|
|
import streamlit as st
|
||
|
|
from state.session_variables import SessionVariables
|
||
|
|
|
||
|
|
|
||
|
|
def create_full_graph_ui(sv: SessionVariables):
|
||
|
|
"""Return graph UI object."""
|
||
|
|
entities = sv.entities.value.copy()
|
||
|
|
communities = sv.communities.value.copy()
|
||
|
|
|
||
|
|
if not communities.empty and not entities.empty:
|
||
|
|
communities_entities = (
|
||
|
|
communities.explode("entity_ids")
|
||
|
|
.merge(
|
||
|
|
entities,
|
||
|
|
left_on="entity_ids",
|
||
|
|
right_on="id",
|
||
|
|
suffixes=("_entities", "_communities"),
|
||
|
|
)
|
||
|
|
.dropna(subset=["x", "y"])
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
communities_entities = pd.DataFrame()
|
||
|
|
|
||
|
|
level = sv.graph_community_level.value
|
||
|
|
communities_entities_filtered = communities_entities[
|
||
|
|
communities_entities["level"] == level
|
||
|
|
]
|
||
|
|
|
||
|
|
graph = (
|
||
|
|
alt.Chart(communities_entities_filtered)
|
||
|
|
.mark_circle()
|
||
|
|
.encode(
|
||
|
|
x="x",
|
||
|
|
y="y",
|
||
|
|
color=alt.Color(
|
||
|
|
"community",
|
||
|
|
scale=alt.Scale(
|
||
|
|
domain=communities_entities_filtered["community"].unique(),
|
||
|
|
scheme="category10",
|
||
|
|
),
|
||
|
|
),
|
||
|
|
size=alt.Size("degree", scale=alt.Scale(range=[50, 1000]), legend=None),
|
||
|
|
tooltip=["id_entities", "type", "description", "community"],
|
||
|
|
)
|
||
|
|
.properties(height=1000)
|
||
|
|
.configure_axis(disable=True)
|
||
|
|
)
|
||
|
|
st.altair_chart(graph, use_container_width=True)
|
||
|
|
return graph
|