torch_geometric.llm.utils.KNNRAGFeatureStore

class KNNRAGFeatureStore

Bases: LocalFeatureStore

A feature store that uses a KNN-based retrieval.

property config: Dict[str, Any]

Get the config for the feature store.

Return type:

Dict[str, Any]

property x: Tensor

Returns the node features.

Return type:

Tensor

property edge_attr: Tensor

Returns the edge attributes.

Return type:

Tensor

retrieve_seed_nodes(query: Union[str, List[str], Tuple[str]]) → Tuple[Union[Tensor, None, str, Tuple[str, Optional[Tensor]]], Tensor]

Retrieves the k_nodes most similar nodes to the given query.

Parameters:

query (Union[str, List[str], Tuple[str]]) – The query or list of queries to search for.

Returns:

Tuple[Union[Tensor, None, str, Tuple[str, Optional[Tensor]]], Tensor] – The indices of the most similar nodes and the encoded query

load_subgraph(sample: Union[SamplerOutput, HeteroSamplerOutput], induced: bool = True) → Union[Data, HeteroData]

Loads a subgraph from the given sample.

Parameters:

• sample (Union[SamplerOutput, HeteroSamplerOutput]) – The sample to load the subgraph from.

• induced (bool, default: True) – Whether to return the induced subgraph. Resets node and edge ids.

Returns:

Union[Data, HeteroData] – The loaded subgraph.