Update graphrag_agent.py

graphrag_agent.py

@@ -45,23 +45,20 @@ HF_TOKEN = os.getenv("HF_TOKEN")
 Entrez.email = ENTREZ_EMAIL
 MAX_TOKENS = 128000
 encoding = tiktoken.get_encoding("cl100k_base")
-# tokenizer = AutoTokenizer.from_pretrained("cambridgeltl/SapBERT-from-PubMedBERT-fulltext",token=HF_TOKEN)
-# model = AutoModel.from_pretrained("cambridgeltl/SapBERT-from-PubMedBERT-fulltext",token=HF_TOKEN).to(DEVICE)
-# model.eval()
+tokenizer = AutoTokenizer.from_pretrained("cambridgeltl/SapBERT-from-PubMedBERT-fulltext",token=HF_TOKEN)
+model = AutoModel.from_pretrained("cambridgeltl/SapBERT-from-PubMedBERT-fulltext",token=HF_TOKEN).to(DEVICE)
+model.eval()
         
-# bi_tokenizer = AutoTokenizer.from_pretrained("BAAI/bge-m3",token=HF_TOKEN)
-# bi_model = AutoModel.from_pretrained("BAAI/bge-m3",token=HF_TOKEN)
-# bi_model.eval()
+bi_tokenizer = AutoTokenizer.from_pretrained("BAAI/bge-m3",token=HF_TOKEN)
+bi_model = AutoModel.from_pretrained("BAAI/bge-m3",token=HF_TOKEN)
+bi_model.eval()
         
-# cross_tokenizer = AutoTokenizer.from_pretrained("BAAI/bge-reranker-v2-m3",token=HF_TOKEN)
-# cross_model = AutoModelForSequenceClassification.from_pretrained("BAAI/bge-reranker-v2-m3",token=HF_TOKEN)
-# cross_model.eval()
-sapbert_client = InferenceClient(
-    provider="hf-inference",
-    api_key=HF_TOKEN,
-)
-bge_client = InferenceClient(provider="hf-inference", api_key=HF_TOKEN)
-cross_client = InferenceClient(provider="hf-inference", api_key=HF_TOKEN)
+cross_tokenizer = AutoTokenizer.from_pretrained("BAAI/bge-reranker-v2-m3",token=HF_TOKEN)
+cross_model = AutoModelForSequenceClassification.from_pretrained("BAAI/bge-reranker-v2-m3",token=HF_TOKEN)
+cross_model.eval()
+#sapbert_client = InferenceClient(provider="hf-inference",api_key=HF_TOKEN)
+#bge_client = InferenceClient(provider="hf-inference", api_key=HF_TOKEN)
+#cross_client = InferenceClient(provider="hf-inference", api_key=HF_TOKEN)
 # ======================== 全局变量 ========================
 faiss_indices = {}
 metadata = {}
@@ -164,7 +161,7 @@ def _extract_json_from_text(text: str) -> Dict[str, Any]:
             return {}
     return {}
 
-# def embed_entity(entity_text: str):
+def embed_entity(entity_text: str):
     if not tokenizer or not model:
         raise ValueError("embedding model not loaded")
     with torch.no_grad():
@@ -176,21 +173,15 @@ def _extract_json_from_text(text: str) -> Dict[str, Any]:
         embedding = outputs.last_hidden_state.mean(dim=1).squeeze().cpu().numpy()
     return embedding
 
-def embed_entity(entity_text: str):
-    """
-    使用 Hugging Face Inference API 获取 SapBERT 嵌入
-    """
-    try:
-        result = sapbert_client.feature_extraction(
-            entity_text,
-            model="cambridgeltl/SapBERT-from-PubMedBERT-fulltext"
-        )
+#def embed_entity(entity_text: str):
+    # """使用 Hugging Face Inference API 获取 SapBERT 嵌入"""
+    # try:result = sapbert_client.feature_extraction(entity_text,model="cambridgeltl/SapBERT-from-PubMedBERT-fulltext")
         # 返回结果通常是 list[list[float]]，取平均或第一 token
-        embedding = [sum(x)/len(x) for x in zip(*result)]  # 对每个维度求平均
-        return embedding
-    except Exception as e:
-        print(f"Embedding error: {e}")
-        return None
+        # embedding = [sum(x)/len(x) for x in zip(*result)]  # 对每个维度求平均
+        # return embedding
+    # except Exception as e:
+        # print(f"Embedding error: {e}")
+        # return None
 
 def search_pubmed(pubmed_query: str, max_results: int = 3) -> str:
     try:
@@ -702,7 +693,7 @@ def whether_to_interact(state):
         return "user_input"
     elif interaction == "sufficient":
         print("决策: 信息充分，进入Neo4j检索。")
-        return "neo4j_retrieval"
+        return "kg_retrieval"
     else:
         return "stop_flow"
 
@@ -722,51 +713,7 @@ with ZipFile(zip_path, "r") as zip_ref:
     zip_ref.extractall("data/")
 
 print("✅ 已成功下载并解压 data.zip")
-def rerank_paths_cloud(query_text, path_kv):
-    try:
-        # 1. query embedding
-        query_emb = bge_client.feature_extraction(query_text, model="BAAI/bge-m3")
-        query_emb = torch.tensor(query_emb[0]).unsqueeze(0)
-        query_emb = F.normalize(query_emb, dim=-1)
-
-        # 2. path embeddings
-        path_keys = list(path_kv.keys())
-        all_cand_embs = []
-        for pk in path_keys:
-            cand_emb = bge_client.feature_extraction(pk, model="BAAI/bge-m3")
-            emb_tensor = torch.tensor(cand_emb[0]).unsqueeze(0)
-            emb_tensor = F.normalize(emb_tensor, dim=-1)
-            all_cand_embs.append(emb_tensor)
-
-        cand_embs = torch.cat(all_cand_embs, dim=0)
-        sim_scores = torch.matmul(query_emb, cand_embs.T).squeeze(0).tolist()
-
-        scored_paths = list(zip(path_keys, sim_scores))
-        scored_paths.sort(key=lambda x: x[1], reverse=True)
-        top100 = scored_paths[:100]
-
-        # 3. cross-encoder rerank
-        pairs = [(query_text, pk) for pk, _ in top100]
-        all_cross_scores = []
-        for q, pk in pairs:
-            input_pair = [(q, pk)]
-            scores = cross_client.text_classification(
-                input_pair,
-                model="BAAI/bge-reranker-v2-m3"
-            )
-            all_cross_scores.append(scores[0]["score"])
-
-        rerank_final = list(zip([p[1] for p in top100], all_cross_scores))
-        rerank_final.sort(key=lambda x: x[1], reverse=True)
-        top30 = rerank_final[:30]
-
-        top30_values = [path_kv[pk] for pk, _ in top30]
-        return {"neo4j_retrieval": top30_values}
 
-    except Exception as e:
-        print(f"rerank error: {e}")
-        fallback_values = list(path_kv.values())[:50]
-        return {"neo4j_retrieval": fallback_values}
 def neo4j_retrieval(state: MyState):
     logger.info("---NODE: neo4j_retrieval---")
     #user_query = [message.content for message in state["messages"] if hasattr(message, 'content')]
@@ -892,7 +839,51 @@ def neo4j_retrieval(state: MyState):
         except Exception as e:
             logger.warning(f"'{entity}'failed in faiss {e}")
             continue
-    return rerank_paths_cloud(query_text, path_kv)
+    try:
+        query_inputs = bi_tokenizer(query_text, return_tensors="pt", truncation=True, max_length=512,padding=True)
+        with torch.no_grad():
+            query_emb = bi_model(**query_inputs).last_hidden_state[:, 0]
+            query_emb = F.normalize(query_emb, dim=-1)
+
+        path_keys = list(path_kv.keys())
+        batch_size = 32
+        all_cand_embs = []
+        with torch.no_grad():
+            for i in range(0, len(path_keys), batch_size):
+                batch = path_keys[i:i + batch_size]
+                cand_inputs = bi_tokenizer(batch, return_tensors="pt", truncation=True, max_length=512,padding=True)
+                cand_embs_batch = bi_model(**cand_inputs).last_hidden_state[:, 0]
+                cand_embs_batch = F.normalize(cand_embs_batch, dim=-1)
+                all_cand_embs.append(cand_embs_batch)
+
+        cand_embs = torch.cat(all_cand_embs, dim=0)
+        sim_scores = torch.matmul(query_emb, cand_embs.T).squeeze(0).tolist()
+        scored_paths = list(zip(path_keys, sim_scores))
+        scored_paths.sort(key=lambda x: x[1], reverse=True)
+
+        top100 = scored_paths[:100]
+        pairs = [(query_text, pk) for pk, _ in top100]
+        all_cross_scores = []
+        cross_batch_size = 16
+        with torch.no_grad():
+            for i in range(0, len(pairs), cross_batch_size):
+                batch_pairs = pairs[i:i + cross_batch_size]
+                inputs = cross_tokenizer(batch_pairs, padding=True, truncation=True,  max_length=512,return_tensors="pt")
+                scores = cross_model(**inputs).logits.view(-1).tolist()
+                all_cross_scores.extend(scores)
+
+        rerank_final = list(zip([p[0] for p in top100], all_cross_scores))
+        rerank_final.sort(key=lambda x: x[1], reverse=True)
+        top30 = rerank_final[:30]
+
+        top30_values = [path_kv[pk] for pk, _ in top30]
+        logger.info(f"Cross-encoder reranked 30 path: {top30_values}")
+        return {"neo4j_retrieval": top30_values}
+
+    except Exception as e:
+        logger.warning(f"rerank error: {e}")
+        fallback_values = list(path_kv.values())[:50]
+        return {"neo4j_retrieval": fallback_values}
 
 
 def decide_router(state: MyState) -> dict:
@@ -1009,7 +1000,7 @@ def build_graphrag_agent():
     builder = StateGraph(MyState)
     builder.add_node("parse_query", parse_query)
     builder.add_node("user_input", user_input)
-    builder.add_node("neo4j_retrieval", neo4j_retrieval)
+    builder.add_node("kg_retrieval", neo4j_retrieval)
     builder.add_node("decide_router", decide_router)
     builder.add_node("api_search", api_search)
     builder.add_node("llm_answer", llm_answer)
@@ -1020,11 +1011,11 @@ def build_graphrag_agent():
             whether_to_interact,
             {
                 "user_input": "user_input",
-                "neo4j_retrieval": "neo4j_retrieval"
+                "kg_retrieval": "kg_retrieval"
             }
         )
     builder.add_edge("user_input", "parse_query")
-    builder.add_edge("neo4j_retrieval", "decide_router")
+    builder.add_edge("kg_retrieval", "decide_router")
     builder.add_conditional_edges(
             "decide_router",
             lambda state: state["route"],