Test corpus

app.py

@@ -26,10 +26,11 @@ with col2:
     st.image("images/OpenWormLogo.png")
 
 
-tab_free, tab_panel, tab_pubs, tab_data, tab_model = st.tabs(["Individual LLMs", 
+tab_free, tab_panel, tab_pubs, tab_data, tab_corpus, tab_model = st.tabs(["Individual LLMs", 
                                          "Panel discussion", 
                                          "Publications",
                                          "Structured data",
+                                         "OpenWorm Corpus",
                                          "Run model"])
 
 
@@ -120,6 +121,25 @@ with tab_data:
 
             st.info(response)
 
+with tab_corpus:
+
+    from corpus_query import run_query
+
+    st.markdown("**Search a (small) corpus of _C. elegans_ literature** Uses: [paper-qa](https://github.com/whitead/paper-qa)")
+
+    with st.form("form_corpus"):
+        
+        text = st.text_input("Enter a query:", "What types of ion channels are present in C. elegans neurons?")
+
+        llm_ver = st.selectbox('Which LLM version should I use?', (LLM_LLAMA2, LLM_GPT4))
+
+        submitted = st.form_submit_button("Submit")
+
+        if submitted:
+            response = run_query(query=text, llm_ver=llm_ver)
+
+            st.info(response)
+
 with tab_model:
 
     st.markdown("**Run a _C. elegans_ cell model**")

corpus_query.py

@@ -0,0 +1,65 @@
+from paperqa import Docs
+import pickle
+
+from ring import LLM_GPT35
+from ring import LLM_GPT4
+from ring import LLM_LLAMA2
+from ring import get_llm
+
+my_docs = ['corpus/MirzakhaliliEtAl2018.pdf',
+           'corpus/NicollettiEtAl2019.pdf',
+           'corpus/ATributetoSydneyBrenner.pdf']
+
+#my_docs = ['corpus/ATributetoSydneyBrenner.pdf']
+
+llm_to_use = LLM_GPT4
+llm = get_llm(llm_to_use, 0.1)
+
+CORPUS_PICKLE = 'corpus.pkl'
+
+def generate_pickle():
+    docs = Docs(llm="langchain",
+                client=llm)
+
+    print('Initialised Docs with LLM: %s'%llm_to_use)
+
+    for d in my_docs:
+        print('Adding %s'%d)
+        docs.add(d)
+
+    print('Saving pickle of docs to: %s'%CORPUS_PICKLE)
+
+    f = open(CORPUS_PICKLE,'wb')
+    pickle.dump(docs, f)
+    f.close()
+ 
+
+def run_query(query, llm_ver):
+
+    f = open(CORPUS_PICKLE,'rb')
+
+    docs = pickle.load(f)
+
+    llm = get_llm(llm_ver, 0)
+
+    docs.set_client(llm)
+    answer = docs.query(query)
+
+    return answer
+
+
+if __name__ == '__main__':
+
+    #generate_pickle()
+
+    print('Running queries')
+
+    queries = ['Where did John Sulston do his postdoc?',"What types of potassium channels are expressed in C. elegans?"]
+
+    for query in queries:
+        print('-------------------')
+        print('Q: %s'%query)
+        #answer = docs.query(query)
+        answer = run_query(query, LLM_GPT4)
+
+        print('A: %s'%answer)

requirements.txt

@@ -8,4 +8,5 @@ wormneuroatlas
 neuromllite
 matplotlib
 pylems
-pyneuroml
+pyneuroml
+paper-qa==4.0.0rc8