Spaces:
Sleeping
Sleeping
File size: 3,796 Bytes
fb793b8 b0964ac |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 |
import os
import pandas as pd
import networkx as nx
from pyvis.network import Network
class CharacterNetworkGenerator:
def __init__(self):
pass
def generate_char_network(df):
windows = 10
entity_relationship = []
for row in df["chars"]:
prev_entity_window = []
for sentence in row:
# each sentence = ["Ted", "Lilly"]
prev_entity_window.append(list(sentence))
# We keep only the last 10 entities as previous.
prev_entity_window = prev_entity_window[-windows:]
# Flatten 2D list into 1D list
prev_entity_flattened = sum(prev_entity_window, [])
# Build relationship for each entity.
for entity in sentence:
# Check each entity with all previous 10 entities.
for entity_in_window in prev_entity_flattened:
# if they aren't same, append them because they are related.
if entity != entity_in_window:
# Sort them because (ted, lilly is same as lilly, ted.)
entity_relationship.append(sorted([entity, entity_in_window]))
relationship_df = pd.DataFrame({"value": entity_relationship})
relationship_df["source"] = relationship_df["value"].apply(lambda x: x[0])
relationship_df["target"] = relationship_df["value"].apply(lambda x: x[1])
relationship_df = relationship_df.groupby(["source", "target"]).count().reset_index()
relationship_df = relationship_df.sort_values("value", ascending=False)
return relationship_df
def draw_char_network(self, df):
df = df.sort_values("value", ascending=False).head(200)
G = nx.from_pandas_edgelist(
df,
source="source",
target="target",
edge_attr="value",
create_using=nx.Graph()
)
net = Network(
notebook=True, width="1000px",
height="700px",
bgcolor="#222222",
font_color="white",
cdn_resources="remote"
)
node_degree = dict(G.degree)
nx.set_node_attributes(G, node_degree, "size")
net.from_nx(G)
html = net.generate_html()
html = html.replace("'", "\"")
output_html = f"""<iframe style="width: 100%; height: 600px;margin:0 auto" name="result" allow="midi; geolocation; microphone; camera;
display-capture; encrypted-media;" sandbox="allow-modals allow-forms
allow-scripts allow-same-origin allow-popups
allow-top-navigation-by-user-activation allow-downloads" allowfullscreen=""
allowpaymentrequest="" frameborder="0" srcdoc='{html}'></iframe>"""
return output_html
def defaultGraph(self, save_path):
if save_path and not save_path.endswith(".html"):
save_path += "network.html"
with open("characters\himym.html", "r") as f:
html = f.read()
html = html.replace("'", "\"")
output_html = f"""<iframe style="width: 100%; height: 600px;margin:0 auto" name="result" allow="midi; geolocation; microphone; camera;
display-capture; encrypted-media;" sandbox="allow-modals allow-forms
allow-scripts allow-same-origin allow-popups
allow-top-navigation-by-user-activation allow-downloads" allowfullscreen=""
allowpaymentrequest="" frameborder="0" srcdoc='{html}'></iframe>"""
try:
if save_path:
with open("../" + save_path, "w") as f:
f.write(output_html)
except:
pass
return output_html |