470 lines
21 KiB
Python
470 lines
21 KiB
Python
from neo4j import GraphDatabase
|
|
import faiss
|
|
import numpy as np
|
|
import random
|
|
from collections import defaultdict
|
|
from typing import List
|
|
import time
|
|
import logging
|
|
from atlas_rag.llm_generator.llm_generator import LLMGenerator
|
|
from atlas_rag.vectorstore.embedding_model import BaseEmbeddingModel
|
|
from atlas_rag.retriever.lkg_retriever.base import BaseLargeKGEdgeRetriever
|
|
import nltk
|
|
from nltk.corpus import stopwords
|
|
nltk.download('stopwords')
|
|
|
|
class LargeKGToGRetriever(BaseLargeKGEdgeRetriever):
|
|
def __init__(self, keyword: str, neo4j_driver: GraphDatabase,
|
|
llm_generator: LLMGenerator, sentence_encoder: BaseEmbeddingModel, filter_encoder: BaseEmbeddingModel,
|
|
node_index: faiss.Index,
|
|
topN : int = 5,
|
|
Dmax : int = 3,
|
|
Wmax : int = 3,
|
|
prune_size: int = 10,
|
|
logger: logging.Logger = None):
|
|
"""
|
|
Initialize the LargeKGToGRetriever for billion-level KG retrieval using Neo4j.
|
|
|
|
Args:
|
|
keyword (str): Identifier for the KG dataset (e.g., 'cc_en').
|
|
neo4j_driver (GraphDatabase): Neo4j driver for database access.
|
|
llm_generator (LLMGenerator): LLM for NER, rating, and reasoning.
|
|
sentence_encoder (BaseEmbeddingModel): Encoder for generating embeddings.
|
|
node_index (faiss.Index): FAISS index for node embeddings.
|
|
logger (Logger, optional): Logger for verbose output.
|
|
"""
|
|
self.keyword = keyword
|
|
self.neo4j_driver = neo4j_driver
|
|
self.llm_generator = llm_generator
|
|
self.sentence_encoder = sentence_encoder
|
|
self.filter_encoder = filter_encoder
|
|
self.node_faiss_index = node_index
|
|
self.verbose = logger is not None
|
|
self.logger = logger
|
|
self.topN = topN
|
|
self.Dmax = Dmax
|
|
self.Wmax = Wmax
|
|
self.prune_size = prune_size
|
|
|
|
def ner(self, text: str) -> List[str]:
|
|
"""
|
|
Extract named entities from the query text using the LLM.
|
|
|
|
Args:
|
|
text (str): The query text.
|
|
|
|
Returns:
|
|
List[str]: List of extracted entities.
|
|
"""
|
|
entities = self.llm_generator.large_kg_ner(text)
|
|
if self.verbose:
|
|
self.logger.info(f"Extracted entities: {entities}")
|
|
return entities
|
|
|
|
def retrieve_topk_nodes(self, query: str, top_k_nodes: int = 5) -> List[str]:
|
|
"""
|
|
Retrieve top-k nodes similar to entities in the query.
|
|
|
|
Args:
|
|
query (str): The user query.
|
|
top_k_nodes (int): Number of nodes to retrieve per entity.
|
|
|
|
Returns:
|
|
List[str]: List of node numeric_ids.
|
|
"""
|
|
start_time = time.time()
|
|
entities = self.ner(query)
|
|
if self.verbose:
|
|
ner_time = time.time() - start_time
|
|
self.logger.info(f"NER took {ner_time:.2f} seconds, entities: {entities}")
|
|
if not entities:
|
|
entities = [query]
|
|
|
|
initial_nodes = []
|
|
for entity in entities:
|
|
entity_embedding = self.sentence_encoder.encode([entity])
|
|
D, I = self.node_faiss_index.search(entity_embedding, 3)
|
|
if self.verbose:
|
|
self.logger.info(f"Entity: {entity}, FAISS Distances: {D}, Indices: {I}")
|
|
if len(I[0]) > 0: # Check if results exist
|
|
initial_nodes.extend([str(i) for i in I[0]])
|
|
# no need filtering as ToG pruning will handle it.
|
|
topk_nodes_ids = list(set(initial_nodes))
|
|
|
|
with self.neo4j_driver.session() as session:
|
|
start_time = time.time()
|
|
query = """
|
|
MATCH (n:Node)
|
|
WHERE n.numeric_id IN $topk_nodes_ids
|
|
RETURN n.numeric_id AS id, n.name AS name
|
|
"""
|
|
result = session.run(query, topk_nodes_ids=topk_nodes_ids)
|
|
topk_nodes_dict = {}
|
|
for record in result:
|
|
topk_nodes_dict[record["id"]] = record["name"]
|
|
if self.verbose:
|
|
neo4j_time = time.time() - start_time
|
|
self.logger.info(f"Neo4j query took {neo4j_time:.2f} seconds, count: {len(topk_nodes_ids)}")
|
|
|
|
if self.verbose:
|
|
self.logger.info(f"Top-k nodes: {topk_nodes_dict}")
|
|
return list(topk_nodes_dict.keys()), list(topk_nodes_dict.values()) # numeric_id of nodes returned
|
|
|
|
def expand_paths(self, P: List[List[str]], PIDS: List[List[str]], PTYPES: List[List[str]], width: int, query: str) -> List[List[str]]:
|
|
"""
|
|
Expand each path by adding neighbors of the last node.
|
|
|
|
Args:
|
|
P (List[List[str]]): Current list of paths, where each path is a list of alternating node_ids and relation types.
|
|
|
|
Returns:
|
|
List[List[str]]: List of expanded paths.
|
|
"""
|
|
last_nodes = []
|
|
last_node_ids = []
|
|
last_node_types = []
|
|
|
|
paths_end_with_text = []
|
|
paths_end_with_text_id = []
|
|
paths_end_with_text_type = []
|
|
|
|
paths_end_with_node = []
|
|
paths_end_with_node_id = []
|
|
paths_end_with_node_type = []
|
|
if self.verbose:
|
|
self.logger.info(f"Expanding paths, current paths: {P}")
|
|
for p, pid, ptype in zip(P, PIDS, PTYPES):
|
|
if not p or not pid or not ptype: # Skip empty paths
|
|
continue
|
|
t = ptype[-1]
|
|
if t == "Text":
|
|
paths_end_with_text.append(p)
|
|
paths_end_with_text_id.append(pid)
|
|
paths_end_with_text_type.append(ptype)
|
|
continue
|
|
last_node = p[-1] # Last node in the path
|
|
last_node_id = pid[-1] # Last node numeric_id
|
|
|
|
last_nodes.append(last_node)
|
|
last_node_ids.append(last_node_id)
|
|
last_node_types.append(t)
|
|
|
|
paths_end_with_node.append(p)
|
|
paths_end_with_node_id.append(pid)
|
|
paths_end_with_node_type.append(ptype)
|
|
|
|
assert len(last_nodes) == len(last_node_ids) == len(last_node_types), "Mismatch in last nodes, ids, and types lengths"
|
|
|
|
if not last_node_ids:
|
|
return paths_end_with_text, paths_end_with_text_id, paths_end_with_text_type
|
|
|
|
with self.neo4j_driver.session() as session:
|
|
# Query Node relationships
|
|
start_time = time.time()
|
|
outgoing_query = """
|
|
CALL apoc.cypher.runTimeboxed(
|
|
"MATCH (n:Node)-[r:Relation]-(m:Node) WHERE n.numeric_id IN $last_node_ids
|
|
WITH n, r, m ORDER BY rand() LIMIT 60000
|
|
RETURN n.numeric_id AS source, n.name AS source_name, r.relation AS rel_type, m.numeric_id AS target, m.name AS target_name, 'Node' AS target_type",
|
|
{last_node_ids: $last_node_ids},
|
|
60000
|
|
)
|
|
YIELD value
|
|
RETURN value.source AS source, value.source_name AS source_name, value.rel_type AS rel_type, value.target AS target, value.target_name AS target_name, value.target_type AS target_type
|
|
"""
|
|
outgoing_result = session.run(outgoing_query, last_node_ids=last_node_ids)
|
|
outgoing = [(record["source"], record['source_name'], record["rel_type"], record["target"], record["target_name"], record["target_type"])
|
|
for record in outgoing_result]
|
|
if self.verbose:
|
|
outgoing_time = time.time() - start_time
|
|
self.logger.info(f"Outgoing relationships query took {outgoing_time:.2f} seconds, count: {len(outgoing)}")
|
|
|
|
# # Query outgoing Node -> Text relationships
|
|
# start_time = time.time()
|
|
# outgoing_text_query = """
|
|
# MATCH (n:Node)-[r:Source]->(t:Text)
|
|
# WHERE n.numeric_id IN $last_node_ids
|
|
# RETURN n.numeric_id AS source, n.name AS source_name, 'from Source' AS rel_type, t.numeric_id as target, t.original_text AS target_name, 'Text' AS target_type
|
|
# """
|
|
# outgoing_text_result = session.run(outgoing_text_query, last_node_ids=last_node_ids)
|
|
# outgoing_text = [(record["source"], record["source_name"], record["rel_type"], record["target"], record["target_name"], record["target_type"])
|
|
# for record in outgoing_text_result]
|
|
# if self.verbose:
|
|
# outgoing_text_time = time.time() - start_time
|
|
# self.logger.info(f"Outgoing Node->Text relationships query took {outgoing_text_time:.2f} seconds, count: {len(outgoing_text)}")
|
|
|
|
last_node_to_new_paths = defaultdict(list)
|
|
last_node_to_new_paths_ids = defaultdict(list)
|
|
last_node_to_new_paths_types = defaultdict(list)
|
|
|
|
|
|
for p, pid, ptype in zip(P, PIDS, PTYPES):
|
|
last_node = p[-1]
|
|
last_node_id = pid[-1]
|
|
# Outgoing Node -> Node
|
|
for source, source_name, rel_type, target, target_name, target_type in outgoing:
|
|
if source == last_node_id and target_name not in p:
|
|
new_path = p + [rel_type, target_name]
|
|
if target_name.lower() in stopwords.words('english'):
|
|
continue
|
|
last_node_to_new_paths[last_node].append(new_path)
|
|
last_node_to_new_paths_ids[last_node].append(pid + [target])
|
|
last_node_to_new_paths_types[last_node].append(ptype + [target_type])
|
|
|
|
# # Outgoing Node -> Text
|
|
# for source, source_name, rel_type, target, target_name, target_type in outgoing_text:
|
|
# if source == last_node_id and target_name not in p:
|
|
# new_path = p + [rel_type, target_name]
|
|
|
|
# last_node_to_new_paths_text[last_node].append(new_path)
|
|
# last_node_to_new_paths_text_ids[last_node].append(pid + [target])
|
|
# last_node_to_new_paths_text_types[last_node].append(ptype + [target_type])
|
|
|
|
# # Incoming Node -> Node
|
|
# for source, rel_type, target, source_name, source_type in incoming:
|
|
# if target == last_node_id and source not in p:
|
|
# new_path = p + [rel_type, source_name]
|
|
|
|
|
|
num_paths = 0
|
|
for last_node, new_paths in last_node_to_new_paths.items():
|
|
num_paths += len(new_paths)
|
|
# for last_node, new_paths in last_node_to_new_paths_text.items():
|
|
# num_paths += len(new_paths)
|
|
new_paths = []
|
|
new_pids = []
|
|
new_ptypes = []
|
|
if self.verbose:
|
|
self.logger.info(f"Number of new paths before filtering: {num_paths}")
|
|
self.logger.info(f"last nodes: {last_node_to_new_paths.keys()}")
|
|
if num_paths > len(last_node_ids) * width:
|
|
# Apply filtering when total paths exceed threshold
|
|
for last_node, new_ps in last_node_to_new_paths.items():
|
|
if len(new_ps) > width:
|
|
path_embeddings = self.filter_encoder.encode(new_ps)
|
|
query_embeddings = self.filter_encoder.encode([query])
|
|
scores = np.dot(path_embeddings, query_embeddings.T).flatten()
|
|
top_indices = np.argsort(scores)[-width:]
|
|
new_paths.extend([new_ps[i] for i in top_indices])
|
|
new_pids.extend([last_node_to_new_paths_ids[last_node][i] for i in top_indices])
|
|
new_ptypes.extend([last_node_to_new_paths_types[last_node][i] for i in top_indices])
|
|
else:
|
|
new_paths.extend(new_ps)
|
|
new_pids.extend(last_node_to_new_paths_ids[last_node])
|
|
new_ptypes.extend(last_node_to_new_paths_types[last_node])
|
|
else:
|
|
# Collect all paths without filtering when total is at or below threshold
|
|
for last_node, new_ps in last_node_to_new_paths.items():
|
|
new_paths.extend(new_ps)
|
|
new_pids.extend(last_node_to_new_paths_ids[last_node])
|
|
new_ptypes.extend(last_node_to_new_paths_types[last_node])
|
|
|
|
if self.verbose:
|
|
self.logger.info(f"Expanded paths count: {len(new_paths)}")
|
|
self.logger.info(f"Expanded paths: {new_paths}")
|
|
return new_paths, new_pids, new_ptypes
|
|
|
|
def path_to_string(self, path: List[str]) -> str:
|
|
"""
|
|
Convert a path to a human-readable string for LLM rating.
|
|
|
|
Args:
|
|
path (List[str]): Path as a list of node_ids and relation types.
|
|
|
|
Returns:
|
|
str: String representation of the path.
|
|
"""
|
|
if len(path) < 1:
|
|
return ""
|
|
path_str = []
|
|
with self.neo4j_driver.session() as session:
|
|
for i in range(0, len(path), 2):
|
|
node_id = path[i]
|
|
result = session.run("MATCH (n:Node {numeric_id: $node_id}) RETURN n.name", node_id=node_id)
|
|
node_name = result.single()["n.name"] if result.single() else node_id
|
|
if i + 1 < len(path):
|
|
rel_type = path[i + 1]
|
|
path_str.append(f"{node_name} ---> {rel_type} --->")
|
|
else:
|
|
path_str.append(node_name)
|
|
return " ".join(path_str).strip()
|
|
|
|
def prune(self, query: str, P: List[List[str]], PIDS: List[List[str]], PTYPES: List[List[str]], topN: int = 5) -> List[List[str]]:
|
|
"""
|
|
Prune paths to keep the top N based on LLM relevance ratings.
|
|
|
|
Args:
|
|
query (str): The user query.
|
|
P (List[List[str]]): List of paths to prune.
|
|
topN (int): Number of paths to retain.
|
|
|
|
Returns:
|
|
List[List[str]]: Top N paths.
|
|
"""
|
|
ratings = []
|
|
path_strings = P
|
|
|
|
# Process paths in chunks of 10
|
|
for i in range(0, len(path_strings), self.prune_size):
|
|
chunk = path_strings[i:i + self.prune_size]
|
|
|
|
# Construct user prompt with the current chunk of paths listed
|
|
user_prompt = f"Please rate the following paths based on how well they help answer the query (1-5, 0 if not relevant).\n\nQuery: {query}\n\nPaths:\n"
|
|
for j, path_str in enumerate(chunk, 1):
|
|
user_prompt += f"{j + i}. {path_str}\n"
|
|
user_prompt += "\nProvide a list of integers, each corresponding to the rating of the path's ability to help answer the query."
|
|
|
|
# Define system prompt to expect a list of integers
|
|
system_prompt = "You are a rating machine that only provides a list of comma-separated integers (0-5) as a response, each rating how well the corresponding path helps answer the query."
|
|
|
|
# Send the prompt to the language model
|
|
messages = [
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": user_prompt}
|
|
]
|
|
|
|
response = self.llm_generator.generate_response(messages, max_new_tokens=1024, temperature=0.0)
|
|
if self.verbose:
|
|
self.logger.info(f"LLM response for chunk {i // self.prune_size + 1}: {response}")
|
|
|
|
# Parse the response into a list of ratings
|
|
rating_str = response.strip()
|
|
chunk_ratings = [int(r) for r in rating_str.split(',') if r.strip().isdigit()]
|
|
if len(chunk_ratings) > len(chunk):
|
|
chunk_ratings = chunk_ratings[:len(chunk)]
|
|
if self.verbose:
|
|
self.logger.warning(f"Received more ratings ({len(chunk_ratings)}) than paths in chunk ({len(chunk)}). Trimming ratings.")
|
|
ratings.extend(chunk_ratings) # Concatenate ratings
|
|
|
|
# Ensure ratings length matches number of paths, padding with 0s if necessary
|
|
if len(ratings) < len(path_strings):
|
|
# self.logger.warning(f"Number of ratings ({len(ratings)}) does not match number of paths ({len(path_strings)}). Padding with 0s.")
|
|
# ratings += [0] * (len(path_strings) - len(ratings))
|
|
# fall back to use filter encoder to get topN
|
|
self.logger.warning(f"Number of ratings ({len(ratings)}) does not match number of paths ({len(path_strings)}). Using filter encoder to get topN paths.")
|
|
path_embeddings = self.filter_encoder.encode(path_strings)
|
|
query_embedding = self.filter_encoder.encode([query])
|
|
scores = np.dot(path_embeddings, query_embedding.T).flatten()
|
|
top_indices = np.argsort(scores)[-topN:]
|
|
top_paths = [path_strings[i] for i in top_indices]
|
|
return top_paths, [PIDS[i] for i in top_indices], [PTYPES[i] for i in top_indices]
|
|
elif len(ratings) > len(path_strings):
|
|
self.logger.warning(f"Number of ratings ({len(ratings)}) exceeds number of paths ({len(path_strings)}). Trimming ratings.")
|
|
ratings = ratings[:len(path_strings)]
|
|
|
|
# Sort indices based on ratings in descending order
|
|
sorted_indices = sorted(range(len(ratings)), key=lambda i: ratings[i], reverse=True)
|
|
|
|
# Filter out indices where the rating is 0
|
|
filtered_indices = [i for i in sorted_indices if ratings[i] > 0]
|
|
|
|
# Take the top N indices from the filtered list
|
|
top_indices = filtered_indices[:topN]
|
|
|
|
# Use the filtered indices to get the top paths, PIDS, and PTYPES
|
|
if self.verbose:
|
|
self.logger.info(f"Top indices after pruning: {top_indices}")
|
|
self.logger.info(f"length of path_strings: {len(path_strings)}")
|
|
top_paths = [path_strings[i] for i in top_indices]
|
|
top_pids = [PIDS[i] for i in top_indices]
|
|
top_ptypes = [PTYPES[i] for i in top_indices]
|
|
|
|
|
|
# Log top paths if verbose mode is enabled
|
|
if self.verbose:
|
|
self.logger.info(f"Pruned to top {topN} paths: {top_paths}")
|
|
|
|
return top_paths, top_pids, top_ptypes
|
|
|
|
def reasoning(self, query: str, P: List[List[str]]) -> bool:
|
|
"""
|
|
Check if the current paths are sufficient to answer the query.
|
|
|
|
Args:
|
|
query (str): The user query.
|
|
P (List[List[str]]): Current list of paths.
|
|
|
|
Returns:
|
|
bool: True if sufficient, False otherwise.
|
|
"""
|
|
triples = []
|
|
with self.neo4j_driver.session() as session:
|
|
for path in P:
|
|
if len(path) < 3:
|
|
continue
|
|
for i in range(0, len(path) - 2, 2):
|
|
node1_name = path[i]
|
|
rel = path[i + 1]
|
|
node2_name = path[i + 2]
|
|
triples.append(f"({node1_name}, {rel}, {node2_name})")
|
|
triples_str = ". ".join(triples)
|
|
prompt = f"Are these triples, along with your knowledge, sufficient to answer the query?\nQuery: {query}\nTriples: {triples_str}"
|
|
messages = [
|
|
{"role": "system", "content": "Answer Yes or No only."},
|
|
{"role": "user", "content": prompt}
|
|
]
|
|
response = self.llm_generator.generate_response(messages,max_new_tokens=512)
|
|
if self.verbose:
|
|
self.logger.info(f"Reasoning result: {response}")
|
|
return "yes" in response.lower()
|
|
|
|
def retrieve_passages(self, query: str) -> List[str]:
|
|
"""
|
|
Retrieve the top N paths to answer the query.
|
|
|
|
Args:
|
|
query (str): The user query.
|
|
topN (int): Number of paths to return.
|
|
Dmax (int): Maximum depth of path expansion.
|
|
Wmax (int): Maximum width of path expansion.
|
|
|
|
Returns:
|
|
List[str]: List of triples as strings.
|
|
"""
|
|
topN = self.topN
|
|
Dmax = self.Dmax
|
|
Wmax = self.Wmax
|
|
if self.verbose:
|
|
self.logger.info(f"Retrieving paths for query: {query}")
|
|
|
|
initial_nodes_ids, initial_nodes = self.retrieve_topk_nodes(query, top_k_nodes=topN)
|
|
if not initial_nodes:
|
|
if self.verbose:
|
|
self.logger.info("No initial nodes found.")
|
|
return []
|
|
|
|
P = [[node] for node in initial_nodes]
|
|
PIDS = [[node_id] for node_id in initial_nodes_ids]
|
|
PTYPES = [["Node"] for _ in initial_nodes_ids] # Assuming all initial nodes are of type 'Node'
|
|
for D in range(Dmax + 1):
|
|
if self.verbose:
|
|
self.logger.info(f"Depth {D}, Current paths: {len(P)}")
|
|
P, PIDS, PTYPES = self.expand_paths(P, PIDS, PTYPES, Wmax, query)
|
|
if not P:
|
|
if self.verbose:
|
|
self.logger.info("No paths to expand.")
|
|
break
|
|
P, PIDS, PTYPES = self.prune(query, P, PIDS, PTYPES, topN)
|
|
if D == Dmax:
|
|
if self.verbose:
|
|
self.logger.info(f"Reached maximum depth {Dmax}, stopping expansion.")
|
|
break
|
|
if self.reasoning(query, P):
|
|
if self.verbose:
|
|
self.logger.info("Paths sufficient, stopping expansion.")
|
|
break
|
|
|
|
# Extract final triples
|
|
triples = []
|
|
with self.neo4j_driver.session() as session:
|
|
for path in P:
|
|
for i in range(0, len(path) - 2, 2):
|
|
node1_name = path[i]
|
|
rel = path[i + 1]
|
|
node2_name = path[i + 2]
|
|
triples.append(f"({node1_name}, {rel}, {node2_name})")
|
|
|
|
if self.verbose:
|
|
self.logger.info(f"Final triples: {triples}")
|
|
return triples, 'N/A' # 'N/A' for passages_score as this retriever does not return passages
|