class Graph:
""" The graph of nodes and edges
...
Attributes
----------
nodes : List[Node]
List of nodes in the graph
num_of_nodes : int
Total number of nodes
num_of_edges : List[int]
List of number of edges of nodes
vectorizer : Vectorizer
Vectorizer object that return vector representation of entities
Methods
-------
add_node(node)
Adds a node to graph
create_graph()
Creates all edges between nodes based on cosine-similarity
"""
def __init__(self, term_dict: Dict[str, Term]):
self.nodes = []
self.num_of_nodes = 0
self.num_of_edges = []
self.vectorizer = Vectorizer(term_dict)
def add_node(self, node: Node):
self.nodes.append(node)
self.num_of_nodes += 1
def create_graph(self):
if type(self.nodes[0].body) == Document:
self.vectorizer.vectorize_docs([node.body for node in self.nodes])
else:
self.vectorizer.vectorize_sentences(
[node.body for node in self.nodes])
self.num_of_edges = [1 for i in range(self.num_of_nodes)]
for i in range(self.num_of_nodes):
for j in range(i + 1, self.num_of_nodes):
node1 = self.nodes[i]
node2 = self.nodes[j]
similarity = cos_sim(node1.body.vector, node2.body.vector)
if similarity > 0.10:
node1.add_edge(Edge(node2.id, similarity))
node2.add_edge(Edge(node1.id, similarity))
self.num_of_edges[i] += 1
self.num_of_edges[j] += 1
