GraphDataset / SubgraphDataset with bibliographic data (OpenAlex)

This notebook shows cvcdocdb’s torch_dataloader on top of the sample bibliographic dataset (cvcdocdb.exemples.load_bibliografia_openalex):

• GraphDataset + GraphDataLoader: flat streaming over nodes (Paper, Author).

• SubgraphDataset + PyGDataLoader: k-hop ego subgraphs as torch_geometric.data.Data objects, with an embedding per Paper (data.x and data.emb).

• MetaPath2Vec: quick embedding training on the heterogeneous Author-Paper graph, following the metapath author -writes-> paper -cites-> paper -written_by-> author.

Since the OpenAlex dataset doesn’t carry real embeddings, we first generate a deterministic embedding (hash of the title) purely for demonstration — and later train a real embedding with MetaPath2Vec, which you can use to replace it.

import importlib.util

# Comprovació de presència del paquet
package_to_check = 'cvcdocdb'
spec = importlib.util.find_spec(package_to_check)

if spec is None:
    print(f'⚠️ {package_to_check} no està instal·lat. Iniciant instal·lació...')
    %pip install -q --upgrade cvcdocdb
    print("✅ Instal·lació completada. L'estat del kernel PODRIA requerir un reinici.")
else:
    print(f'✅ {package_to_check} ja està present al sistema. Saltant instal·lació.')

import importlib.util

for package_to_check in ('torch', 'torch_geometric'):
    spec = importlib.util.find_spec(package_to_check)
    if spec is None:
        print(f'⚠️ {package_to_check} no està instal·lat. Iniciant instal·lació...')
        %pip install -q {package_to_check}
        print("✅ Instal·lació completada. L'estat del kernel PODRIA requerir un reinici.")
    else:
        print(f'✅ {package_to_check} ja està present al sistema. Saltant instal·lació.')

1. Load the bibliographic dataset into NetworkXGraph

from cvcdocdb import NetworkXGraph
from cvcdocdb.exemples import load_bibliografia_openalex

store = NetworkXGraph()
stats = load_bibliografia_openalex(store, query="graph database", per_page=15)
print("Dataset stats:", stats)

2. Add a demo embedding to each Paper

load_bibliografia_openalex doesn’t generate embeddings. Here we compute a deterministic one (hash of the title -> normalised vector) and add it as the embedding attribute of each Paper node with insertNode(..., update=True), which merges the attribute without touching the rest of the node.

import hashlib

from cvcdocdb.base import Node

EMB_DIM = 16


def toy_embedding(text: str, dim: int = EMB_DIM) -> list[float]:
    """Deterministic embedding (demo only) derived from the hash of the text."""
    digest = hashlib.sha256(text.encode("utf-8")).digest()
    raw = [digest[i % len(digest)] / 255.0 for i in range(dim)]
    norm = sum(v * v for v in raw) ** 0.5 or 1.0
    return [v / norm for v in raw]


paper_ids = store.find_nodes_by_property("main_label", "Paper")
for nid in paper_ids:
    attrs = store.get_node_attrs(nid)
    title = attrs.get("title", "")
    openalex_id = attrs.get("openalex_id")
    embedding = toy_embedding(title)
    store.insertNode(
        Node(pk={"openalex_id": openalex_id}, main_label="Paper", embedding=embedding),
        update=True,
    )

print(f"Added embeddings to {len(paper_ids)} papers (dim={EMB_DIM}).")

3. Flat streaming with GraphDataset / GraphDataLoader

from cvcdocdb.torch_dataloader import GraphDataset, GraphDataLoader

flat_ds = GraphDataset(store, label_filter="Paper")
flat_loader = GraphDataLoader(flat_ds, batch_size=8, num_workers=0)

batch = next(iter(flat_loader))
print("node_id:", batch["node_id"])
print("main_label:", batch["main_label"][:3], "...")

4. Ego subgraphs with embeddings via SubgraphDataset / PyGDataLoader

Each Paper is a seed; the k-hop subgraph includes its co-authors and cited/citing papers. vector_attr="embedding" merges the embedding into data.x and also keeps it accessible separately on data.emb.

from cvcdocdb.torch_dataloader import SubgraphDataset, PyGDataLoader

sub_ds = SubgraphDataset(
    store,
    hops=1,
    node_attrs=["year"],
    vector_attr="embedding",
    label_filter="Paper",
    missing=0.0,
)
sub_loader = PyGDataLoader(sub_ds, batch_size=4)

pyg_batch = next(iter(sub_loader))
print("x:", pyg_batch.x.shape)          # [N, 1 (year) + 16 (embedding)]
print("emb:", pyg_batch.emb.shape)      # [N, 16]
print("edge_index:", pyg_batch.edge_index.shape)
print("node_ids:", pyg_batch.node_ids[:10])

5. Training embeddings with MetaPath2Vec

Instead of section 2’s “toy” embedding, we train a real embedding with torch_geometric.nn.MetaPath2Vec on the full heterogeneous Author/Paper graph, following the metapath:

Author --AUTHORED--> Paper --CITES--> Paper --AUTHORED_BY--> Author

to_hetero_edge_index_dict loads the whole graph in a single pass over the edges (no node attributes, as lightweight as possible) and returns the edge_index_dict, the node count per type, and the original-id -> local-index mapping that MetaPath2Vec needs.

Since the dataset is small (an OpenAlex sample of ~15 papers), training is brief (few epochs) and purely illustrative — on a real dataset you should increase walk_length, walks_per_node, context_size, and the number of epochs.

import torch
from torch_geometric.nn import MetaPath2Vec

from cvcdocdb.torch_dataloader import to_hetero_edge_index_dict

edge_index_dict, num_nodes_dict, node_maps = to_hetero_edge_index_dict(store)
print("Node types:", num_nodes_dict)
print("Relation types:", list(edge_index_dict))

# Add the reverse Paper->Author relation so the metapath can close the cycle.
authored = edge_index_dict[("Author", "AUTHORED", "Paper")]
edge_index_dict[("Paper", "AUTHORED_BY", "Author")] = authored.flip(0)

if ("Paper", "CITES", "Paper") not in edge_index_dict:
    # Small sample: there may be no intra-sample citations. Add self-loops
    # so the paper->paper step of the metapath stays traversable.
    n_papers = num_nodes_dict["Paper"]
    idx = torch.arange(n_papers, dtype=torch.long)
    edge_index_dict[("Paper", "CITES", "Paper")] = torch.stack([idx, idx])

metapath = [
    ("Author", "AUTHORED", "Paper"),
    ("Paper", "CITES", "Paper"),
    ("Paper", "AUTHORED_BY", "Author"),
]

model = MetaPath2Vec(
    edge_index_dict,
    embedding_dim=16,
    metapath=metapath,
    walk_length=4,
    context_size=2,
    walks_per_node=3,
    num_negative_samples=3,
    num_nodes_dict=num_nodes_dict,
    sparse=True,
)

loader = model.loader(batch_size=8, shuffle=True)
optimizer = torch.optim.SparseAdam(list(model.parameters()), lr=0.01)

model.train()
for epoch in range(5):
    total_loss = 0.0
    for pos_rw, neg_rw in loader:
        optimizer.zero_grad()
        loss = model.loss(pos_rw, neg_rw)
        loss.backward()
        optimizer.step()
        total_loss += float(loss)
    print(f"epoch {epoch + 1}: loss={total_loss / max(len(loader), 1):.4f}")

paper_embeddings = model("Paper").detach()
print("Paper embeddings (MetaPath2Vec):", paper_embeddings.shape)

With the learned embedding we can, for example, look at which sample papers are most similar to each other (cosine similarity) — a typical use of these embeddings is to feed them back into the graph (insertNode(..., update=True)) to use them with SubgraphDataset(vector_attr=...) as in section 4.

import torch.nn.functional as F

sample = paper_embeddings[: min(5, len(paper_ids))]
sim = F.cosine_similarity(sample.unsqueeze(1), sample.unsqueeze(0), dim=-1)
print("Cosine similarity between the first papers:")
print(sim)

6. Link prediction: predicting citations between papers

With Paper embeddings already trained, let’s run a simple link prediction example: predicting whether a CITES relation exists between two papers. Since MetaPath2Vec only learns node embeddings, we derive an edge representation by combining both endpoints with edge_embeddings(..., op="dot") — the dot product is directly a similarity/existence score.

We split the CITES edges by year (split_edges_by_node_property): training on “older” paper pairs and validating on “recent” ones. Since we also need negative examples (pairs with no citation), we use torch_geometric.utils.negative_sampling.

import torch
from torch_geometric.utils import negative_sampling

from cvcdocdb.torch_dataloader import edge_embeddings, split_edges_by_node_property

cites_edge_index = edge_index_dict[("Paper", "CITES", "Paper")]

# Publication year per paper, aligned with node_maps["Paper"] (local index).
local_to_orig_paper = {v: k for k, v in node_maps["Paper"].items()}
years = torch.tensor([
    float((store.get_node_attrs(local_to_orig_paper[i]) or {}).get("year") or 0)
    for i in range(num_nodes_dict["Paper"])
])

cutoff = years.median()
splits = split_edges_by_node_property(
    cites_edge_index, years,
    train=lambda s, d: torch.maximum(s, d) <= cutoff,
    test=lambda s, d: torch.minimum(s, d) >= cutoff,
)
print("Cutoff (median year):", cutoff.item())
print("CITES edges — train:", splits["train"].shape[1], " test:", splits["test"].shape[1])

n_papers = num_nodes_dict["Paper"]


def score_split(pos_edge_index, name):
    if pos_edge_index.shape[1] == 0:
        print(f"{name}: no positive edges, skipping.")
        return
    neg_edge_index = negative_sampling(
        edge_index=cites_edge_index,
        num_nodes=n_papers,
        num_neg_samples=pos_edge_index.shape[1],
    )
    pos_scores = edge_embeddings(paper_embeddings, pos_edge_index, op="dot")
    neg_scores = edge_embeddings(paper_embeddings, neg_edge_index, op="dot")

    preds = torch.cat([pos_scores, neg_scores])
    labels = torch.cat([torch.ones_like(pos_scores), torch.zeros_like(neg_scores)])
    accuracy = ((preds > 0).float() == labels).float().mean()

    print(f"{name}: mean positive score={pos_scores.mean():.3f}  "
          f"negative={neg_scores.mean():.3f}  accuracy(threshold=0)={accuracy:.2%}")


score_split(splits["train"], "train")
score_split(splits["test"], "test")

store.close()