Visualise NLP dependency trees

nlp_dependency_tree.py

#!/usr/bin/env python3
#
# Requirements:
# pip3 install spacy graphviz
# python3 -m spacy download en_core_web_lg
#
import spacy
from graphviz import Digraph

nlp = spacy.load('en_core_web_lg')

class DependencyTree:
    
    def __init__(self, text):
        self.text = text
        self.dot, self.graph = self.make_tree(text)
    
    def make_tree(self, text):
        dot = Digraph()
        graph = []
        self.doc = nlp(text)
        for token in self.doc:
            s, label = str(token.i), f"{token.text} {token.pos_}"
            dot.node(s, label)
            graph.append((int(s), 'label', label))
            for t in token.children:
                s, p, o = str(token.i), t.dep_, str(t.i)
                dot.edge(s, o, label=p)
                graph.append((int(s), p, int(o)))
        return dot, graph
        
tree = DependencyTree("Today is very hot, but it's a good day.")

# visualise tree in Jupyter
tree.dot

# export to PNG
tree.dot.render('my_tree', format='png', cleanup=True)

# list edges
tree.graph

# spaCy object
tree.doc