restructure

app.py

@@ -1,117 +1,36 @@
 import streamlit as st
-import re
+import model_class
 
-from langchain.embeddings.openai import OpenAIEmbeddings
-from langchain.vectorstores import Pinecone
-from langchain.chat_models import ChatOpenAI
-from langchain.chains import RetrievalQA
-from langchain.output_parsers import OutputFixingParser
-from langchain.schema import OutputParserException
 
-import pinecone
-import json
+# import streamlit as st
 
-model_name = 'text-embedding-ada-002'
+col1, col2 = st.columns(2)
 
-st.title("ICD DEMO")
+with col1:
+#    st.header("A cat")
+#    st.image("https://static.streamlit.io/examples/cat.jpg")
+    # text_box
+    gptkey = st.text_input(label='Enter GPT KEY',placeholder="Enter GPT-API-KEY",label_visibility="collapsed")
 
 
 
-gpt_api_key = st.text_input(label='Enter GPT KEY',placeholder="Enter GPT-API-KEY",label_visibility="collapsed")
-pinecone_api_key=st.text_input(label='Enter PINECONE KEY',placeholder="Enter Pinecone API-KEY",label_visibility="collapsed")
+with col2:
+    pineconekey = st.text_input(label='Enter PINECONE KEY',placeholder="Enter Pinecone API-KEY",label_visibility="collapsed")
 
-embeddings = OpenAIEmbeddings(
-    model=model_name,
-    openai_api_key=gpt_api_key,
-)
 
+# text_box
+# query = st.text_input(label='Enter Query',placeholder="Enter Query",label_visibility="collapsed")
 
-pinecone.init(
-    api_key=pinecone_api_key,
-    environment='gcp-starter',
-)
-index = pinecone.Index('vectordb')
+query = st.text_area(label='Enter Query',placeholder="Enter Query",height=200)
 
+if st.button('Submit'):
+# if query:
+    model_class.model_obj.setAPIKEY(gptkey, pineconekey)
+    model_class.model_obj.initializer()
 
-vectorstore = Pinecone(
-    index=index,
-    embedding_function=embeddings.embed_query,
-    text_key='text',
-)
+    resp = model_class.model_obj.get_response(query)
 
+    # for i in resp:
+        # st.write(i)
 
-
-def get_response(instruction, query):
-    """
-    This function takes in an instruction and a query, and returns a response and a list of results.
-    instruction: str
-    query: str
-    Returns: str, list
-    """
-
-    results = vectorstore.similarity_search(query, k=5)
-    llm = ChatOpenAI(
-        openai_api_key=gpt_api_key,
-        model_name='gpt-3.5-turbo',
-        temperature=0.0,
-        request_timeout=1000
-    )
-    qa = RetrievalQA.from_chain_type(
-        llm=llm,
-        chain_type='stuff',
-        retriever=vectorstore.as_retriever(),
-    )
-    response = qa.run(str(instruction) + str(query))
-
-    return response, results
-    
-
-instruction = """Given the progress notes below, your task is to carefully identify and list all the diagnosis, paying attention to the specific details such as laterality, severity, type, cause, and progression stage, that could influence to find the corresponding International Classification of Diseases (ICD) codes.
-Please exclude any conditions that the patient explicitly denies (e.g., phrases like 'denies,' 'negative for,' etc).
-Following the extraction process, compile the identified conditions in a list, prioritizing conditions of higher severity or urgency at the top, and present the data in a JSON format in descending order based on their priority or severity.
-For example, below is the sample output:
-{
-  "Diseases": [
-    {
-      "Disease": "Fatty Liver",
-      "Laterality": "Not specified",
-      "Severity": "Not specified",
-      "Type": "Not specified",
-      "Cause": "Alcholic",
-      "Progression Stage": "Not specified",
-      "ICD" : "<ICD for Fattly Liver>"
-    },
-    {
-      "Disease": "Leg Fracture",
-      "Laterality": "Right",
-      "Severity": "Not specified",
-      "Type": "Not specified",
-      "Cause": "Accident",
-      "Progression Stage": "Not specified",
-      "ICD" : "<ICD for Leg Fracture>
-     }
-  ]
-}
-"""
-
-if "messages" not in st.session_state:
-    st.session_state.messages = []
-
-for message in st.session_state.messages:
-    with st.chat_message(message["role"]):
-        st.markdown(message["content"])
-
-
-if prompt := st.chat_input("Enter the progress note here"):
-    st.chat_message("user").markdown(prompt)
-    st.session_state.messages.append({"role": "user", "content": prompt})
-    response = get_response(instruction, prompt)
-
-    with st.chat_message("assistant"):
-        output=json.loads(response[0])["Diseases"]
-        output_json_dump = json.dumps(output, indent=4)
-
-        for i in eval(output_json_dump):
-            st.write(i)
-
-    st.session_state.messages.append({"role": "assistant", "content": response})
+    st.write(eval(resp))

instructions.txt

@@ -0,0 +1,26 @@
+Given the progress notes below, your task is to carefully identify and list all the diagnosis, paying attention to the specific details such as laterality, severity, type, cause, and progression stage, that could influence to find the corresponding International Classification of Diseases (ICD) codes.
+Please exclude any conditions that the patient explicitly denies (e.g., phrases like 'denies,' 'negative for,' etc).
+Following the extraction process, compile the identified conditions in a list, prioritizing conditions of higher severity or urgency at the top, and present the data in a JSON format in descending order based on their priority or severity.
+For example, below is the sample output:
+{
+  "Diseases": [
+    {
+      "Disease": "Fatty Liver",
+      "Laterality": "Not specified",
+      "Severity": "Not specified",
+      "Type": "Not specified",
+      "Cause": "Alcholic",
+      "Progression Stage": "Not specified",
+      "ICD" : "<ICD for Fattly Liver>"
+    },
+    {
+      "Disease": "Leg Fracture",
+      "Laterality": "Right",
+      "Severity": "Not specified",
+      "Type": "Not specified",
+      "Cause": "Accident",
+      "Progression Stage": "Not specified",
+      "ICD" : "<ICD for Leg Fracture>
+     }
+  ]
+}

model_class.py

@@ -0,0 +1,99 @@
+import re
+
+from langchain.embeddings.openai import OpenAIEmbeddings
+from langchain.vectorstores import Pinecone
+from langchain.chat_models import ChatOpenAI
+from langchain.chains import RetrievalQA
+# from langchain.output_parsers import OutputFixingParser
+# from langchain.schema import OutputParserException
+
+import pinecone
+
+
+
+
+class model:
+    def __init__(self) -> None:
+        self.gpt_api = None
+        self.pinecone_api = None
+        self.model_name = 'text-embedding-ada-002'
+
+        self.instructions = self.getInstructions()
+
+        self.embeddings = None
+        self.index = None
+        self.vectorstore = None
+    
+    def setAPIKEY(self, key1, key2):
+        self.gpt_api = key1
+        self.pinecone_api = key2
+
+        
+    
+    def initializer(self):
+        self.embeddings = OpenAIEmbeddings(
+            model=self.model_name,
+            openai_api_key=self.gpt_api,
+        )
+
+        pinecone.init(
+            api_key=self.pinecone_api,
+            environment='gcp-starter',
+        )
+        self.index = pinecone.Index('vectordb') 
+
+        # print("Initialized", self.index)
+        # print(self.instructions)
+
+
+        self.vectorstore = Pinecone(
+            index=self.index,
+            embedding_function=self.embeddings.embed_query,
+            text_key='text',
+        )
+    
+    def getInstructions(self):
+        with open("instructions.txt", "r") as f:
+            instructions = f.read()
+        return instructions
+
+    def get_response(self, query):
+        """
+        This function takes in an instruction and a query, and returns a response and a list of results.
+        instruction: str
+        query: str
+        Returns: str, list
+        """
+
+        results = self.vectorstore.similarity_search(query, k=5)
+        llm = ChatOpenAI(
+            openai_api_key=self.gpt_api,
+            model_name='gpt-3.5-turbo',
+            temperature=0.0,
+            request_timeout=1000
+        )
+        qa = RetrievalQA.from_chain_type(
+            llm=llm,
+            chain_type='stuff',
+            retriever=self.vectorstore.as_retriever(),
+        )
+        response = qa.run(str(self.instructions) + str(query))
+
+        return response
+
+
+
+# Main
+
+model_obj = model()
+
+
+# response.setAPIKEY("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", "ce1dd6d6-6783-44ba-ac68-57cff520df1e")
+# response.initializer()
+
+
+# query="""
+# 59 y/o male patient with h/o dyspepsia, esophageal ulcer, and overweight is here for a follow up after a colonoscopy/EGD. Colonoscopy show examined portion of the ileum was normal, a single colonic ulcer in the cecum, erythematous mucosa in the rectum, and internal hemorrhoids. Advise patient to repeat colonoscopy based on pathology for surveillance. EGD show LA grade C esophagitis, esophageal ulcers, gastroparesis, gastritis, and a normal duodenum. Path report show no significant histopathology on esophagus or stomach. Patient advised to repeat EGD in 3 months to check for healing and for surveillance. Results discussed with patient. Significance of findings and importance of surveillance explained to patient. Anti reflux measures including increase physical activity explained to patient. All patient questions were answered. Patient understand and agrees to plan.
+# """
+
+# print(response.get_response(query))