Bug fix: find_result cache breaks down upon concurrent users

README.md

@@ -4,7 +4,7 @@ emoji: 📖
 colorFrom: blue
 colorTo: green
 sdk: gradio
-sdk_version: 4.36.0
+sdk_version: 4.44.1
 app_file: app.py
 pinned: false
 license: cc-by-nc-sa-4.0

app.py

@@ -150,16 +150,16 @@ def search_docs(index_desc, query, maxnum, max_disp_len, max_clause_freq, max_di
         docs.append([])
     return tuple([latency, tokenization_info, message] + metadatas + docs)
 
-find_result = None
-
-def search_docs_new(index_desc, query, max_disp_len, max_clause_freq, max_diff_tokens):
-    global find_result
+def search_docs_new(index_desc, query, max_disp_len, max_clause_freq, max_diff_tokens, state):
     if ' AND ' in query or ' OR ' in query: # CNF query
         find_result = process('find_cnf', index_desc, query=query, max_clause_freq=max_clause_freq, max_diff_tokens=max_diff_tokens)
         find_result['type'] = 'cnf'
     else: # simple query
         find_result = process('find', index_desc, query=query)
         find_result['type'] = 'simple'
+
+    state = find_result
+
     latency = '' if 'latency' not in find_result else f'{find_result["latency"]:.3f}'
     tokenization_info = format_tokenization_info(find_result)
     if 'error' in find_result:
@@ -167,7 +167,7 @@ def search_docs_new(index_desc, query, max_disp_len, max_clause_freq, max_diff_t
         idx = gr.Number(minimum=0, maximum=0, step=1, value=0, interactive=False)
         metadata = ''
         doc = []
-        return latency, tokenization_info, message, idx, metadata, doc
+        return latency, tokenization_info, message, idx, metadata, doc, state
 
     if ' AND ' in query or ' OR ' in query: # CNF query
         ptrs_by_shard = find_result['ptrs_by_shard']
@@ -183,21 +183,20 @@ def search_docs_new(index_desc, query, max_disp_len, max_clause_freq, max_diff_t
         idx = gr.Number(minimum=0, maximum=0, step=1, value=0, interactive=False)
         metadata = ''
         doc = []
-        return latency, tokenization_info, message, idx, metadata, doc
+        return latency, tokenization_info, message, idx, metadata, doc, state
     idx = random.randint(0, cnt_retrievable-1)
-    metadata, doc = get_another_doc(index_desc, idx, max_disp_len)
+    metadata, doc = get_another_doc(index_desc, idx, max_disp_len, state)
     idx = gr.Number(minimum=0, maximum=cnt_retrievable-1, step=1, value=idx, interactive=True)
-    return latency, tokenization_info, message, idx, metadata, doc
+    return latency, tokenization_info, message, idx, metadata, doc, state
 
-def clear_search_docs_new():
-    global find_result
-    find_result = None
+def clear_search_docs_new(state):
+    state = None
     idx = gr.Number(minimum=0, maximum=0, step=1, value=0, interactive=False)
-    return idx
+    return idx, state
 
-def get_another_doc(index_desc, idx, max_disp_len):
-    global find_result
-    if not (type(idx) == int and 0 <= idx and idx < find_result['cnt']):
+def get_another_doc(index_desc, idx, max_disp_len, state):
+    find_result = state
+    if find_result is None or not (type(idx) == int and 0 <= idx and idx < find_result['cnt']):
         metadata = ''
         doc = []
         return metadata, doc
@@ -230,10 +229,10 @@ def get_another_doc(index_desc, idx, max_disp_len):
 with gr.Blocks() as demo:
     with gr.Column():
         gr.HTML(
-            '''<h1 text-align="center">Infini-gram: An Engine for n-gram / ∞-gram Language Modeling with Trillion-Token Corpora</h1>
+            '''<h1 text-align="center">Infini-gram: An Efficient Search Engine over the Massive Pretraining Datasets of Language Models</h1>
 
-            <p style='font-size: 16px;'>This is an engine that processes n-gram / ∞-gram queries on massive text corpora. Please first select the corpus and the type of query, then enter your query and submit.</p>
-            <p style='font-size: 16px;'>The engine is developed by <a href="https://liujch1998.github.io">Jiacheng (Gary) Liu</a> and documented in our paper: <a href="https://huggingface.co/papers/2401.17377">Infini-gram: Scaling Unbounded n-gram Language Models to a Trillion Tokens</a>. Feel free to check out our <a href="https://infini-gram.io">Project Homepage</a>.</p>
+            <p style='font-size: 16px;'>This engine does exact-match search over several open pretraining datasets of language models. Please first select the corpus and the type of query, then enter your query and submit.</p>
+            <p style='font-size: 16px;'>The engine is developed by <a href="https://liujch1998.github.io">Jiacheng Liu</a> and documented in our paper: <a href="https://huggingface.co/papers/2401.17377">Infini-gram: Scaling Unbounded n-gram Language Models to a Trillion Tokens</a>. Feel free to check out our <a href="https://infini-gram.io">Project Homepage</a>.</p>
             <p style='font-size: 16px;'><b>API Endpoint:</b> If you'd like to issue batch queries to infini-gram, you may invoke our API endpoint. Please refer to the <a href="https://infini-gram.io/api_doc">API documentation</a>.</p>
             <p style='font-size: 16px;'><b>Note:</b> The query is <b>case-sensitive</b>. Your query will be tokenized with the Llama-2 tokenizer (unless otherwise specified).</p>
             '''
@@ -482,10 +481,29 @@ with gr.Blocks() as demo:
                                 search_docs_new_idx = gr.Slider(label='', minimum=0, maximum=0, step=1, value=0, interactive=False)
                                 search_docs_new_metadata = gr.Textbox(label='Metadata', lines=3, max_lines=3, interactive=False)
                                 search_docs_new_output = gr.HighlightedText(label='Document', show_legend=False, color_map={"-": "red", "0": "green", "1": "cyan", "2": "blue", "3": "magenta"})
+                    search_docs_state = gr.State(value=None)
                     search_docs_new_clear.add([search_docs_new_query, search_docs_new_latency, search_docs_new_tokenized, search_docs_new_message, search_docs_new_idx, search_docs_new_metadata, search_docs_new_output])
-                    search_docs_new_clear.click(clear_search_docs_new, inputs=[], outputs=[search_docs_new_idx], api_name=False)
-                    search_docs_new_submit.click(search_docs_new, inputs=[index_desc, search_docs_new_query, search_docs_new_max_disp_len, search_docs_new_max_clause_freq, search_docs_new_max_diff_tokens], outputs=[search_docs_new_latency, search_docs_new_tokenized, search_docs_new_message, search_docs_new_idx, search_docs_new_metadata, search_docs_new_output], api_name=False)
-                    search_docs_new_idx.input(get_another_doc, inputs=[index_desc, search_docs_new_idx, search_docs_new_max_disp_len], outputs=[search_docs_new_metadata, search_docs_new_output], api_name=False)
+                    search_docs_new_clear.click(
+                        clear_search_docs_new,
+                        inputs=[search_docs_state],
+                        outputs=[search_docs_new_idx, search_docs_state]
+                    )
+                    search_docs_new_submit.click(
+                        search_docs_new,
+                        inputs=[index_desc, search_docs_new_query, search_docs_new_max_disp_len,
+                               search_docs_new_max_clause_freq, search_docs_new_max_diff_tokens,
+                               search_docs_state],
+                        outputs=[search_docs_new_latency, search_docs_new_tokenized,
+                                search_docs_new_message, search_docs_new_idx,
+                                search_docs_new_metadata, search_docs_new_output,
+                                search_docs_state]
+                    )
+                    search_docs_new_idx.input(
+                        get_another_doc,
+                        inputs=[index_desc, search_docs_new_idx, search_docs_new_max_disp_len,
+                               search_docs_state],
+                        outputs=[search_docs_new_metadata, search_docs_new_output]
+                    )
 
         with gr.Row():
             gr.Markdown('''