Update graphrag_agent.py

graphrag_agent.py

@@ -42,9 +42,17 @@ HF_TOKEN = os.getenv("HF_TOKEN")
 Entrez.email = ENTREZ_EMAIL
 MAX_TOKENS = 128000
 encoding = tiktoken.get_encoding("cl100k_base")
-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)
+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()
+        
+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()
 # ======================== 全局变量 ========================
 faiss_indices = {}
 metadata = {}
@@ -147,6 +155,18 @@ def _extract_json_from_text(text: str) -> Dict[str, Any]:
             return {}
     return {}
 
+def embed_entity(entity_text: str):
+    if not tokenizer or not model:
+        raise ValueError("embedding model not loaded")
+    with torch.no_grad():
+        inputs = tokenizer(
+            entity_text, return_tensors="pt",
+            padding=True, truncation=True, max_length=64
+        ).to(DEVICE)
+        outputs = model(**inputs)
+        embedding = outputs.last_hidden_state.mean(dim=1).squeeze().cpu().numpy()
+    return embedding
+
 def embed_entity_cloud(entity_text: str):
     """
     使用 Hugging Face Inference API 获取 SapBERT 嵌入
@@ -163,7 +183,6 @@ def embed_entity_cloud(entity_text: str):
         print(f"Embedding error: {e}")
         return None
 
-
 def search_pubmed(pubmed_query: str, max_results: int = 3) -> str:
     try:
         handle = Entrez.esearch(db="pubmed", term=pubmed_query, retmax=max_results)
@@ -864,7 +883,51 @@ def neo4j_retrieval(state: MyState):
             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:
     print("---EDGE: decide_router---")