Update app.py

app.py

@@ -19,6 +19,13 @@ def half_circle_layout(n_nodes, sentence_node=True):
   pos[n_nodes - 1] = (0, -0.25)
   return pos
 
+def get_adjacency_matrix(jraph_graph: jraph.GraphsTuple):
+  nodes, edges, receivers, senders, _, _, _ = jraph_graph
+  adj_mat = jnp.zeros((len(nodes), len(nodes)))
+  for i in range(len(receivers)):
+    adj_mat = adj_mat.at[senders[i], receivers[i]].set(1)
+  return adj_mat
+
 def dependency_parser(sentences):
   return [nlp(sentence) for sentence in sentences]