Graph Learning for NLP: Text-to-Graph Pipelines
Language models process text as sequences. But language is also inherently relational: words modify other words, sentences reference shared entities, documents share themes and authors. Representing text as a graph rather than a sequence opens up a different set of inductive biases. Graph neural networks trained on text graphs can capture non-local dependencies and structural patterns that are harder to express in positional sequence models.
This article covers the main text-to-graph representations, how to convert text data into TGraphX Graph objects, how to initialize node features from pre-trained embeddings, and what a GNN over text graphs can and cannot do compared to standard sequence models.
Why Represent Text as a Graph
Three common graph representations of text motivate the approach differently:
Dependency parse graphs represent a sentence's syntactic structure. Each word is a node and directed edges encode grammatical dependencies (subject, object, modifier). For tasks like relation extraction, where the relationship between two entities depends on the syntactic path between them, a GNN that can traverse this path is more natural than a sequence model that must learn to attend across arbitrary positions.
Entity co-occurrence graphs represent a document collection. Each entity (person, organization, location) is a node, and edges connect entities that appear in the same sentence or document. This graph captures which entities are discussed together, making it useful for knowledge base population, entity clustering, and topic modeling.
Document graphs (TextRank-style) represent a single document as a graph over sentences or paragraphs. Nodes are sentences; edges connect sentences that share vocabulary or semantic similarity. This is one of the original graph-based NLP methods, used for extractive summarization.
Prerequisites
This pipeline requires:
- TGraphX installed (pip install tgraphx)
- A text processing library (spaCy for dependency parsing, or a sentence encoder)
- Basic familiarity with PyTorch tensor operations
The text-to-graph conversion is handled with standard Python; the GNN part uses TGraphX. Text encoding is done with external libraries — TGraphX does not bundle a sentence encoder.
Building a Dependency Parse Graph
Using spaCy for dependency parsing, each token becomes a node and dependency edges become graph edges:
import torch
import spacy
from tgraphx import Graph
nlp = spacy.load("en_core_web_sm")
def sentence_to_dep_graph(sentence: str, embedding_model) -> Graph:
doc = nlp(sentence)
tokens = list(doc)
num_nodes = len(tokens)
# Dependency edges: each token points to its head
sources = []
targets = []
for token in tokens:
if token.head != token: # skip root self-loops
sources.append(token.i)
targets.append(token.head.i)
# Make undirected
edge_index = torch.tensor(
[sources + targets, targets + sources],
dtype=torch.long
)
# Node features: word embeddings from external encoder
token_texts = [token.text for token in tokens]
node_features = embedding_model.encode(token_texts) # [N, D]
node_features = torch.tensor(node_features, dtype=torch.float)
return Graph(
node_features=node_features,
edge_index=edge_index,
)
# Example usage (requires a sentence encoder like SentenceTransformers)
# from sentence_transformers import SentenceTransformer
# encoder = SentenceTransformer("all-MiniLM-L6-v2")
# g = sentence_to_dep_graph("The researchers published their findings.", encoder)
Building an Entity Co-Occurrence Graph
For a document collection, construct a graph where entities are nodes and edges connect entities that co-occur in the same sentence:
import torch
import spacy
from collections import defaultdict
from tgraphx import Graph
nlp = spacy.load("en_core_web_sm")
def build_cooccurrence_graph(documents: list, embedding_model) -> Graph:
# Extract entities and their co-occurrences
entity_set = set()
cooccurrences = defaultdict(int)
for doc_text in documents:
doc = nlp(doc_text)
for sent in doc.sents:
sent_entities = [ent.text.lower() for ent in sent.ents]
entity_set.update(sent_entities)
for i, e1 in enumerate(sent_entities):
for e2 in sent_entities[i+1:]:
pair = tuple(sorted([e1, e2]))
cooccurrences[pair] += 1
# Build index
entity_list = sorted(entity_set)
entity_to_idx = {e: i for i, e in enumerate(entity_list)}
num_entities = len(entity_list)
# Build edge index (only include edges above a minimum co-occurrence count)
min_cooccurrence = 2
sources, targets = [], []
for (e1, e2), count in cooccurrences.items():
if count >= min_cooccurrence:
i, j = entity_to_idx[e1], entity_to_idx[e2]
sources.extend([i, j]) # undirected
targets.extend([j, i])
if not sources:
# Fallback: create a single isolated edge if no co-occurrences meet threshold
sources, targets = [0], [0]
edge_index = torch.tensor([sources, targets], dtype=torch.long)
# Encode entity names as node features
node_features = torch.tensor(
embedding_model.encode(entity_list),
dtype=torch.float
)
return Graph(
node_features=node_features,
edge_index=edge_index,
), entity_list
# The returned entity_list maps node index -> entity name for interpretability
Node Features from Transformer Embeddings
The quality of text graph learning depends heavily on node feature initialization. Randomly initialized features tell the GNN nothing about the text; pre-trained embeddings give the GNN a strong starting point.
A practical pattern is to use sentence-transformers or static fastText embeddings for word-level features, or BERT-style contextual embeddings for sentence/entity features. These are external to TGraphX:
# Pattern: pre-encode text to tensors, then pass to Graph
# This works with any embedding library that returns numpy arrays or tensors
import numpy as np
def encode_texts_to_tensor(texts: list, model) -> torch.Tensor:
"""
Convert a list of text strings to a feature matrix.
Model can be any encoder with an .encode() method.
Returns: [N, D] tensor
"""
embeddings = model.encode(texts, batch_size=64, show_progress_bar=False)
if isinstance(embeddings, np.ndarray):
return torch.tensor(embeddings, dtype=torch.float)
return embeddings.float()
For very large vocabularies, it is practical to pre-encode all text features once and save them to disk, then load them at training time rather than re-encoding every epoch.
A Complete Text Classification Pipeline
Here is a minimal end-to-end pipeline for document-level classification using entity co-occurrence graphs:
import torch
import torch.nn as nn
import torch.nn.functional as F
from tgraphx import Graph, GraphBatch
from tgraphx.layers.gat import TensorGATLayer
from tgraphx.layers.pooling import GlobalMeanPool
class TextGraphClassifier(nn.Module):
def __init__(self, in_dim, hidden_dim, num_classes, num_heads=4):
super().__init__()
self.gat1 = TensorGATLayer(
in_channels=in_dim,
out_channels=hidden_dim,
num_heads=num_heads,
)
self.gat2 = TensorGATLayer(
in_channels=hidden_dim * num_heads,
out_channels=hidden_dim,
num_heads=1,
)
self.pool = GlobalMeanPool()
self.classifier = nn.Sequential(
nn.Linear(hidden_dim, hidden_dim // 2),
nn.ReLU(),
nn.Dropout(0.3),
nn.Linear(hidden_dim // 2, num_classes),
)
def forward(self, x, edge_index, batch):
x = F.elu(self.gat1(x, edge_index))
x = F.elu(self.gat2(x, edge_index))
x = self.pool(x, batch)
return self.classifier(x)
# Assume graph_list is a list of (Graph, label) pairs
# Each graph is a document's entity co-occurrence graph
model = TextGraphClassifier(in_dim=384, hidden_dim=64, num_classes=4)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
def collate_text_graphs(batch_list):
graphs = [item[0] for item in batch_list]
labels = torch.tensor([item[1] for item in batch_list])
return GraphBatch.from_graphs(graphs), labels
from torch.utils.data import DataLoader
loader = DataLoader(
graph_dataset, batch_size=16, shuffle=True, collate_fn=collate_text_graphs
)
for epoch in range(50):
model.train()
for batched_graphs, labels in loader:
optimizer.zero_grad()
logits = model(batched_graphs.node_features, batched_graphs.edge_index, batched_graphs.batch)
loss = F.cross_entropy(logits, labels)
loss.backward()
optimizer.step()
Graph Construction Choices and Their Effects
The choice of how to build the graph has large effects on what the GNN can learn:
Sliding window co-occurrence creates edges between any two words within a fixed window size (e.g., 3 words). This is fast to compute and captures local context, but the resulting graphs are very dense for long documents.
Semantic similarity edges connect nodes (sentences or entities) with cosine similarity above a threshold. This produces sparser, semantically meaningful graphs but requires computing all pairwise similarities — O(N^2) for N nodes.
Dependency edges produce sparse graphs tied to syntactic structure. They work well for sentence-level tasks but require a parser, which adds a preprocessing dependency.
Citation or hyperlink edges use external link structure rather than content-derived edges. These are often available in document corpora (Wikipedia, scientific papers) and can be very informative.
There is no single best choice. The appropriate graph construction depends on the task, the document type, and how much preprocessing cost is acceptable.
Limitations vs Sequence Models
Text graph approaches have real tradeoffs compared to transformer-based sequence models:
Transformers are pre-trained at massive scale. A BERT model trained on 3.3 billion words has seen linguistic patterns across nearly the entire web. A GNN trained on your text graph dataset sees only that dataset's structure. The pre-training advantage of sequence models is enormous and difficult to match with graph methods alone.
Graph construction introduces irreversible information loss. Once you have converted text to a graph, the sequential order of words is gone (unless explicitly encoded as positional features). For tasks where word order matters (e.g., sentiment about "not good" vs "good"), the graph representation may discard the ordering information that makes the difference.
GNNs over text graphs are transductive by default. A graph built from a fixed document collection cannot easily incorporate new documents without rebuilding the graph. Sequence models handle new text naturally as new input.
Entity extraction quality caps graph quality. If your named entity recognizer has 85% precision and 80% recall, your entity co-occurrence graph contains 15% spurious edges and is missing 20% of real edges. GNN performance is bounded by preprocessing quality.
Dependency parse graphs require language-specific parsers. Extending to multilingual settings requires parsers for each language, adding infrastructure complexity.
The strongest use case for text graphs is when you have a fixed document collection with rich relational structure (citation graphs, knowledge bases, forum discussions) and you want to combine document content with network structure. In that setting, GNNs can genuinely outperform sequence-only models because they directly model the relational information.
Connecting to TGraphX Modules
Once text is represented as a Graph object, the full TGraphX ecosystem is available:
- Use
tgraphx.mining.communitiesfor document clustering via community detection - Use
tgraphx.mining.link_predictionto predict entity relationships - Use
tgraphx.explainto identify which entities or words drove a classification decision - Use
tgraphx.dashboardto generate shareable reports from text classification experiments
For the graph mining tools referenced above, see the graph mining overview. For link prediction over entity graphs, see the link prediction scoring functions guide.
Frequently Asked Questions
Can I use TGraphX for knowledge graph completion from text? Yes, with some setup. The tgraphx.kg module handles knowledge graph link prediction; the text-to-graph pipeline provides the entity embeddings. See the knowledge graph embedding tutorial for the KG side of this pipeline.
How do I handle documents of very different lengths? Variable-size graphs are handled naturally by TGraphX's GraphBatch. Longer documents produce larger graphs, which is fine. The pooling step in the classifier aggregates nodes regardless of graph size.
Should I use GNN or BERT for text classification? For most standard text classification benchmarks, BERT and its variants outperform GNN-based methods. GNN approaches are competitive or superior on tasks that explicitly involve network structure — citation classification, knowledge base completion, dialogue act recognition in conversations.
What node feature dimension should I use? Match the output dimension of your text encoder. Common sizes: 384 (all-MiniLM-L6-v2), 768 (BERT-base), 1536 (larger models). The GNN's first layer will project from this dimension to your chosen hidden dimension.