pyg-nightly 2.6.0.dev20240912__py3-none-any.whl → 2.7.0.dev20240914__py3-none-any.whl

{pyg_nightly-2.6.0.dev20240912.dist-info → pyg_nightly-2.7.0.dev20240914.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pyg-nightly
-Version: 2.6.0.dev20240912
+Version: 2.7.0.dev20240914
 Summary: Graph Neural Network Library for PyTorch
 Keywords: deep-learning,pytorch,geometric-deep-learning,graph-neural-networks,graph-convolutional-networks
 Author-email: Matthias Fey <matthias@pyg.org>

{pyg_nightly-2.6.0.dev20240912.dist-info → pyg_nightly-2.7.0.dev20240914.dist-info}/RECORD

@@ -1,4 +1,4 @@
-torch_geometric/__init__.py,sha256=zJh2N-U_IS2TvPl3dt8Cas7iUpxUQ0vxgqsO85xR8cA,1904
+torch_geometric/__init__.py,sha256=n23o0W0D6f8VcOq-d2ljeW7QxvqdEiX9hzH_YvgmLK0,1904
 torch_geometric/_compile.py,sha256=0HAdz6MGmyrgi4g6P-PorTg8dPIKx3Jo4zVJavrlfX0,1139
 torch_geometric/_onnx.py,sha256=V9ffrIKSqhDw6xUZ12lkuSfNs48cQp2EeJ6Z19GfnVw,349
 torch_geometric/backend.py,sha256=lVaf7aLoVaB3M-UcByUJ1G4T4FOK6LXAg0CF4W3E8jo,1575
@@ -7,10 +7,10 @@ torch_geometric/config_store.py,sha256=zdMzlgBpUmBkPovpYQh5fMNwTZLDq2OneqX47QEx7
 torch_geometric/debug.py,sha256=cLyH9OaL2v7POyW-80b19w-ctA7a_5EZsS4aUF1wc2U,1295
 torch_geometric/deprecation.py,sha256=dWRymDIUkUVI2MeEmBG5WF4R6jObZeseSBV9G6FNfjc,858
 torch_geometric/device.py,sha256=tU5-_lBNVbVHl_kUmWPwiG5mQ1pyapwMF4JkmtNN3MM,1224
-torch_geometric/edge_index.py,sha256=Kjkk1kW9C0D2xZPdvTKXQUA4WnqqkHWlUJXpwPG-JWc,70052
+torch_geometric/edge_index.py,sha256=r4_24Rhm3YCK0BF-kzLvL7PlY_1tWcXrBDIr7JPDjkw,70048
 torch_geometric/experimental.py,sha256=JbtNNEXjFGI8hZ9raM6-qrZURP6Z5nlDK8QicZUIbz0,4756
 torch_geometric/home.py,sha256=EV54B4Dmiv61GDbkCwtCfWGWJ4eFGwZ8s3KOgGjwYgY,790
-torch_geometric/index.py,sha256=ZDUq2LTumN1UyYkNF3tYglW8DZ1G-s2ejAIKaNcvfgI,24054
+torch_geometric/index.py,sha256=9ChzWFCwj2slNcVBOgfV-wQn-KscJe_y7502w-Vf76w,24045
 torch_geometric/inspector.py,sha256=9M61T9ruSid5-r2aelRAeX9g_7AZ1VMnYAB2KozM71E,19267
 torch_geometric/isinstance.py,sha256=truZjdU9PxSvjJ6k0d_CLJ2iOpen2o8U-54pbUbNRyE,935
 torch_geometric/lazy_loader.py,sha256=SM0UcXtIdiFge75MKBAWXedoiSOdFDOV0rm1PfoF9cE,908
@@ -53,7 +53,7 @@ torch_geometric/data/temporal.py,sha256=WOJ6gFrTLikaLhUvotyUF5ql14FkE5Ox3hNkdSp6
 torch_geometric/data/view.py,sha256=XjkVSc-UWZFCT4DlXLShZtO8duhFQkS9gq88zZXANsk,1089
 torch_geometric/data/lightning/__init__.py,sha256=w3En1tJfy3kSqe1MycpOyZpHFO3fxBCgNCUOznPA3YU,178
 torch_geometric/data/lightning/datamodule.py,sha256=Bn9iaIfE4NWDDWWMqCvBeZ4bIW1Silx_Ol5CPJCliaQ,29242
-torch_geometric/datasets/__init__.py,sha256=U8ieW-6Xb4Ha1YwjoMqsEEOYziLAweJk5vxx9TPgXqs,5816
+torch_geometric/datasets/__init__.py,sha256=fey-955PyCQXGBeUTNPWwU5uK3PJOEvaY1_fDt1SxXc,5880
 torch_geometric/datasets/actor.py,sha256=H8srMdo5qo8eg4LDxEdYcxZi49I_HVDcr8R_pb2W99Q,4461
 torch_geometric/datasets/airfrans.py,sha256=7Yt0Xs7jx2NotPT4_9GbpLRWRXYSS5g_4zSENoB_9hs,5684
 torch_geometric/datasets/airports.py,sha256=HSZdi6KM_yavppaUl0uWyZ93BEsrtDf9InjPPu9zaUE,3903
@@ -149,6 +149,7 @@ torch_geometric/datasets/tosca.py,sha256=nUSF8NQT1GlkwWQLshjWmr8xORsvRHzzIqhUyDC
 torch_geometric/datasets/tu_dataset.py,sha256=14OSaXBgVwT1dX2h1wZ3xVIwoo0GQBEfR3yWh6Q0VF0,7847
 torch_geometric/datasets/twitch.py,sha256=qfEerf-Uaojx2ZvegENowdG4E7RoUT_HUO9xtULadvo,3658
 torch_geometric/datasets/upfd.py,sha256=crqO8uQNz1wC1JOn4prSs8iOGv9LuLK3dZf_KUV9tUE,7010
+torch_geometric/datasets/web_qsp_dataset.py,sha256=OusHv0DcvDgCjUbBtkhPzwm2pdPlyG98BSzaQPv_GP8,8451
 torch_geometric/datasets/webkb.py,sha256=beC1kWeW7cIjYwWyaINQSk-3lmVR85Lus7cKZniHp8Y,4879
 torch_geometric/datasets/wikics.py,sha256=iTzYif1WvbMXnMdhPMfvrkVaAbnM009WiB_f_JWZqhU,3879
 torch_geometric/datasets/wikidata.py,sha256=9mYShF_HlpTmcdLpiaP_tYJ9eQtUOu5vRPvohN6RXqI,4979
@@ -426,7 +427,7 @@ torch_geometric/nn/models/deep_graph_infomax.py,sha256=u6j-5-iHBASDCZ776dyfCI1N8
 torch_geometric/nn/models/deepgcn.py,sha256=tIgT03cj8MghYlxEozpoGvGG_CwpJrGDxv1Z0CVIUts,4339
 torch_geometric/nn/models/dimenet.py,sha256=Kc5p-rB5q-0e8lY22l-OdQTscTxJh2lTEpeRFMdL4RY,36186
 torch_geometric/nn/models/dimenet_utils.py,sha256=xP_nbzkSSL25GC3rrZ9KP8x9QZ59S-CZuHzCmQ-K0fI,5062
-torch_geometric/nn/models/g_retriever.py,sha256=uH_aYrFbFNHaAeKQn_LtUgP5ajutLYYD8N9UvSKcpfk,7271
+torch_geometric/nn/models/g_retriever.py,sha256=VueRImNJlh1WvRWcsSXliSw8RlxlzWlu2WSFs_VQaJc,7749
 torch_geometric/nn/models/gnnff.py,sha256=15dkiLgy0LmH1hnUrpeoHioIp4BPTfjpVATpnGRt9E0,7860
 torch_geometric/nn/models/graph_mixer.py,sha256=mthMeCOikR8gseEsu4oJ3Cd9C35zHSv1p32ROwnG-6s,9246
 torch_geometric/nn/models/graph_unet.py,sha256=WFb7d_DBByMGyXh3AdK2CKNmvMmSKsSUt8l8UnSOovs,5395
@@ -449,8 +450,8 @@ torch_geometric/nn/models/signed_gcn.py,sha256=J40CnedFIqtKI1LhW1ITSEFRbA_XiJZL6
 torch_geometric/nn/models/tgn.py,sha256=kEGdfLJybkbMT4UMoAh2nCzfX3_nDjfm1cicuPHEwAM,11878
 torch_geometric/nn/models/visnet.py,sha256=97OFMCsPDEI5BCSi7RhoRcU2CNRp7zck2tEzrltFZj4,43192
 torch_geometric/nn/nlp/__init__.py,sha256=JJESTA7w_K8v60XbCd25IqmrKKHLz5OiNexMHYGV2mE,138
-torch_geometric/nn/nlp/llm.py,sha256=KwSXgI55FuHLR_9vhgekDXMaRUodPQceHPD7OCp2KN4,11639
-torch_geometric/nn/nlp/sentence_transformer.py,sha256=DzbQO8wgR34BkKpXfMqQu61hMrK94W2MBa3bZ4fDmVs,3114
+torch_geometric/nn/nlp/llm.py,sha256=a5YkJA32Ok2PmWFEJ0VJD0HfsauDpxosIwlij6wqwJo,11728
+torch_geometric/nn/nlp/sentence_transformer.py,sha256=JrTN3W1srdkNX7qYDGB08mY5615i5nfEJSTHAdd5EuA,3260
 torch_geometric/nn/norm/__init__.py,sha256=u2qIDrkbeuObGVXSAIftAlvSd6ouGTtxznCfD-59UiA,669
 torch_geometric/nn/norm/batch_norm.py,sha256=sJKrinHGwA-noIgteg1RD2W06rd0zskD-rXuY-36glY,8283
 torch_geometric/nn/norm/diff_group_norm.py,sha256=b57XvNekrUYGDjNJlGeqvaMGNJmHwopSF0_yyBWlLuA,4722
@@ -617,6 +618,6 @@ torch_geometric/utils/undirected.py,sha256=H_nfpI0_WluOG6VfjPyldvcjL4w5USAKWu2x5
 torch_geometric/visualization/__init__.py,sha256=PyR_4K5SafsJrBr6qWrkjKr6GBL1b7FtZybyXCDEVwY,154
 torch_geometric/visualization/graph.py,sha256=SvbdVx5Zmuy_WSSA4-WWCkqAcCSHVe84mjMfsEWbZCs,4813
 torch_geometric/visualization/influence.py,sha256=CWMvuNA_Nf1sfbJmQgn58yS4OFpeKXeZPe7kEuvkUBw,477
-pyg_nightly-2.6.0.dev20240912.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
-pyg_nightly-2.6.0.dev20240912.dist-info/METADATA,sha256=D-HQbnicYK-Tr_IxBWo2XuDsqcPG2_K5zrYJXg91xnQ,63068
-pyg_nightly-2.6.0.dev20240912.dist-info/RECORD,,
+pyg_nightly-2.7.0.dev20240914.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
+pyg_nightly-2.7.0.dev20240914.dist-info/METADATA,sha256=Pz5f49zwvSFV-FkxufA5cJDayObU27GV4gtdnv1KN0g,63068
+pyg_nightly-2.7.0.dev20240914.dist-info/RECORD,,

torch_geometric/__init__.py

@@ -30,7 +30,7 @@ from .lazy_loader import LazyLoader
 contrib = LazyLoader('contrib', globals(), 'torch_geometric.contrib')
 graphgym = LazyLoader('graphgym', globals(), 'torch_geometric.graphgym')
 
-__version__ = '2.6.0.dev20240912'
+__version__ = '2.7.0.dev20240914'
 
 __all__ = [
     'Index',

torch_geometric/datasets/__init__.py

@@ -61,7 +61,6 @@ from .gemsec import GemsecDeezer
 from .twitch import Twitch
 from .airports import Airports
 from .lrgb import LRGBDataset
-from .neurograph import NeuroGraphDataset
 from .malnet_tiny import MalNetTiny
 from .omdb import OMDB
 from .polblogs import PolBlogs
@@ -76,6 +75,8 @@ from .jodie import JODIEDataset
 from .wikidata import Wikidata5M
 from .myket import MyketDataset
 from .brca_tgca import BrcaTcga
+from .neurograph import NeuroGraphDataset
+from .web_qsp_dataset import WebQSPDataset
 
 from .dbp15k import DBP15K
 from .aminer import AMiner
@@ -188,6 +189,7 @@ homo_datasets = [
     'MyketDataset',
     'BrcaTcga',
     'NeuroGraphDataset',
+    'WebQSPDataset',
 ]
 
 hetero_datasets = [

torch_geometric/datasets/web_qsp_dataset.py

@@ -0,0 +1,239 @@
+# Code adapted from the G-Retriever paper: https://arxiv.org/abs/2402.07630
+from typing import Any, Dict, List, Tuple, no_type_check
+
+import numpy as np
+import torch
+from torch import Tensor
+from tqdm import tqdm
+
+from torch_geometric.data import Data, InMemoryDataset
+from torch_geometric.nn.nlp import SentenceTransformer
+
+
+@no_type_check
+def retrieval_via_pcst(
+    data: Data,
+    q_emb: Tensor,
+    textual_nodes: Any,
+    textual_edges: Any,
+    topk: int = 3,
+    topk_e: int = 3,
+    cost_e: float = 0.5,
+) -> Tuple[Data, str]:
+    c = 0.01
+    if len(textual_nodes) == 0 or len(textual_edges) == 0:
+        desc = textual_nodes.to_csv(index=False) + "\n" + textual_edges.to_csv(
+            index=False,
+            columns=["src", "edge_attr", "dst"],
+        )
+        return data, desc
+
+    from pcst_fast import pcst_fast
+
+    root = -1
+    num_clusters = 1
+    pruning = 'gw'
+    verbosity_level = 0
+    if topk > 0:
+        n_prizes = torch.nn.CosineSimilarity(dim=-1)(q_emb, data.x)
+        topk = min(topk, data.num_nodes)
+        _, topk_n_indices = torch.topk(n_prizes, topk, largest=True)
+
+        n_prizes = torch.zeros_like(n_prizes)
+        n_prizes[topk_n_indices] = torch.arange(topk, 0, -1).float()
+    else:
+        n_prizes = torch.zeros(data.num_nodes)
+
+    if topk_e > 0:
+        e_prizes = torch.nn.CosineSimilarity(dim=-1)(q_emb, data.edge_attr)
+        topk_e = min(topk_e, e_prizes.unique().size(0))
+
+        topk_e_values, _ = torch.topk(e_prizes.unique(), topk_e, largest=True)
+        e_prizes[e_prizes < topk_e_values[-1]] = 0.0
+        last_topk_e_value = topk_e
+        for k in range(topk_e):
+            indices = e_prizes == topk_e_values[k]
+            value = min((topk_e - k) / sum(indices), last_topk_e_value - c)
+            e_prizes[indices] = value
+            last_topk_e_value = value * (1 - c)
+        # reduce the cost of the edges such that at least one edge is selected
+        cost_e = min(cost_e, e_prizes.max().item() * (1 - c / 2))
+    else:
+        e_prizes = torch.zeros(data.num_edges)
+
+    costs = []
+    edges = []
+    virtual_n_prizes = []
+    virtual_edges = []
+    virtual_costs = []
+    mapping_n = {}
+    mapping_e = {}
+    for i, (src, dst) in enumerate(data.edge_index.t().numpy()):
+        prize_e = e_prizes[i]
+        if prize_e <= cost_e:
+            mapping_e[len(edges)] = i
+            edges.append((src, dst))
+            costs.append(cost_e - prize_e)
+        else:
+            virtual_node_id = data.num_nodes + len(virtual_n_prizes)
+            mapping_n[virtual_node_id] = i
+            virtual_edges.append((src, virtual_node_id))
+            virtual_edges.append((virtual_node_id, dst))
+            virtual_costs.append(0)
+            virtual_costs.append(0)
+            virtual_n_prizes.append(prize_e - cost_e)
+
+    prizes = np.concatenate([n_prizes, np.array(virtual_n_prizes)])
+    num_edges = len(edges)
+    if len(virtual_costs) > 0:
+        costs = np.array(costs + virtual_costs)
+        edges = np.array(edges + virtual_edges)
+
+    vertices, edges = pcst_fast(edges, prizes, costs, root, num_clusters,
+                                pruning, verbosity_level)
+
+    selected_nodes = vertices[vertices < data.num_nodes]
+    selected_edges = [mapping_e[e] for e in edges if e < num_edges]
+    virtual_vertices = vertices[vertices >= data.num_nodes]
+    if len(virtual_vertices) > 0:
+        virtual_vertices = vertices[vertices >= data.num_nodes]
+        virtual_edges = [mapping_n[i] for i in virtual_vertices]
+        selected_edges = np.array(selected_edges + virtual_edges)
+
+    edge_index = data.edge_index[:, selected_edges]
+    selected_nodes = np.unique(
+        np.concatenate(
+            [selected_nodes, edge_index[0].numpy(), edge_index[1].numpy()]))
+
+    n = textual_nodes.iloc[selected_nodes]
+    e = textual_edges.iloc[selected_edges]
+    desc = n.to_csv(index=False) + '\n' + e.to_csv(
+        index=False, columns=['src', 'edge_attr', 'dst'])
+
+    mapping = {n: i for i, n in enumerate(selected_nodes.tolist())}
+    src = [mapping[i] for i in edge_index[0].tolist()]
+    dst = [mapping[i] for i in edge_index[1].tolist()]
+
+    data = Data(
+        x=data.x[selected_nodes],
+        edge_index=torch.tensor([src, dst]),
+        edge_attr=data.edge_attr[selected_edges],
+    )
+
+    return data, desc
+
+
+class WebQSPDataset(InMemoryDataset):
+    r"""The WebQuestionsSP dataset of the `"The Value of Semantic Parse
+    Labeling for Knowledge Base Question Answering"
+    <https://aclanthology.org/P16-2033/>`_ paper.
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        split (str, optional): If :obj:`"train"`, loads the training dataset.
+            If :obj:`"val"`, loads the validation dataset.
+            If :obj:`"test"`, loads the test dataset. (default: :obj:`"train"`)
+        force_reload (bool, optional): Whether to re-process the dataset.
+            (default: :obj:`False`)
+    """
+    def __init__(
+        self,
+        root: str,
+        split: str = "train",
+        force_reload: bool = False,
+    ) -> None:
+        super().__init__(root, force_reload=force_reload)
+
+        if split not in {'train', 'val', 'test'}:
+            raise ValueError(f"Invalid 'split' argument (got {split})")
+
+        path = self.processed_paths[['train', 'val', 'test'].index(split)]
+        self.load(path)
+
+    @property
+    def processed_file_names(self) -> List[str]:
+        return ['train_data.pt', 'val_data.pt', 'test_data.pt']
+
+    def process(self) -> None:
+        import datasets
+        import pandas as pd
+
+        datasets = datasets.load_dataset('rmanluo/RoG-webqsp')
+
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        model_name = 'sentence-transformers/all-roberta-large-v1'
+        model = SentenceTransformer(model_name).to(device)
+        model.eval()
+
+        for dataset, path in zip(
+            [datasets['train'], datasets['validation'], datasets['test']],
+                self.processed_paths,
+        ):
+            questions = [example["question"] for example in dataset]
+            question_embs = model.encode(
+                questions,
+                batch_size=256,
+                output_device='cpu',
+            )
+
+            data_list = []
+            for i, example in enumerate(tqdm(dataset)):
+                raw_nodes: Dict[str, int] = {}
+                raw_edges = []
+                for tri in example["graph"]:
+                    h, r, t = tri
+                    h = h.lower()
+                    t = t.lower()
+                    if h not in raw_nodes:
+                        raw_nodes[h] = len(raw_nodes)
+                    if t not in raw_nodes:
+                        raw_nodes[t] = len(raw_nodes)
+                    raw_edges.append({
+                        "src": raw_nodes[h],
+                        "edge_attr": r,
+                        "dst": raw_nodes[t]
+                    })
+                nodes = pd.DataFrame([{
+                    "node_id": v,
+                    "node_attr": k,
+                } for k, v in raw_nodes.items()])
+                edges = pd.DataFrame(raw_edges)
+
+                nodes.node_attr = nodes.node_attr.fillna("")
+                x = model.encode(
+                    nodes.node_attr.tolist(),
+                    batch_size=256,
+                    output_device='cpu',
+                )
+                edge_attr = model.encode(
+                    edges.edge_attr.tolist(),
+                    batch_size=256,
+                    output_device='cpu',
+                )
+                edge_index = torch.tensor([
+                    edges.src.tolist(),
+                    edges.dst.tolist(),
+                ])
+
+                question = f"Question: {example['question']}\nAnswer: "
+                label = ('|').join(example['answer']).lower()
+                data = Data(
+                    x=x,
+                    edge_index=edge_index,
+                    edge_attr=edge_attr,
+                )
+                data, desc = retrieval_via_pcst(
+                    data,
+                    question_embs[i],
+                    nodes,
+                    edges,
+                    topk=3,
+                    topk_e=5,
+                    cost_e=0.5,
+                )
+                data.question = question
+                data.label = label
+                data.desc = desc
+                data_list.append(data)
+
+            self.save(data_list, path)

torch_geometric/edge_index.py

@@ -173,7 +173,7 @@ class EdgeIndex(Tensor):
     :meth:`EdgeIndex.fill_cache_`, and are maintained and adjusted over its
     lifespan (*e.g.*, when calling :meth:`EdgeIndex.flip`).
 
-    This representation ensures for optimal computation in GNN message passing
+    This representation ensures optimal computation in GNN message passing
     schemes, while preserving the ease-of-use of regular COO-based :pyg:`PyG`
     workflows.
 

torch_geometric/index.py

@@ -106,7 +106,7 @@ class Index(Tensor):
     :meth:`Index.fill_cache_`, and are maintaned and adjusted over its
     lifespan.
 
-    This representation ensures for optimal computation in GNN message passing
+    This representation ensures optimal computation in GNN message passing
     schemes, while preserving the ease-of-use of regular COO-based :pyg:`PyG`
     workflows.
 
@@ -120,7 +120,7 @@ class Index(Tensor):
         assert index.is_sorted
 
         # Flipping order:
-        edge_index.flip(0)
+        index.flip(0)
         >>> Index([[2, 1, 1, 0], dim_size=3)
         assert not index.is_sorted
 

torch_geometric/nn/models/g_retriever.py

@@ -3,7 +3,6 @@ from typing import List, Optional
 import torch
 from torch import Tensor
 
-from torch_geometric.nn.models import GAT
 from torch_geometric.nn.nlp.llm import BOS, LLM, MAX_NEW_TOKENS
 from torch_geometric.utils import scatter
 
@@ -43,7 +42,6 @@ class GRetriever(torch.nn.Module):
         llm: LLM,
         gnn: torch.nn.Module,
         use_lora: bool = False,
-        gnn_to_use=GAT,
         mlp_out_channels: int = 4096,
     ) -> None:
         super().__init__()
@@ -126,7 +124,15 @@ class GRetriever(torch.nn.Module):
         """
         x = self.encode(x, edge_index, batch, edge_attr)
         x = self.projector(x)
-        xs = x.split(x.size(0), dim=0)
+        xs = x.split(1, dim=0)
+
+        # Handle questions without node features:
+        batch_unique = batch.unique()
+        batch_size = len(question)
+        if len(batch_unique) < batch_size:
+            xs = [
+                xs[i] if i in batch_unique else None for i in range(batch_size)
+            ]
 
         (
             inputs_embeds,
@@ -174,7 +180,15 @@ class GRetriever(torch.nn.Module):
         """
         x = self.encode(x, edge_index, batch, edge_attr)
         x = self.projector(x)
-        xs = x.split(x.size(0), dim=0)
+        xs = x.split(1, dim=0)
+
+        # Handle questions without node features:
+        batch_unique = batch.unique()
+        batch_size = len(question)
+        if len(batch_unique) < batch_size:
+            xs = [
+                xs[i] if i in batch_unique else None for i in range(batch_size)
+            ]
 
         inputs_embeds, attention_mask, _ = self.llm._get_embeds(
             question, additional_text_context, xs)

torch_geometric/nn/nlp/llm.py

@@ -1,3 +1,4 @@
+import warnings
 from contextlib import nullcontext
 from typing import Any, Dict, List, Optional
 
@@ -85,6 +86,7 @@ class LLM(torch.nn.Module):
         self.word_embedding = self.llm.model.get_input_embeddings()
 
         if 'max_memory' not in kwargs:  # Pure CPU:
+            warnings.warn("LLM is being used on CPU, which may be slow")
             self.device = torch.device('cpu')
             self.autocast_context = nullcontext()
         else:

torch_geometric/nn/nlp/sentence_transformer.py

@@ -54,8 +54,11 @@ class SentenceTransformer(torch.nn.Module):
         self,
         text: List[str],
         batch_size: Optional[int] = None,
-        output_device: Optional[torch.device] = None,
+        output_device: Optional[Union[torch.device, str]] = None,
     ) -> Tensor:
+        is_empty = len(text) == 0
+        text = ['dummy'] if is_empty else text
+
         batch_size = len(text) if batch_size is None else batch_size
 
         embs: List[Tensor] = []
@@ -70,11 +73,13 @@ class SentenceTransformer(torch.nn.Module):
             emb = self(
                 input_ids=token.input_ids.to(self.device),
                 attention_mask=token.attention_mask.to(self.device),
-            ).to(output_device or 'cpu')
+            ).to(output_device)
 
             embs.append(emb)
 
-        return torch.cat(embs, dim=0) if len(embs) > 1 else embs[0]
+        out = torch.cat(embs, dim=0) if len(embs) > 1 else embs[0]
+        out = out[:0] if is_empty else out
+        return out
 
     def __repr__(self) -> str:
         return f'{self.__class__.__name__}(model_name={self.model_name})'