working on ent buttons to update query

app.py

@@ -1,7 +1,6 @@
 import streamlit as st
 import textract
 import tempfile
-from typing import List
 import spacy
 from spacy.tokens import DocBin, Doc, Span
 from collections import Counter
@@ -13,6 +12,8 @@ with open("style.css") as f:
     st.markdown("<style>" + f.read() + "</style>", unsafe_allow_html=True)
     
 st.title('Index and Search a Collection of Documents')
+if 'query' not in st.session_state:
+    st.session_state['query'] = ''
 
 @st.cache
 def download_model(select_model:str):
@@ -22,7 +23,7 @@ def download_model(select_model:str):
     except Exception as e:
         return False
 
-def search_docs(query:str, documents:List[Doc], nlp) -> List[Span]:
+def search_docs(query:str, documents:list[Doc], nlp) -> list[Span]:
     terms = query.split('|')
     patterns = [nlp.make_doc(text) for text in terms]
     matcher = PhraseMatcher(nlp.vocab)
@@ -36,6 +37,9 @@ def search_docs(query:str, documents:List[Doc], nlp) -> List[Span]:
             
     return results
 
+def update_query(arg:str):
+    st.session_state.query = arg
+
 models = srsly.read_json('models.json')
 models[''] = [] #require the user to choose a language
 languages = models.keys()
@@ -53,8 +57,10 @@ if language:
             
 
             uploaded_files = st.file_uploader("Select files to process", accept_multiple_files=True)
-            query = st.sidebar.text_input(label="Enter your query (use | to separate search terms)", value="...")
+            st.session_state.query = st.sidebar.text_input(label="Enter your query (use | to separate search terms)", value="...")
+            
             documents = []
+            all_ents = []
             for uploaded_file in uploaded_files:
                 file_type = uploaded_file.type
                 file_suffix = '.' + uploaded_file.name.split('.')[-1]
@@ -66,24 +72,24 @@ if language:
                     doc = nlp(text)
                     doc.user_data['filename'] = uploaded_file.name
                     documents.append(doc)
-                    ent_freq = Counter([ent.label_ for ent in doc.ents])
-                    for key, value in ent_freq.items():
-                        if st.sidebar.button(key, key=key):
-                            #There are multiple identical st.button widgets with key='DATE'.
-                            #To fix this, please make sure that the key argument is unique for each st.button you create.
-                            # need to pool results from several files, then create buttons and search links
-                            st.sidebar.write(value)
-                            for ent in doc.ents:
-                                if ent.label_ == key:
-                                    #TODO Make clickable button that updates query
-                                    #TODO also return distinct ents
-                                    st.sidebar.write(ent.text)
+                    for ent in doc.ents:
+                        all_ents.append(ent)
+                    
                         
                 except Exception as e:
                     st.error(e)
-            
+
+            ents_container = st.container()   
+            label_freq = Counter([ent.label_ for ent in all_ents])
+            for key, value in label_freq.items():
+                if st.sidebar.button(key, key=key):
+                    st.sidebar.write(value)
+                    text_freq = Counter([ent.text for ent in all_ents if ent.label_ == key])
+                    for text in text_freq.keys():
+                        st.sidebar.button(f'{text} ({text_freq[text]})', on_click=update_query, args=(text, ))
+                        
             results_container = st.container()   
-            results = search_docs(query, documents,nlp)
+            results = search_docs(st.session_state.query, documents,nlp)
             for result in results:
                 doc = result.doc
                 sent_before = doc[result.sent.start:result.start]