implement RAG

Author: marswen

FullTabulation-ICD-11-MMS-zh.xlsx

@@ -3,8 +3,10 @@
 import json
 import yaml
 import prompts
+import pandas as pd
 import streamlit as st
 import streamlit_authenticator as stauth
+from icd import SemanticSearch
 from dotenv import load_dotenv
 from yaml.loader import SafeLoader
 from langchain.prompts import PromptTemplate
@@ -15,8 +17,18 @@
 
 
 load_dotenv()
-with open(os.path.join(os.path.dirname(__file__), 'code_book.txt')) as f:
-    code_book = f.readlines()
+term_df = pd.read_excel('ICD-10-ICD-O.xlsx')
+term_df = term_df.loc[~term_df['Coding System'].isin(['ICD-O-3行为学编码', 'ICD-O-3组织学等级和分化程度编码'])]
+term_map = dict(zip(term_df['Code'], term_df['释义']))
+icd10_semantic_search = SemanticSearch('./vs/icd10')
+icdo3_semantic_search = SemanticSearch('./vs/icdo3')
+
+
+def search_reference0(code):
+    code_comps = code.split('-')
+    if len(code_comps) > 1:
+        code = '-'.join(code.split('-')[1:])
+    return '\n'.join([line for line in code_book if re.search(re.escape(code), line, re.I) is not None])
 
 
 def search_reference(code):
@@ -26,27 +38,53 @@ def search_reference(code):
     return '\n'.join([line for line in code_book if re.search(re.escape(code), line, re.I) is not None])
 
 
-def generate(context, chat_llm, callbacks):
+def generate(context, chat_llm, callbacks, output_container):
     system_message = SystemMessage(content=prompts.prompt1)
-    initial_user_prompt = PromptTemplate(template=prompts.prompt2, input_variables=['diagnosis']).format(context=context)
+    related_icd10 = icd10_semantic_search.search(context, k=5)
+    related_icdo3 = icdo3_semantic_search.search(context, k=5)
+    related_code_context = '\n'.join([f'{x[0]["Code"]}\n{x[0]["释义"]}' for x in related_icd10 + related_icdo3])
+    initial_user_prompt = PromptTemplate(template=prompts.prompt2, input_variables=['context', 'related_codes']).format(context=context, related_codes=related_code_context)
     initial_user_message = HumanMessage(content=initial_user_prompt)
+    output_container.chat_message("user").write(system_message.content.replace('\n', '\n\n'))
+    output_container.chat_message("user").write(initial_user_message.content.replace('\n', '\n\n'))
     initial_result = chat_llm([system_message, initial_user_message], callbacks=callbacks)
     second_user_message = HumanMessage(content=prompts.prompt3)
+    output_container.chat_message("user").write(second_user_message.content.replace('\n', '\n\n'))
     second_result = chat_llm([system_message, initial_user_message, initial_result, second_user_message], callbacks=callbacks)
     code_result = second_result
     try_cnt = 0
     while True:
         format_user_prompt = HumanMessage(content=prompts.prompt4)
+        output_container.chat_message("user").write(format_user_prompt.content.replace('\n', '\n\n'))
         format_result = chat_llm([code_result, format_user_prompt], callbacks=callbacks)
         json_text = re.search('```json(.+)```', format_result.content, re.DOTALL)
         if json_text is not None:
             json_data = json.loads(json_text.group(1))
             references = ''
             for code in json_data['code'][:3]:
-                ref = search_reference(code)
-                references += f'{code}:\n{ref}\n\n'
+                if re.search('[\:：]', code) is not None:
+                    code = re.split('[\:：]', code)[-1].strip()
+                if re.search('^[A-Z]-', code):
+                    code = '-'.join(code.split('-')[1:])
+                ref = term_map.get(code)
+                if ref is not None:
+                    references += f'{code}:\n{ref}\n\n'
+                if re.search('\d\.\-', code):
+                    related_codes = [x for x in term_map.keys() if x.startswith(code.strip('-'))]
+                    for rel in related_codes:
+                        ref = term_map.get(rel)
+                        references += f'{rel}:\n{ref}\n\n'
+                if re.search('\d\.\d\-\d', code):
+                    code_compo = re.search('(.*\d\.)(\d)\-(\d)', code)
+                    related_codes = [code_compo.group(1) + str(x) for x in
+                                     range(int(code_compo.group(2)), int(code_compo.group(3)) + 1)]
+                    related_codes = [x for x in related_codes if x in term_map]
+                    for rel in related_codes:
+                        ref = term_map.get(rel)
+                        references += f'{rel}:\n{ref}\n\n'
             refine_user_prompt = PromptTemplate(template=prompts.prompt5, input_variables=['references']).format(references=references)
             refine_user_message = HumanMessage(content=refine_user_prompt)
+            output_container.chat_message("user").write(refine_user_message.content.replace('\n', '\n\n'))
             refine_result = chat_llm([system_message, initial_user_message, initial_result, second_user_message, code_result, refine_user_message], callbacks=callbacks)
             code_result = refine_result
             if '"confirmed": true' in code_result.content:
@@ -55,6 +93,7 @@ def generate(context, chat_llm, callbacks):
         if try_cnt > 5:
             break
     format_user_prompt = HumanMessage(content=prompts.prompt4)
+    output_container.chat_message("user").write(format_user_prompt.content.replace('\n', '\n\n'))
     format_result = chat_llm([code_result, format_user_prompt], callbacks=callbacks)
     return format_result.content
 
@@ -74,8 +113,18 @@ def demo_page():
         st_callback = CustomStreamlitCallbackHandler(output_container)
         std_callback = StreamingStdOutCallbackHandler()
         callbacks = [st_callback, std_callback]
-        result = generate(raw_input, chat_llm, callbacks=callbacks)
+        result = generate(raw_input, chat_llm, callbacks, output_container)
         st.markdown(result)
+    st.write("""
+    <hr style="border: none; border-top: 1px solid #ccc;">
+    <p style="text-align: center; font-size: 12px;">
+        沪ICP备18007075号-2
+    </p>
+    """, unsafe_allow_html=True)
+    # st.write("""
+    # ---
+    # 沪ICP备18007075号-2
+    # """)
 
 
 if __name__ == '__main__':

ICD-10-ICD-O.xlsx

@@ -0,0 +1,4 @@
+from modelscope.hub.snapshot_download import snapshot_download
+
+model_dir = snapshot_download('AI-ModelScope/bge-large-zh-v1.5', cache_dir='./')
+

medcodegpt.py

@@ -59,6 +59,9 @@
 
 根据实际情况，肿瘤报告的诊断名称往往不能与手册中的名称完全符合，应选择其中符合程度最高的编码名称。
 部分疾病（如部分肝癌、黑色素瘤、间皮瘤和淋巴瘤、白血病等）编码较特殊，本手册已在《ICD-10与ICD-O-3解剖部位编码》中列出了ICD-10与ICD-O-3不同的解剖部位编码和/或对应的ICD-O-3形态学编码；同样在《ICD-O-3形态学编码》中也已经列出了对应的解剖部位编码，请使用者酌情参考。
+
+以下编码可能与目标编码相关，但最终结果不限于以下编码，仅供参考。
+{related_codes}
 '''
 
 prompt3 = '''

models/download_models.py

@@ -5,3 +5,4 @@ python-dotenv==1.0.1
 PyYAML==6.0.1
 streamlit==1.32.0
 streamlit_authenticator==0.3.2
+sentence_transformers==3.0.1