Use Nvidia NIM API as LLM Client.

SNOMED-CT_Assistant.py

@@ -14,16 +14,28 @@ remote = True
 
 if remote:
     with st.sidebar:
-        if 'OPENAI_API_TOKEN' in st.secrets:
+        if 'NVIDIA_NIM_KEY' in st.secrets:
             st.success('API key already provided!', icon='✅')
-            openai_api_key = st.secrets['OPENAI_API_TOKEN']
+            nvidia_nim_key = st.secrets['NVIDIA_NIM_KEY']
 else:
     load_dotenv()
     openai_api_key = os.environ.get("OpenAI_API_KEY")
+    nvidia_nim_key = os.environ.get("NVIDIA_NIM_KEY")
 
 st.title("🏥 SNOMED-CT Assistant")
 st.caption("👩‍⚕️ A smart medical assistant with SNOMED-CT knowledge.")
 
+# Chroma DB Client
+chroma_client = chromadb.PersistentClient(path="snomed_ct_id_term_1410k")
+collection = chroma_client.get_or_create_collection(name="snomed_ct_id_term")
+
+# NIM Client Configuration
+client = OpenAI(
+    base_url = "https://integrate.api.nvidia.com/v1",
+    api_key = nvidia_nim_key
+)
+model_tag = "meta/llama3-70b-instruct"
+
 # System prompt
 system_prompt = """You are a medical expert with rich experience in SNOMED-CT professional knowledge. 
 You are skilled at assisting medical professionals and answering questions in the medical field. 
@@ -35,19 +47,19 @@ Please refuse to answer inquiries and requests unrelated to the medical field, i
 As an experienced professional, you possess deep expertise in the field of SNOMED CT Entity Linking. 
 You have a thorough understanding of the relevant workflows and critical aspects involved, encompassing:
 - Adept handling of electronic medical record (EMR) data processing
-- Entity Identification, Proficient entity recognition capabilities, identifying and extracting relevant medical concepts from unstructured text
+- *Entity Identification*, Proficient entity recognition capabilities, identifying and extracting relevant medical concepts from unstructured text
 - Skilled Entity Mapping, accurately linking identified entities to their corresponding SNOMED CT concepts
 - Seamless integration and output of clinical terminology, ensuring the accurate representation and utilization of standardized medical language
 - Patiently and professionally respond to all SNOMED CT related inquiries, even if the user repeats questions.
 - Demonstrate deep expertise in the standard SNOMED CT Entity Linking workflow, which involves:
-  1. Performing Entity Identification to extract relevant medical terminology from the input.
-  2. Conducting Entity Mapping to link the identified entities to their corresponding SNOMED CT concepts.
-- Present the results in a tabular format only with the following 3 columns: "Identified Entity", "SNOMED CT Concept IDs", "SNOMED CT Descriptions".
-
-Here is the practical entity linking process example:
-- the input text in EHRs: "Patient referred for a biopsy to investigate potential swelling in upper larynx."
-- the identified entity: "biopsy", "larynx"
+- **All YOU CAN DO** : Performing **Entity Identification** : Try to extract relevant medical terminology from the medical text input.
+ 
+Here is the practical entity identification process example:
+- the input text will the part of EHRs record: "Patient referred for a biopsy to investigate potential swelling in upper larynx."
+- if the identified entity: "biopsy", "larynx"
 - response the identified entities with JSON format: {"identified_entity" : ["biopsy", "larynx"]}
+- If no identifiable entity is found in the input text, return an empty list: {"identified_entity" : []}
+- **DON't response the other format besides JSON**
 - During Entity Identification processing, if the original medical text data clearly contains commonly used medical abbreviations, convert the abbreviations into their full names, and provide the original abbreviations in parentheses for easy reference.
 - For example: "The patient has the multiple disease, including T2D, CAD, HTN, CKD etc. decreased T3 and T4 levels."
 - T2D: "Type 2 Diabetes Mellitus", CAD: "Coronary Artery Disease", HTN: "Hypertension", CKD: "Chronic Kidney Disease", T3: "Triiodothyronine", T4: "Thyroxine"
@@ -65,6 +77,8 @@ Numbers or units related symbols are not included in this range and can be ignor
 Output Format Requirements (Must follow):
 - As default, only process "Entity Identification", and find out the entity related to SNOMED CT terms.
 - Present the results in JSON format, like:  {"identified_entity" : ["biopsy", "larynx"]}
+- If no identifiable entity is found in the input text, return an empty list: {"identified_entity" : []}
+- **DON't response the other format besides JSON**
 """
 
 
@@ -90,10 +104,6 @@ def generate_entity_identification_prompt(medical_text):
 def generate_entity_mapping_prompt(entity, query_result_dict):
     return f"""Help me to do "SNOMED-CT Entity Mapping" process with entity: {entity} and query result \n {query_result_dict} \n , output with table format, including 5 columns: "Identified Entity", "Distance", "IDs", "SNOMED CT Concept IDs", "SNOMED CT Descriptions"  \n """
 
-# Chroma DB Client
-chroma_client = chromadb.PersistentClient(path="snomed_ct_id_term_1410k")
-collection = chroma_client.get_or_create_collection(name="snomed_ct_id_term")
-
 # Func: query chrome_db
 def query_chroma_db(query_text, query_number):
     results = collection.query(
@@ -108,19 +118,17 @@ def get_dict_from_chroma_results(results):
     result_dict = {'ids': results['ids'][0], 'concept_ids': [ str(sub['concept_id']) for sub in results['metadatas'][0] ], 'distances': results['distances'][0], 'descriptions': results['documents'][0]}
     return result_dict
 
-
-# OpenAI Client Configuration
-client = OpenAI(api_key=openai_api_key)
-model_tag = "gpt-3.5-turbo"
-
-# Chat Session with OpenAI API
+# Chat Session with NIM API
 def chat_input(prompt, med_text):
     st.session_state.messages.append({"role": "user", "content": med_text})
     st.chat_message("user").write(med_text)
     with st.spinner("Thinking..."):
         entity_identification_response = client.chat.completions.create(
-            model=model_tag, response_format={ "type": "json_object" }, messages=st.session_state.messages, temperature=0.5)
+            model=model_tag, 
+            messages=st.session_state.messages, 
+            temperature=0.5)
         msg = entity_identification_response.choices[0].message.content
+        print("NIM output" + msg)
         entity_list = json.loads(msg)["identified_entity"]
         st.session_state.messages.append({"role": "assistant", "content": msg})
         st.chat_message("assistant").write(msg)
@@ -163,7 +171,7 @@ def entity_mapping_result_to_table(entity, results_dict):
 
 if "messages" not in st.session_state:
     st.session_state["messages"] = [{"role": "system", "content": system_prompt},
-                                    {"role": "assistant", "content": "👩‍⚕️ 您好，我是您的專業醫學助理。請問有任何我可以協助你的地方嗎?"}]
+                                    {"role": "assistant", "content": "👩‍⚕️ Hello, I am your professional medical assistant. Is there anything I can assist you with?"}]
 
 for msg in st.session_state.messages:
     if msg["role"] == "system":