知識圖譜增強的KG-RAG框架

昨天我們聊到KG在RAG中如何發揮作用，今天我們來看一個具體的例子。 我們找到一篇論文： https://arxiv.org/abs/2311.17330 ，論文的研究人員開發了一種名爲知識圖譜增強的提示生成（KG-RAG）框架（https://github.com/BaranziniLab/KG_RAG），該框架利用生物醫學知識圖譜SPOKE與大型語言模型相結合，有效的提升了LLM在醫療領域的問答效果。

KG-RAG框架介紹

KG-RAG框架，較好的結合了生物醫學知識圖譜SPOKE和LLM的優勢。SPOKE是一個開放知識圖譜，提供數據下載和開放API，整合了超過40個公開可用的生物醫學知識源，涵蓋了基因、蛋白質、藥物、化合物、疾病等概念和概念之間的關係，可以爲LLM提供一個強大的醫療領域知識。

研究人員對KG-RAG框架進行了廣泛的測試，包括單跳和雙跳提示、藥物再利用查詢、生物醫學真假問題和多項選擇題。結果表明，KG-RAG顯著提高了LLMs的性能，特別是在具有挑戰性的多項選擇題數據集上，LLMs都取得了較大的提升。此外，KG-RAG還能夠提供有意義的藥物再利用建議，並在回答中體現出對臨牀試驗必要性的謹慎態度。

相關測試結果：

工作原理

KG-RAG框架的工作原理包括以下步驟：

1. 實體識別：從用戶輸入的query中識別出疾病實體，然後在SPOKE知識圖譜中找到相應的節點。
2. 上下文提取：從SPOKE知識圖譜中提取與疾病節點相關的上下文信息，並將其轉換爲自然語言。
3. 提示組裝：將提取的上下文與原始prompt結合。
4. 文本生成：使用LLM（如Llama-2-13b、GPT-3.5-Turbo或GPT-4）生成有意義的生物醫學文本。

實體識別

區別於用小模型去做NER，KG-RAG裏使用LLM識別實體。

1. 實體抽取（Disease Entity Extraction）

在KG-RAG框架中，這一過程是通過零樣本提示（zero-shot prompting）實現的。研究人員設計了一個高效的抽取prompt，引導大型語言模型（如GPT-3.5-Turbo）從輸入文本中提取疾病實體，並將結果以JSON格式返回。

def disease_entity_extractor_v2(text):  
    chat_model_id, chat_deployment_id = get_gpt35()  
    prompt_updated = system_prompts["DISEASE_ENTITY_EXTRACTION"] + "\n" + "Sentence : " + text  
    resp = get_GPT_response(prompt_updated, system_prompts["DISEASE_ENTITY_EXTRACTION"], chat_model_id, chat_deployment_id, temperature=0)  
    try:  
        entity_dict = json.loads(resp)  
        return entity_dict["Diseases"]  
    except:  
        return None

這裏的DISEASE_ENTITY_EXTRACTION:

  You are an expert disease entity extractor from a sentence and report it as JSON in the following format:  
  Diseases: <List of extracted entities>  
  Please report only Diseases. Do not report any other entities like Genes, Proteins, Enzymes etc.

2. 實體鏈接（Entity Matching to SPOKE）

疾病實體抽取出來後，下一步就是將這些實體與SPOKE知識圖譜中的疾病實體進行匹配，也就是傳統NLP任務中的實體鏈接，KG-RAG這個框架中採用的方法是，用語義相似度的方式來做。

• 實體embedding計算：首先，使用Embedding模型（如'all-MiniLM-L6-v2'）爲SPOKE知識圖譜中的所有疾病概念節點計算embedding向量
• 將計算出的疾病embedding存儲在向量數據庫（如'Chroma'）中，以便快速檢索。
• 語義搜索匹配：將LLM提取的疾病實體與向量數據庫中的疾病實體進行比較，選擇最相似的

當然，如果零樣本方法未能識別出疾病實體，採取的辦法是直接拿原始query去匹配，取top 5。

最終，實體匹配過程會輸出與輸入文本提示中的疾病實體最相關的SPOKE知識圖譜節點。這些節點及其相關信息將用於後續的上下文提取和文本生成步驟。通過這種方法，KG-RAG框架能夠有效地從專業文本中提取和識別疾病實體，並將其與豐富的生物醫學知識庫相連接，從而生成準確、可靠的生物醫學相關信息。

子圖查詢與剪枝

子圖查詢

在得到具體的實體後，緊接着就是從KG中去查詢這個實體關聯的子圖，這些信息通常以三元組（Subject, Predicate, Object）的形式存在，表示不同的生物醫學關係。通常情況下，可以查詢1~3跳內的三元組信息，這裏藉助圖數據庫可以比較容易的實現。

得到的三元組信息，LLM可能不太能比較好的理解，這裏就需要將三元組轉換成自然語言，以便與輸入提示結合並用於後續的文本生成。舉個例子：

(Disease hypertension, ASSOCIATES_DaG, Gene VHL) → `Disease hypertension associates Gene VHL`

上下文剪枝

在KG-RAG框架中，Context Pruning（上下文剪枝）是一個關鍵步驟，就和dfs遍歷時，需要剪枝來減少遍歷時間一樣，這裏的剪枝可以減少給LLM的信息，減少token數量，同時過濾掉一些無用信心，還能提升LLM回答的精確性。

Context Pruning的具體做法還是會基於embedding來計算語義相似度，大概就是使用embedding模型計算三元組和query的cos相似度，最後選擇策略：

• 條件一：上下文關聯的餘弦相似度必須大於所有提取上下文關聯的相似度分佈的75%分位
• 條件二：餘弦相似度的最小值必須達到0.5

通過這個0.5 和 75%，可以有效減少給LLM的無效信息，有助於提高後續文本生成的準確性和相關性。

提示組裝與文本生成

這裏就簡單了，就是和question一起，組合爲propmt，再加上SYSTEM_PROMPT，送給LLM回答：

question = row["text"]  
#檢索
context = retrieve_context(question, vectorstore, embedding_function_for_context_retrieval, node_context_df, context_volume, QUESTION_VS_CONTEXT_SIMILARITY_PERCENTILE_THRESHOLD, QUESTION_VS_CONTEXT_MINIMUM_SIMILARITY, edge_evidence)  
# 
enriched_prompt = "Context: "+ context + "\n" + "Question: " + question  
output = get_GPT_response(enriched_prompt, SYSTEM_PROMPT, CHAT_MODEL_ID, CHAT_DEPLOYMENT_ID, temperature=TEMPERATURE)  
if not output:
enriched_prompt = "Context: "+ context + "\n" + "Question: "+ question

這裏的SYSTEM_PROMPT：

# One-Hop Validation  
SINGLE_DISEASE_ENTITY_VALIDATION: |  
  You are an expert biomedical researcher. For answering the Question at the end, you need to first read the Context provided.  
  Then give your final answer by considering the context and your inherent knowledge on the topic. Give your answer in the following JSON format:  
    {Compounds: <list of compounds>, Diseases: <list of diseases>}  
  
# Two-Hop Validation  
TWO_DISEASE_ENTITY_VALIDATION: |  
  You are an expert biomedical researcher. For answering the Question at the end, you need to first read the Context provided.  
  Then give your final answer by considering the context and your inherent knowledge on the topic. Give your answer in the following JSON format:  
    {Nodes: <list of nodes>}

運行的結果舉例：

question = 'Does drug dependence have any genetic factors? Do you have any statistical evidence from trustworthy sources for this?'
KG_RAG_FLAG = True  
EDGE_EVIDENCE_FLAG = True   
generate_response(question, LLM_TO_USE, KG_RAG_FLAG, evidence_flag=EDGE_EVIDENCE_FLAG, temperature=TEMPERATURE)

Yes, drug dependence does have genetic factors. The genes KAT2B and SLC25A16 have been associated with drug dependence. This information is backed by statistical evidence from the GWAS Catalog, with p-values of 4e-10 and 1e-09 respectively.

KG-RAG 在應用中落地思考

KG-RAG 給出瞭如何結合KG來做RAG的一個有效方案，但這裏再工業場景中落地，還有很多是我們細緻去思考的。比如NER實體識別這裏，通過LLM來抽取，再來做entity link，這裏的效率肯定是感人的，其實這裏傳統的bert模型就可以了，成本可以忽略不計。

再則，剪枝這裏，原始的實現效率是很低的，這裏的embedding模型也需要專門去微調訓練。三元組轉換成自然語言，這裏也是有講究，如何生成更通順的自然語言，更好的做法LLM+人工，確定好模版，通過模版生成。另外，是先是被實體，然後去查詢實體的關聯子圖，還是全圖查詢，通過實體來過濾，都是可以考慮的點。

總結

KG-RAG框架通過結合生物醫學知識圖譜和LLM，爲生物醫學領域的問題提供了通用的解決方案。不僅提高了模型的性能，而且簡化了流程，使其更具成本效益和時間效率。

在其他領域如何去應用KG做RAG，一方面可以擴展該框架，另外一方面，也要結合自己的實際場景去定製具體的策略。