Sync changes

app.py

@@ -1,42 +1,23 @@
 import gradio as gr
 import datetime
 import json
-import os
 import requests
-import time
 from constants import *
 
-API_IPADDR = os.environ.get('API_IPADDR', None)
-default_concurrency_limit = os.environ.get('default_concurrency_limit', 10)
-max_size = os.environ.get('max_size', 100)
-max_threads = os.environ.get('max_threads', 40)
-debug = (os.environ.get('debug', 'False') != 'False')
-
-last_query_time_by_ip = {}
-
-def process(corpus_desc, query_desc, query, ret_num, request: gr.Request):
-    global last_query_time_by_ip
-    ip = request.client.host if request else ''
+def process(query_type, corpus_desc, engine_desc, query, maxnum, request: gr.Request):
     timestamp = datetime.datetime.now().strftime('%Y%m%d-%H%M%S')
-    t = time.time()
-    last_query_time = 0 if ip == '' else last_query_time_by_ip.get(ip, 0)
-    blocked = (t - last_query_time < MIN_QUERY_INTERVAL_SECONDS)
-
     corpus = CORPUS_BY_DESC[corpus_desc]
-    query_type = QUERY_TYPE_BY_DESC[query_desc]
+    engine = ENGINE_BY_DESC[engine_desc]
     data = {
         'timestamp': timestamp,
-        'ip': ip,
-        'blocked': blocked,
-        'corpus': corpus,
         'query_type': query_type,
+        'corpus': corpus,
+        'engine': engine,
         'query': query,
     }
+    if maxnum is not None:
+        data['maxnum'] = maxnum
     print(json.dumps(data))
-    if blocked:
-        return tuple([f'You queried too frequently. Please try again in {MIN_QUERY_INTERVAL_SECONDS} seconds.'] + [''] * (ret_num - 1))
-    if ip != '':
-        last_query_time_by_ip[ip] = t
     if API_IPADDR is None:
         raise ValueError(f'API_IPADDR envvar is not set!')
     response = requests.post(f'http://{API_IPADDR}:5000/', json=data)
@@ -44,58 +25,35 @@ def process(corpus_desc, query_desc, query, ret_num, request: gr.Request):
         result = response.json()
     else:
         raise ValueError(f'HTTP error {response.status_code}: {response.json()}')
-    if debug:
+    if DEBUG:
         print(result)
     return result
 
-def process_1(corpus_desc, query_desc, query, request: gr.Request):
-    return process(corpus_desc, query_desc, query, 1, request)
-def process_2(corpus_desc, query_desc, query, request: gr.Request):
-    return process(corpus_desc, query_desc, query, 2, request)
-def process_3(corpus_desc, query_desc, query, request: gr.Request):
-    return process(corpus_desc, query_desc, query, 3, request)
-
-def process_ard_cnf_multi(corpus_desc, query_desc, query, maxnum, request: gr.Request):
-    global last_query_time_by_ip
-    ip = request.client.host if request else ''
-    timestamp = datetime.datetime.now().strftime('%Y%m%d-%H%M%S')
-    t = time.time()
-    last_query_time = 0 if ip == '' else last_query_time_by_ip.get(ip, 0)
-    blocked = (t - last_query_time < MIN_QUERY_INTERVAL_SECONDS)
-
-    corpus = CORPUS_BY_DESC[corpus_desc]
-    query_type = QUERY_TYPE_BY_DESC[query_desc]
-    timestamp = datetime.datetime.now().strftime('%Y%m%d-%H%M%S')
-    data = {
-        'timestamp': timestamp,
-        'ip': ip,
-        'blocked': blocked,
-        'corpus': corpus,
-        'query_type': query_type,
-        'query': query,
-        'maxnum': maxnum,
-    }
-    print(json.dumps(data))
-    if blocked:
-        return tuple([f'You queried too frequently. Please try again in {MIN_QUERY_INTERVAL_SECONDS} seconds.'] + [''] * 11)
-    if ip != '':
-        last_query_time_by_ip[ip] = t
-    if API_IPADDR is None:
-        raise ValueError(f'API_IPADDR envvar is not set!')
-    response = requests.post(f'http://{API_IPADDR}:5000/', json=data)
-    if response.status_code == 200:
-        result = response.json()
-    else:
-        raise ValueError(f'HTTP error {response.status_code}: {response.json()}')
-    if debug:
-        print(result)
-    if len(result) != 3:
-        raise ValueError(f'Invalid result: {result}')
-    outputs, output_tokens, message = result[0], result[1], result[2]
+def count(corpus_desc, engine_desc, query, request: gr.Request):
+    result = process('count', corpus_desc, engine_desc, query, None, request)
+    return result.get('latency', ''), result.get('tokenized', ''), result.get('count', '')
+def prob(corpus_desc, engine_desc, query, request: gr.Request):
+    result = process('prob', corpus_desc, engine_desc, query, None, request)
+    return result.get('latency', ''), result.get('tokenized', ''), result.get('probability', '')
+def ntd(corpus_desc, engine_desc, query, request: gr.Request):
+    result = process('ntd', corpus_desc, engine_desc, query, None, request)
+    return result.get('latency', ''), result.get('tokenized', ''), result.get('distribution', '')
+def infgram_prob(corpus_desc, engine_desc, query, request: gr.Request):
+    result = process('infgram_prob', corpus_desc, engine_desc, query, None, request)
+    return result.get('latency', ''), result.get('tokenized', ''), result.get('longest_suffix', ''), result.get('probability', '')
+def infgram_ntd(corpus_desc, engine_desc, query, request: gr.Request):
+    result = process('infgram_ntd', corpus_desc, engine_desc, query, None, request)
+    return result.get('latency', ''), result.get('tokenized', ''), result.get('longest_suffix', ''), result.get('distribution', '')
+def search_docs(corpus_desc, engine_desc, query, maxnum, request: gr.Request):
+    result = process('search_docs', corpus_desc, engine_desc, query, maxnum, request)
+    outputs = result.get('outputs', [])
     outputs = outputs[:maxnum]
     while len(outputs) < 10:
         outputs.append([])
-    return message, output_tokens, outputs[0], outputs[1], outputs[2], outputs[3], outputs[4], outputs[5], outputs[6], outputs[7], outputs[8], outputs[9]
+    return result.get('latency', ''), result.get('tokenized', ''), result.get('message', ''), outputs[0], outputs[1], outputs[2], outputs[3], outputs[4], outputs[5], outputs[6], outputs[7], outputs[8], outputs[9]
+def analyze_document(corpus_desc, engine_desc, query, request: gr.Request):
+    result = process('analyze_document', corpus_desc, engine_desc, query, None, request)
+    return result.get('latency', ''), result.get('html', '')
 
 with gr.Blocks() as demo:
     with gr.Column():
@@ -103,196 +61,181 @@ with gr.Blocks() as demo:
             '''<h1 text-align="center">Infini-gram: An Engine for n-gram / ∞-gram Language Models with Trillion-Token Corpora</h1>
 
             <p style='font-size: 16px;'>This is an engine that processes n-gram / ∞-gram queries on a text corpus. 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://arxiv.org/abs/2401.17377">Infini-gram: Scaling Unbounded n-gram Language Models to a Trillion Tokens</a></p>
-            <p style='font-size: 16px;'>HF Paper Page: <a href="https://huggingface.co/papers/2401.17377">https://huggingface.co/papers/2401.17377</a></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://arxiv.org/abs/2401.17377">Infini-gram: Scaling Unbounded n-gram Language Models to a Trillion Tokens</a>. HF Paper Page: <a href="https://huggingface.co/papers/2401.17377">https://huggingface.co/papers/2401.17377</a></p>
             <p style='font-size: 16px;'><b>Note: We kindly ask you not to programmatically submit queries to the API at the moment. We will release a more stable API soon. Thank you :)</b></p>
             '''
         )
         with gr.Row():
             with gr.Column(scale=1):
                 corpus_desc = gr.Radio(choices=CORPUS_DESCS, label='Corpus', value=CORPUS_DESCS[0])
-            with gr.Column(scale=3):
-                query_desc = gr.Radio(
-                    choices=QUERY_DESCS, label='Query Type', value=QUERY_DESCS[0],
-                )
+                engine_desc = gr.Radio(choices=ENGINE_DESCS, label='Engine', value=ENGINE_DESCS[0])
 
-        with gr.Row(visible=True) as row_1:
-            with gr.Column():
-                gr.HTML('<h2>1. Count an n-gram</h2>')
-                gr.HTML('<p style="font-size: 16px;">This counts the number of times an n-gram appears in the corpus. If you submit an empty input, it will return the total number of tokens in the corpus.</p>')
-                gr.HTML('<p style="font-size: 16px;">Example query: <b>natural language processing</b> (the output is Cnt(natural language processing))</p>')
-                with gr.Row():
-                    with gr.Column(scale=1):
-                        count_input = gr.Textbox(placeholder='Enter a string (an n-gram) here', label='Query', interactive=True)
+            with gr.Column(scale=5):
+                with gr.Tab('1. Count an n-gram'):
+                    with gr.Column():
+                        gr.HTML('<h2>1. Count an n-gram</h2>')
+                        gr.HTML('<p style="font-size: 16px;">This counts the number of times an n-gram appears in the corpus. If you submit an empty input, it will return the total number of tokens in the corpus.</p>')
+                        gr.HTML('<p style="font-size: 16px;">Example query: <b>natural language processing</b> (the output is Cnt(natural language processing))</p>')
                         with gr.Row():
-                            count_clear = gr.ClearButton(value='Clear', variant='secondary', visible=True)
-                            count_submit = gr.Button(value='Submit', variant='primary', visible=True)
-                        count_output_tokens = gr.Textbox(label='Tokenized', lines=2, interactive=False)
-                    with gr.Column(scale=1):
-                        count_output = gr.Label(label='Count', num_top_classes=0)
+                            with gr.Column(scale=1):
+                                count_query = gr.Textbox(placeholder='Enter a string (an n-gram) here', label='Query', interactive=True)
+                                with gr.Row():
+                                    count_clear = gr.ClearButton(value='Clear', variant='secondary', visible=True)
+                                    count_submit = gr.Button(value='Submit', variant='primary', visible=True)
+                                count_latency = gr.Textbox(label='Latency (milliseconds)', interactive=False, lines=1)
+                                count_tokenized = gr.Textbox(label='Tokenized', lines=2, interactive=False)
+                            with gr.Column(scale=1):
+                                count_count = gr.Label(label='Count', num_top_classes=0)
+                    count_clear.add([count_query, count_latency, count_tokenized, count_count])
+                    count_submit.click(count, inputs=[corpus_desc, engine_desc, count_query], outputs=[count_latency, count_tokenized, count_count], api_name=False)
 
-        with gr.Row(visible=False) as row_2:
-            with gr.Column():
-                gr.HTML('<h2>2. Compute the probability of the last token in an n-gram</h2>')
-                gr.HTML('<p style="font-size: 16px;">This computes the n-gram probability of the last token conditioned on the previous tokens (i.e. (n-1)-gram)).</p>')
-                gr.HTML('<p style="font-size: 16px;">Example query: <b>natural language processing</b> (the output is P(processing | natural language), by counting the appearance of the 3-gram "natural language processing" and the 2-gram "natural language", and take the division between the two)</p>')
-                gr.HTML('<p style="font-size: 16px;">Note: The (n-1)-gram needs to exist in the corpus. If the (n-1)-gram is not found in the corpus, an error message will appear.</p>')
-                with gr.Row():
-                    with gr.Column(scale=1):
-                        ngram_input = gr.Textbox(placeholder='Enter a string (an n-gram) here', label='Query', interactive=True)
+                with gr.Tab('2. Prob of the last token'):
+                    with gr.Column():
+                        gr.HTML('<h2>2. Compute the probability of the last token in an n-gram</h2>')
+                        gr.HTML('<p style="font-size: 16px;">This computes the n-gram probability of the last token conditioned on the previous tokens (i.e. (n-1)-gram)).</p>')
+                        gr.HTML('<p style="font-size: 16px;">Example query: <b>natural language processing</b> (the output is P(processing | natural language), by counting the appearance of the 3-gram "natural language processing" and the 2-gram "natural language", and take the division between the two)</p>')
+                        gr.HTML('<p style="font-size: 16px;">Note: The (n-1)-gram needs to exist in the corpus. If the (n-1)-gram is not found in the corpus, an error message will appear.</p>')
                         with gr.Row():
-                            ngram_clear = gr.ClearButton(value='Clear', variant='secondary', visible=True)
-                            ngram_submit = gr.Button(value='Submit', variant='primary', visible=True)
-                        ngram_output_tokens = gr.Textbox(label='Tokenized', lines=2, interactive=False)
-                    with gr.Column(scale=1):
-                        ngram_output = gr.Label(label='Probability', num_top_classes=0)
+                            with gr.Column(scale=1):
+                                prob_query = gr.Textbox(placeholder='Enter a string (an n-gram) here', label='Query', interactive=True)
+                                with gr.Row():
+                                    prob_clear = gr.ClearButton(value='Clear', variant='secondary', visible=True)
+                                    prob_submit = gr.Button(value='Submit', variant='primary', visible=True)
+                                prob_latency = gr.Textbox(label='Latency (milliseconds)', interactive=False, lines=1)
+                                prob_tokenized = gr.Textbox(label='Tokenized', lines=2, interactive=False)
+                            with gr.Column(scale=1):
+                                prob_probability = gr.Label(label='Probability', num_top_classes=0)
+                    prob_clear.add([prob_query, prob_latency, prob_tokenized, prob_probability])
+                    prob_submit.click(prob, inputs=[corpus_desc, engine_desc, prob_query], outputs=[prob_latency, prob_tokenized, prob_probability], api_name=False)
 
-        with gr.Row(visible=False) as row_3:
-            with gr.Column():
-                gr.HTML('<h2>3. Compute the next-token distribution of an (n-1)-gram</h2>')
-                gr.HTML('<p style="font-size: 16px;">This is an extension of the Query 2: It interprets your input as the (n-1)-gram and gives you the full next-token distribution.</p>')
-                gr.HTML('<p style="font-size: 16px;">Example query: <b>natural language</b> (the output is P(* | natural language), for the top-10 tokens *)</p>')
-                gr.HTML(f'<p style="font-size: 16px;">Note: The (n-1)-gram needs to exist in the corpus. If the (n-1)-gram is not found in the corpus, an error message will appear. If the (n-1)-gram appears more than {MAX_CNT_FOR_NTD} times in the corpus, the result will be approximate.</p>')
-                with gr.Row():
-                    with gr.Column(scale=1):
-                        ntd_input = gr.Textbox(placeholder='Enter a string (an (n-1)-gram) here', label='Query', interactive=True)
+                with gr.Tab('3. Next-token distribution'):
+                    with gr.Column():
+                        gr.HTML('<h2>3. Compute the next-token distribution of an (n-1)-gram</h2>')
+                        gr.HTML('<p style="font-size: 16px;">This is an extension of the Query 2: It interprets your input as the (n-1)-gram and gives you the full next-token distribution.</p>')
+                        gr.HTML('<p style="font-size: 16px;">Example query: <b>natural language</b> (the output is P(* | natural language), for the top-10 tokens *)</p>')
+                        gr.HTML(f'<p style="font-size: 16px;">Note: The (n-1)-gram needs to exist in the corpus. If the (n-1)-gram is not found in the corpus, an error message will appear. If the (n-1)-gram appears more than {MAX_CNT_FOR_NTD} times in the corpus, the result will be approximate.</p>')
                         with gr.Row():
-                            ntd_clear = gr.ClearButton(value='Clear', variant='secondary', visible=True)
-                            ntd_submit = gr.Button(value='Submit', variant='primary', visible=True)
-                        ntd_output_tokens = gr.Textbox(label='Tokenized', lines=2, interactive=False)
-                    with gr.Column(scale=1):
-                        ntd_output = gr.Label(label='Distribution', num_top_classes=10)
+                            with gr.Column(scale=1):
+                                ntd_query = gr.Textbox(placeholder='Enter a string (an (n-1)-gram) here', label='Query', interactive=True)
+                                with gr.Row():
+                                    ntd_clear = gr.ClearButton(value='Clear', variant='secondary', visible=True)
+                                    ntd_submit = gr.Button(value='Submit', variant='primary', visible=True)
+                                ntd_latency = gr.Textbox(label='Latency (milliseconds)', interactive=False, lines=1)
+                                ntd_tokenized = gr.Textbox(label='Tokenized', lines=2, interactive=False)
+                            with gr.Column(scale=1):
+                                ntd_distribution = gr.Label(label='Distribution', num_top_classes=10)
+                    ntd_clear.add([ntd_query, ntd_latency, ntd_tokenized, ntd_distribution])
+                    ntd_submit.click(ntd, inputs=[corpus_desc, engine_desc, ntd_query], outputs=[ntd_latency, ntd_tokenized, ntd_distribution], api_name=False)
 
-        with gr.Row(visible=False) as row_4:
-            with gr.Column():
-                gr.HTML('<h2>4. Compute the ∞-gram probability of the last token</h2>')
-                gr.HTML('<p style="font-size: 16px;">This computes the ∞-gram probability of the last token conditioned on the previous tokens. Compared to Query 2 (which uses your entire input for n-gram modeling), here we take the longest suffix that we can find in the corpus.</p>')
-                gr.HTML('<p style="font-size: 16px;">Example query: <b>I love natural language processing</b> (the output is P(processing | natural language), because "natural language" appears in the corpus but "love natural language" doesn\'t; in this case the effective n = 3)</p>')
-                gr.HTML('<p style="font-size: 16px;">Note: It may be possible that the effective n = 1, in which case it reduces to the uni-gram probability of the last token.</p>')
-                with gr.Row():
-                    with gr.Column(scale=1):
-                        infgram_input = gr.Textbox(placeholder='Enter a string here', label='Query', interactive=True)
+                with gr.Tab('4. ∞-gram prob'):
+                    with gr.Column():
+                        gr.HTML('<h2>4. Compute the ∞-gram probability of the last token</h2>')
+                        gr.HTML('<p style="font-size: 16px;">This computes the ∞-gram probability of the last token conditioned on the previous tokens. Compared to Query 2 (which uses your entire input for n-gram modeling), here we take the longest suffix that we can find in the corpus.</p>')
+                        gr.HTML('<p style="font-size: 16px;">Example query: <b>I love natural language processing</b> (the output is P(processing | natural language), because "natural language" appears in the corpus but "love natural language" doesn\'t; in this case the effective n = 3)</p>')
+                        gr.HTML('<p style="font-size: 16px;">Note: It may be possible that the effective n = 1, in which case it reduces to the uni-gram probability of the last token.</p>')
                         with gr.Row():
-                            infgram_clear = gr.ClearButton(value='Clear', variant='secondary', visible=True)
-                            infgram_submit = gr.Button(value='Submit', variant='primary', visible=True)
-                        infgram_output_tokens = gr.Textbox(label='Tokenized', lines=2, interactive=False)
-                        infgram_longest_suffix = gr.Textbox(label='Longest Found Suffix', interactive=False)
-                    with gr.Column(scale=1):
-                        infgram_output = gr.Label(label='Probability', num_top_classes=0)
+                            with gr.Column(scale=1):
+                                infgram_prob_query = gr.Textbox(placeholder='Enter a string here', label='Query', interactive=True)
+                                with gr.Row():
+                                    infgram_prob_clear = gr.ClearButton(value='Clear', variant='secondary', visible=True)
+                                    infgram_prob_submit = gr.Button(value='Submit', variant='primary', visible=True)
+                                infgram_prob_latency = gr.Textbox(label='Latency (milliseconds)', interactive=False, lines=1)
+                                infgram_prob_tokenized = gr.Textbox(label='Tokenized', lines=2, interactive=False)
+                                infgram_prob_longest_suffix = gr.Textbox(label='Longest Found Suffix', interactive=False)
+                            with gr.Column(scale=1):
+                                infgram_prob_probability = gr.Label(label='Probability', num_top_classes=0)
+                    infgram_prob_clear.add([infgram_prob_query, infgram_prob_latency, infgram_prob_tokenized, infgram_prob_longest_suffix, infgram_prob_probability])
+                    infgram_prob_submit.click(infgram_prob, inputs=[corpus_desc, engine_desc, infgram_prob_query], outputs=[infgram_prob_latency, infgram_prob_tokenized, infgram_prob_longest_suffix, infgram_prob_probability], api_name=False)
 
-        with gr.Row(visible=False) as row_5:
-            with gr.Column():
-                gr.HTML('<h2>5. Compute the ∞-gram next-token distribution</h2>')
-                gr.HTML('<p style="font-size: 16px;">This is similar to Query 3, but with ∞-gram instead of n-gram.</p>')
-                gr.HTML('<p style="font-size: 16px;">Example query: <b>I love natural language</b> (the output is P(* | natural language), for the top-10 tokens *)</p>')
-                with gr.Row():
-                    with gr.Column(scale=1):
-                        infntd_input = gr.Textbox(placeholder='Enter a string here', label='Query', interactive=True)
+                with gr.Tab('5. ∞-gram next-token distribution'):
+                    with gr.Column():
+                        gr.HTML('<h2>5. Compute the ∞-gram next-token distribution</h2>')
+                        gr.HTML('<p style="font-size: 16px;">This is similar to Query 3, but with ∞-gram instead of n-gram.</p>')
+                        gr.HTML('<p style="font-size: 16px;">Example query: <b>I love natural language</b> (the output is P(* | natural language), for the top-10 tokens *)</p>')
                         with gr.Row():
-                            infntd_clear = gr.ClearButton(value='Clear', variant='secondary', visible=True)
-                            infntd_submit = gr.Button(value='Submit', variant='primary', visible=True)
-                        infntd_output_tokens = gr.Textbox(label='Tokenized', lines=2, interactive=False)
-                        infntd_longest_suffix = gr.Textbox(label='Longest Found Suffix', interactive=False)
-                    with gr.Column(scale=1):
-                        infntd_output = gr.Label(label='Distribution', num_top_classes=10)
+                            with gr.Column(scale=1):
+                                infgram_ntd_query = gr.Textbox(placeholder='Enter a string here', label='Query', interactive=True)
+                                with gr.Row():
+                                    infgram_ntd_clear = gr.ClearButton(value='Clear', variant='secondary', visible=True)
+                                    infgram_ntd_submit = gr.Button(value='Submit', variant='primary', visible=True)
+                                infgram_ntd_latency = gr.Textbox(label='Latency (milliseconds)', interactive=False, lines=1)
+                                infgram_ntd_tokenized = gr.Textbox(label='Tokenized', lines=2, interactive=False)
+                                infgram_ntd_longest_suffix = gr.Textbox(label='Longest Found Suffix', interactive=False)
+                            with gr.Column(scale=1):
+                                infgram_ntd_distribution = gr.Label(label='Distribution', num_top_classes=10)
+                    infgram_ntd_clear.add([infgram_ntd_query, infgram_ntd_latency, infgram_ntd_tokenized, infgram_ntd_longest_suffix, infgram_ntd_distribution])
+                    infgram_ntd_submit.click(infgram_ntd, inputs=[corpus_desc, engine_desc, infgram_ntd_query], outputs=[infgram_ntd_latency, infgram_ntd_tokenized, infgram_ntd_longest_suffix, infgram_ntd_distribution], api_name=False)
 
-        # with gr.Row(visible=False) as row_6:
-        #     with gr.Column():
-        #         gr.HTML(f'''<h2>6. Searching for document containing n-gram(s)</h2>
-        #                     <p style="font-size: 16px;">This displays a random document in the corpus that satisfies your query. You can simply enter an n-gram, in which case the document displayed would contain your n-gram. You can also connect multiple n-gram terms with the AND/OR operators, in the <a href="https://en.wikipedia.org/wiki/Conjunctive_normal_form">CNF format</a>, in which case the displayed document contains n-grams such that it satisfies this logical constraint.</p>
-        #                     <p style="font-size: 16px;">Example queries:</p>
-        #                     <ul style="font-size: 16px;">
-        #                         <li><b>natural language processing</b> (the displayed document would contain "natural language processing")</li>
-        #                         <li><b>natural language processing AND deep learning</b> (the displayed document would contain both "natural language processing" and "deep learning")</li>
-        #                         <li><b>natural language processing OR artificial intelligence AND deep learning OR machine learning</b> (the displayed document would contain at least one of "natural language processing" / "artificial intelligence", and also at least one of "deep learning" / "machine learning")</li>
-        #                     </ul>
-        #                     <p style="font-size: 16px;">If you want another random document, simply hit the Submit button again :)</p>
-        #                     <p style="font-size: 16px;">A few notes:</p>
-        #                     <ul style="font-size: 16px;">
-        #                         <li>When you write a query in CNF, note that <b>OR has higher precedence than AND</b> (which is contrary to conventions in boolean algebra).</li>
-        #                         <li>If the document is too long, it will be truncated to {MAX_OUTPUT_DOC_TOKENS} tokens.</li>
-        #                         <li>We can only include documents where all terms (or clauses) are separated by no more than {MAX_DIFF_TOKENS} tokens.</li>
-        #                         <li>If you query for two or more clauses, and a clause has more than {MAX_CLAUSE_FREQ_FAST_APPROX_PER_SHARD} matches (per shard), we will estimate the count from a random subset of all documents containing that clause. This might cause a zero count on conjuction of some simple n-grams (e.g., <b>birds AND oil</b>).</li>
-        #                         <li>The number of found documents may contain duplicates (e.g., if a document contains your query term twice, it may be counted twice).</li>
-        #                     </ul>
-        #                     <p style="font-size: 16px;">❗️WARNING: Corpus may contain problematic contents such as PII, toxicity, hate speech, and NSFW text. This tool is merely presenting selected text from the corpus, without any post-hoc safety filtering. It is NOT creating new text. This is a research prototype through which we can expose and examine existing problems with massive text corpora. Please use with caution. Don't be evil :)</p>
-        #                 ''')
-        #         with gr.Row():
-        #             with gr.Column(scale=1):
-        #                 ard_cnf_input = gr.Textbox(placeholder='Enter a query here', label='Query', interactive=True)
-        #                 with gr.Row():
-        #                     ard_cnf_clear = gr.ClearButton(value='Clear', variant='secondary', visible=True)
-        #                     ard_cnf_submit = gr.Button(value='Submit', variant='primary', visible=True)
-        #                 ard_cnf_output_tokens = gr.Textbox(label='Tokenized', lines=2, interactive=False)
-        #             with gr.Column(scale=1):
-        #                 ard_cnf_output_message = gr.Label(label='Message', num_top_classes=0)
-        #                 ard_cnf_output = gr.HighlightedText(label='Document', show_legend=False, color_map={"-": "red", "0": "green", "1": "cyan", "2": "blue", "3": "magenta"})
-
-        with gr.Row(visible=False) as row_6a:
-            with gr.Column():
-                gr.HTML(f'''<h2>6. Searching for documents containing n-gram(s)</h2>
-                            <p style="font-size: 16px;">This displays a few random documents in the corpus that satisfies your query. You can simply enter an n-gram, in which case the document displayed would contain your n-gram. You can also connect multiple n-gram terms with the AND/OR operators, in the <a href="https://en.wikipedia.org/wiki/Conjunctive_normal_form">CNF format</a>, in which case the displayed document contains n-grams such that it satisfies this logical constraint.</p>
-                            <p style="font-size: 16px;">Example queries:</p>
-                            <ul style="font-size: 16px;">
-                                <li><b>natural language processing</b> (the displayed document would contain "natural language processing")</li>
-                                <li><b>natural language processing AND deep learning</b> (the displayed document would contain both "natural language processing" and "deep learning")</li>
-                                <li><b>natural language processing OR artificial intelligence AND deep learning OR machine learning</b> (the displayed document would contain at least one of "natural language processing" / "artificial intelligence", and also at least one of "deep learning" / "machine learning")</li>
-                            </ul>
-                            <p style="font-size: 16px;">If you want another batch of random documents, simply hit the Submit button again :)</p>
-                            <p style="font-size: 16px;">A few notes:</p>
-                            <ul style="font-size: 16px;">
-                                <li>When you write a query in CNF, note that <b>OR has higher precedence than AND</b> (which is contrary to conventions in boolean algebra).</li>
-                                <li>If the document is too long, it will be truncated to {MAX_OUTPUT_DOC_TOKENS} tokens.</li>
-                                <li>We can only include documents where all terms (or clauses) are separated by no more than {MAX_DIFF_TOKENS} tokens.</li>
-                                <li>If you query for two or more clauses, and a clause has more than {MAX_CLAUSE_FREQ_FAST_APPROX_PER_SHARD} matches (per shard), we will estimate the count from a random subset of all documents containing that clause. This might cause a zero count on conjuction of some simple n-grams (e.g., <b>birds AND oil</b>).</li>
-                                <li>The number of found documents may contain duplicates (e.g., if a document contains your query term twice, it may be counted twice).</li>
-                            </ul>
-                            <p style="font-size: 16px;">❗️WARNING: Corpus may contain problematic contents such as PII, toxicity, hate speech, and NSFW text. This tool is merely presenting selected text from the corpus, without any post-hoc safety filtering. It is NOT creating new text. This is a research prototype through which we can expose and examine existing problems with massive text corpora. Please use with caution. Don't be evil :)</p>
-                        ''')
-                with gr.Row():
-                    with gr.Column(scale=1):
-                        ard_cnf_multi_input = gr.Textbox(placeholder='Enter a query here', label='Query', interactive=True)
-                        ard_cnf_multi_maxnum = gr.Slider(minimum=1, maximum=10, value=1, step=1, label='Number of documents to Display')
+                with gr.Tab('6. Search documents'):
+                    with gr.Column():
+                        gr.HTML(f'''<h2>6. Search for documents containing n-gram(s)</h2>
+                                    <p style="font-size: 16px;">This displays a few random documents in the corpus that satisfies your query. You can simply enter an n-gram, in which case the document displayed would contain your n-gram. You can also connect multiple n-gram terms with the AND/OR operators, in the <a href="https://en.wikipedia.org/wiki/Conjunctive_normal_form">CNF format</a>, in which case the displayed document contains n-grams such that it satisfies this logical constraint.</p>
+                                    <p style="font-size: 16px;">Example queries:</p>
+                                    <ul style="font-size: 16px;">
+                                        <li><b>natural language processing</b> (the displayed document would contain "natural language processing")</li>
+                                        <li><b>natural language processing AND deep learning</b> (the displayed document would contain both "natural language processing" and "deep learning")</li>
+                                        <li><b>natural language processing OR artificial intelligence AND deep learning OR machine learning</b> (the displayed document would contain at least one of "natural language processing" / "artificial intelligence", and also at least one of "deep learning" / "machine learning")</li>
+                                    </ul>
+                                    <p style="font-size: 16px;">If you want another batch of random documents, simply hit the Submit button again :)</p>
+                                    <p style="font-size: 16px;">A few notes:</p>
+                                    <ul style="font-size: 16px;">
+                                        <li>When you write a query in CNF, note that <b>OR has higher precedence than AND</b> (which is contrary to conventions in boolean algebra).</li>
+                                        <li>If the document is too long, it will be truncated to {MAX_OUTPUT_DOC_TOKENS} tokens.</li>
+                                        <li>We can only include documents where all terms (or clauses) are separated by no more than {MAX_DIFF_TOKENS} tokens.</li>
+                                        <li>If you query for two or more clauses, and a clause has more than {MAX_CLAUSE_FREQ_FAST_APPROX_PER_SHARD} matches (per shard), we will estimate the count from a random subset of all documents containing that clause. This might cause a zero count on conjuction of some simple n-grams (e.g., <b>birds AND oil</b>).</li>
+                                        <li>The number of found documents may contain duplicates (e.g., if a document contains your query term twice, it may be counted twice).</li>
+                                    </ul>
+                                    <p style="font-size: 16px;">❗️WARNING: Corpus may contain problematic contents such as PII, toxicity, hate speech, and NSFW text. This tool is merely presenting selected text from the corpus, without any post-hoc safety filtering. It is NOT creating new text. This is a research prototype through which we can expose and examine existing problems with massive text corpora. Please use with caution. Don't be evil :)</p>
+                                ''')
                         with gr.Row():
-                            ard_cnf_multi_clear = gr.ClearButton(value='Clear', variant='secondary', visible=True)
-                            ard_cnf_multi_submit = gr.Button(value='Submit', variant='primary', visible=True)
-                        ard_cnf_multi_output_tokens = gr.Textbox(label='Tokenized', lines=2, interactive=False)
-                    with gr.Column(scale=1):
-                        ard_cnf_multi_output_message = gr.Label(label='Message', num_top_classes=0)
-                        with gr.Tab(label='1'):
-                            ard_cnf_multi_output_0 = gr.HighlightedText(label='Document', show_legend=False, color_map={"-": "red", "0": "green", "1": "cyan", "2": "blue", "3": "magenta"})
-                        with gr.Tab(label='2'):
-                            ard_cnf_multi_output_1 = gr.HighlightedText(label='Document', show_legend=False, color_map={"-": "red", "0": "green", "1": "cyan", "2": "blue", "3": "magenta"})
-                        with gr.Tab(label='3'):
-                            ard_cnf_multi_output_2 = gr.HighlightedText(label='Document', show_legend=False, color_map={"-": "red", "0": "green", "1": "cyan", "2": "blue", "3": "magenta"})
-                        with gr.Tab(label='4'):
-                            ard_cnf_multi_output_3 = gr.HighlightedText(label='Document', show_legend=False, color_map={"-": "red", "0": "green", "1": "cyan", "2": "blue", "3": "magenta"})
-                        with gr.Tab(label='5'):
-                            ard_cnf_multi_output_4 = gr.HighlightedText(label='Document', show_legend=False, color_map={"-": "red", "0": "green", "1": "cyan", "2": "blue", "3": "magenta"})
-                        with gr.Tab(label='6'):
-                            ard_cnf_multi_output_5 = gr.HighlightedText(label='Document', show_legend=False, color_map={"-": "red", "0": "green", "1": "cyan", "2": "blue", "3": "magenta"})
-                        with gr.Tab(label='7'):
-                            ard_cnf_multi_output_6 = gr.HighlightedText(label='Document', show_legend=False, color_map={"-": "red", "0": "green", "1": "cyan", "2": "blue", "3": "magenta"})
-                        with gr.Tab(label='8'):
-                            ard_cnf_multi_output_7 = gr.HighlightedText(label='Document', show_legend=False, color_map={"-": "red", "0": "green", "1": "cyan", "2": "blue", "3": "magenta"})
-                        with gr.Tab(label='9'):
-                            ard_cnf_multi_output_8 = gr.HighlightedText(label='Document', show_legend=False, color_map={"-": "red", "0": "green", "1": "cyan", "2": "blue", "3": "magenta"})
-                        with gr.Tab(label='10'):
-                            ard_cnf_multi_output_9 = gr.HighlightedText(label='Document', show_legend=False, color_map={"-": "red", "0": "green", "1": "cyan", "2": "blue", "3": "magenta"})
+                            with gr.Column(scale=2):
+                                search_docs_query = gr.Textbox(placeholder='Enter a query here', label='Query', interactive=True)
+                                search_docs_maxnum = gr.Slider(minimum=1, maximum=10, value=1, step=1, label='Number of documents to Display')
+                                with gr.Row():
+                                    search_docs_clear = gr.ClearButton(value='Clear', variant='secondary', visible=True)
+                                    search_docs_submit = gr.Button(value='Submit', variant='primary', visible=True)
+                                search_docs_latency = gr.Textbox(label='Latency (milliseconds)', interactive=False, lines=1)
+                                search_docs_tokenized = gr.Textbox(label='Tokenized', lines=2, interactive=False)
+                            with gr.Column(scale=3):
+                                search_docs_message = gr.Label(label='Message', num_top_classes=0)
+                                with gr.Tab(label='1'):
+                                    search_docs_output_0 = gr.HighlightedText(label='Document', show_legend=False, color_map={"-": "red", "0": "green", "1": "cyan", "2": "blue", "3": "magenta"})
+                                with gr.Tab(label='2'):
+                                    search_docs_output_1 = gr.HighlightedText(label='Document', show_legend=False, color_map={"-": "red", "0": "green", "1": "cyan", "2": "blue", "3": "magenta"})
+                                with gr.Tab(label='3'):
+                                    search_docs_output_2 = gr.HighlightedText(label='Document', show_legend=False, color_map={"-": "red", "0": "green", "1": "cyan", "2": "blue", "3": "magenta"})
+                                with gr.Tab(label='4'):
+                                    search_docs_output_3 = gr.HighlightedText(label='Document', show_legend=False, color_map={"-": "red", "0": "green", "1": "cyan", "2": "blue", "3": "magenta"})
+                                with gr.Tab(label='5'):
+                                    search_docs_output_4 = gr.HighlightedText(label='Document', show_legend=False, color_map={"-": "red", "0": "green", "1": "cyan", "2": "blue", "3": "magenta"})
+                                with gr.Tab(label='6'):
+                                    search_docs_output_5 = gr.HighlightedText(label='Document', show_legend=False, color_map={"-": "red", "0": "green", "1": "cyan", "2": "blue", "3": "magenta"})
+                                with gr.Tab(label='7'):
+                                    search_docs_output_6 = gr.HighlightedText(label='Document', show_legend=False, color_map={"-": "red", "0": "green", "1": "cyan", "2": "blue", "3": "magenta"})
+                                with gr.Tab(label='8'):
+                                    search_docs_output_7 = gr.HighlightedText(label='Document', show_legend=False, color_map={"-": "red", "0": "green", "1": "cyan", "2": "blue", "3": "magenta"})
+                                with gr.Tab(label='9'):
+                                    search_docs_output_8 = gr.HighlightedText(label='Document', show_legend=False, color_map={"-": "red", "0": "green", "1": "cyan", "2": "blue", "3": "magenta"})
+                                with gr.Tab(label='10'):
+                                    search_docs_output_9 = gr.HighlightedText(label='Document', show_legend=False, color_map={"-": "red", "0": "green", "1": "cyan", "2": "blue", "3": "magenta"})
+                    search_docs_clear.add([search_docs_query, search_docs_latency, search_docs_tokenized, search_docs_message, search_docs_output_0, search_docs_output_1, search_docs_output_2, search_docs_output_3, search_docs_output_4, search_docs_output_5, search_docs_output_6, search_docs_output_7, search_docs_output_8, search_docs_output_9])
+                    search_docs_submit.click(search_docs, inputs=[corpus_desc, engine_desc, search_docs_query, search_docs_maxnum], outputs=[search_docs_latency, search_docs_tokenized, search_docs_message, search_docs_output_0, search_docs_output_1, search_docs_output_2, search_docs_output_3, search_docs_output_4, search_docs_output_5, search_docs_output_6, search_docs_output_7, search_docs_output_8, search_docs_output_9], api_name=False)
 
-        with gr.Row(visible=False) as row_7:
-            with gr.Column():
-                gr.HTML('<h2>7. Analyze an (AI-generated) document using ∞-gram</h2>')
-                gr.HTML('<p style="font-size: 16px;">This analyzes the document you entered using the ∞-gram. Each token is highlighted where (1) the color represents its ∞-gram probability (red is 0.0, blue is 1.0), and (2) the alpha represents the effective n (higher alpha means higher n).</p>')
-                gr.HTML('<p style="font-size: 16px;">If you hover over a token, the tokens preceding it are each highlighted where (1) the color represents the n-gram probability of your selected token, with the n-gram starting from that highlighted token (red is 0.0, blue is 1.0), and (2) the alpha represents the count of the (n-1)-gram starting from that highlighted token (and up to but excluding your selected token) (higher alpha means higher count).</p>')
-                with gr.Row():
-                    with gr.Column(scale=1):
-                        doc_analysis_input = gr.Textbox(placeholder='Enter a document here', label='Query', interactive=True, lines=10)
+                with gr.Tab('7. Analyze an (AI-generated) document using ∞-gram', visible=False):
+                    with gr.Column():
+                        gr.HTML('<h2>7. Analyze an (AI-generated) document using ∞-gram</h2>')
+                        gr.HTML('<p style="font-size: 16px;">This analyzes the document you entered using the ∞-gram. Each token is highlighted where (1) the color represents its ∞-gram probability (red is 0.0, blue is 1.0), and (2) the alpha represents the effective n (higher alpha means higher n).</p>')
+                        gr.HTML('<p style="font-size: 16px;">If you hover over a token, the tokens preceding it are each highlighted where (1) the color represents the n-gram probability of your selected token, with the n-gram starting from that highlighted token (red is 0.0, blue is 1.0), and (2) the alpha represents the count of the (n-1)-gram starting from that highlighted token (and up to but excluding your selected token) (higher alpha means higher count).</p>')
                         with gr.Row():
-                            doc_analysis_clear = gr.ClearButton(value='Clear', variant='secondary', visible=True)
-                            doc_analysis_submit = gr.Button(value='Submit', variant='primary', visible=True)
-                    with gr.Column(scale=1):
-                        doc_analysis_output = gr.HTML(value='', label='Analysis')
+                            with gr.Column(scale=1):
+                                analyze_document_query = gr.Textbox(placeholder='Enter a document here', label='Query', interactive=True, lines=10)
+                                with gr.Row():
+                                    analyze_document_clear = gr.ClearButton(value='Clear', variant='secondary', visible=True)
+                                    analyze_document_submit = gr.Button(value='Submit', variant='primary', visible=True)
+                            with gr.Column(scale=1):
+                                analyze_document_html = gr.HTML(value='', label='Analysis')
+                    analyze_document_clear.add([analyze_document_query, analyze_document_html])
+                    analyze_document_submit.click(analyze_document, inputs=[corpus_desc, engine_desc, analyze_document_query], outputs=[analyze_document_html], api_name=False)
 
         with gr.Row():
             gr.Markdown('''
@@ -307,65 +250,25 @@ If you find this tool useful, please kindly cite our paper:
 ```
 ''')
 
-    count_clear.add([count_input, count_output, count_output_tokens])
-    ngram_clear.add([ngram_input, ngram_output, ngram_output_tokens])
-    ntd_clear.add([ntd_input, ntd_output, ntd_output_tokens])
-    infgram_clear.add([infgram_input, infgram_output, infgram_output_tokens])
-    infntd_clear.add([infntd_input, infntd_output, infntd_output_tokens, infntd_longest_suffix])
-    # ard_cnf_clear.add([ard_cnf_input, ard_cnf_output, ard_cnf_output_tokens, ard_cnf_output_message])
-    ard_cnf_multi_clear.add([ard_cnf_multi_input, ard_cnf_multi_output_tokens, ard_cnf_multi_output_message, ard_cnf_multi_output_0, ard_cnf_multi_output_1, ard_cnf_multi_output_2, ard_cnf_multi_output_3, ard_cnf_multi_output_4, ard_cnf_multi_output_5, ard_cnf_multi_output_6, ard_cnf_multi_output_7, ard_cnf_multi_output_8, ard_cnf_multi_output_9])
-    doc_analysis_clear.add([doc_analysis_input, doc_analysis_output])
-
-    count_submit.click(process_2, inputs=[corpus_desc, query_desc, count_input], outputs=[count_output, count_output_tokens], api_name=False)
-    ngram_submit.click(process_2, inputs=[corpus_desc, query_desc, ngram_input], outputs=[ngram_output, ngram_output_tokens], api_name=False)
-    ntd_submit.click(process_2, inputs=[corpus_desc, query_desc, ntd_input], outputs=[ntd_output, ntd_output_tokens], api_name=False)
-    infgram_submit.click(process_3, inputs=[corpus_desc, query_desc, infgram_input], outputs=[infgram_output, infgram_output_tokens, infgram_longest_suffix], api_name=False)
-    infntd_submit.click(process_3, inputs=[corpus_desc, query_desc, infntd_input], outputs=[infntd_output, infntd_output_tokens, infntd_longest_suffix], api_name=False)
-    # ard_cnf_submit.click(process, inputs=[corpus_desc, query_desc, ard_cnf_input], outputs=[ard_cnf_output, ard_cnf_output_tokens, ard_cnf_output_message], api_name=False)
-    ard_cnf_multi_submit.click(process_ard_cnf_multi, inputs=[corpus_desc, query_desc, ard_cnf_multi_input, ard_cnf_multi_maxnum], outputs=[ard_cnf_multi_output_message, ard_cnf_multi_output_tokens, ard_cnf_multi_output_0, ard_cnf_multi_output_1, ard_cnf_multi_output_2, ard_cnf_multi_output_3, ard_cnf_multi_output_4, ard_cnf_multi_output_5, ard_cnf_multi_output_6, ard_cnf_multi_output_7, ard_cnf_multi_output_8, ard_cnf_multi_output_9], api_name=False)
-    doc_analysis_submit.click(process_1, inputs=[corpus_desc, query_desc, doc_analysis_input], outputs=[doc_analysis_output], api_name=False)
-
-    def update_query_desc(selection):
-        return {
-            row_1: gr.Row(visible=(selection == QUERY_DESC_BY_TYPE['count'])),
-            row_2: gr.Row(visible=(selection == QUERY_DESC_BY_TYPE['compute_prob'])),
-            row_3: gr.Row(visible=(selection == QUERY_DESC_BY_TYPE['get_next_token_distribution_approx'])),
-            row_4: gr.Row(visible=(selection == QUERY_DESC_BY_TYPE['compute_infgram_prob'])),
-            row_5: gr.Row(visible=(selection == QUERY_DESC_BY_TYPE['get_infgram_next_token_distribution_approx'])),
-            # row_6: gr.Row(visible=(selection == QUERY_DESC_BY_TYPE['get_a_random_document_from_cnf_query_fast_approx'])),
-            row_6a: gr.Row(visible=(selection == QUERY_DESC_BY_TYPE['get_random_documents_from_cnf_query_fast_approx'])),
-            # row_7: gr.Row(visible=(selection == QUERY_DESC_BY_TYPE['analyze_document'])),
-        }
-    query_desc.change(fn=update_query_desc, inputs=query_desc, outputs=[
-        row_1,
-        row_2,
-        row_3,
-        row_4,
-        row_5,
-        # row_6,
-        row_6a,
-        # row_7,
-    ])
-
 for d in demo.dependencies:
     d['api_name'] = False
 for d in demo.config['dependencies']:
     d['api_name'] = False
-if debug:
-    print(demo.dependencies)
-    print(demo.config['dependencies'])
+# if DEBUG:
+#     print(demo.dependencies)
+#     print(demo.config['dependencies'])
 
 demo.queue(
-    default_concurrency_limit=default_concurrency_limit,
-    max_size=max_size,
+    default_concurrency_limit=DEFAULT_CONCURRENCY_LIMIT,
+    max_size=MAX_SIZE,
     api_open=False,
 ).launch(
-    max_threads=max_threads,
-    debug=debug,
+    max_threads=MAX_THREADS,
+    debug=DEBUG,
     show_api=False,
 )
 
-for d in gr.context.Context.root_block.dependencies:
-    d['api_name'] = False
-if debug:
-    print(gr.context.Context.root_block.dependencies)
+# for d in gr.context.Context.root_block.dependencies:
+#     d['api_name'] = False
+# if DEBUG:
+#     print(gr.context.Context.root_block.dependencies)

constants.py

@@ -1,27 +1,24 @@
 import os
 
+# options
 CORPUS_BY_DESC = {
     'RedPajama (LLaMA tokenizer), 1.4T tokens': 'v3_rpj_llama_c4',
     'Pile-val (LLaMA tokenizer), 390M tokens': 'v3_pileval_llama',
-    'Pile-val (GPT-2 tokenizer), 380M tokens': 'v3_pileval',
+    'Pile-val (GPT-2 tokenizer), 380M tokens': 'v3_pileval_gpt2',
 }
 CORPUS_DESCS = list(CORPUS_BY_DESC.keys())
-QUERY_TYPE_BY_DESC = {
-    '1. Count an n-gram': 'count',
-    '2. Compute the probability of the last token in an n-gram': 'compute_prob',
-    '3. Compute the next-token distribution of an (n-1)-gram': 'get_next_token_distribution_approx',
-    '4. Compute the ∞-gram probability of the last token': 'compute_infgram_prob',
-    '5. Compute the ∞-gram next-token distribution': 'get_infgram_next_token_distribution_approx',
-    # '6. Searching for document containing n-gram(s)': 'get_a_random_document_from_cnf_query_fast_approx',
-    '6. Searching for documents containing n-gram(s)': 'get_random_documents_from_cnf_query_fast_approx',
-    # '7. Analyze an (AI-generated) document using ∞-gram': 'analyze_document',
+ENGINE_BY_DESC = {
+    'Python': 'python',
+    'C++ (🚀 Fast, Experimental)': 'c++',
 }
-QUERY_DESC_BY_TYPE = {v: k for k, v in QUERY_TYPE_BY_DESC.items()}
-QUERY_DESCS = list(QUERY_TYPE_BY_DESC.keys())
+ENGINE_DESCS = list(ENGINE_BY_DESC.keys())
+ENGINES = list(ENGINE_BY_DESC.values())
 
+# engine
 MAX_QUERY_CHARS = int(os.environ.get('MAX_QUERY_CHARS', 1000))
 MAX_INPUT_DOC_TOKENS = int(os.environ.get('MAX_INPUT_DOC_TOKENS', 1000))
 MAX_OUTPUT_DOC_TOKENS = int(os.environ.get('MAX_OUTPUT_DOC_TOKENS', 5000))
+MAX_OUTPUT_NUM_DOCS = int(os.environ.get('MAX_OUTPUT_NUM_DOCS', 10))
 MAX_CNT_FOR_NTD = int(os.environ.get('MAX_CNT_FOR_NTD', 1000))
 MAX_CLAUSE_FREQ = int(os.environ.get('MAX_CLAUSE_FREQ', 10000))
 MAX_CLAUSE_FREQ_FAST = int(os.environ.get('MAX_CLAUSE_FREQ_FAST', 1000000))
@@ -31,4 +28,14 @@ MAX_DIFF_BYTES = 2 * MAX_DIFF_TOKENS
 MAX_CLAUSES_IN_CNF = int(os.environ.get('MAX_CLAUSES_IN_CNF', 4))
 MAX_TERMS_IN_DISJ_CLAUSE = int(os.environ.get('MAX_TERMS_IN_DISJ_CLAUSE', 4))
 
+# HF demo
+API_IPADDR = os.environ.get('API_IPADDR', None)
+DEFAULT_CONCURRENCY_LIMIT = os.environ.get('DEFAULT_CONCURRENCY_LIMIT', 10)
+MAX_SIZE = os.environ.get('MAX_SIZE', 100)
+MAX_THREADS = os.environ.get('MAX_THREADS', 40)
+DEBUG = (os.environ.get('DEBUG', 'False') != 'False')
 MIN_QUERY_INTERVAL_SECONDS = int(os.environ.get('MIN_QUERY_INTERVAL_SECONDS', 5))
+
+# C++ engine
+CPP_PORT = int(os.environ.get('CPP_PORT', 3786))
+SOCKET_OUT_BUFFER_SIZE = 65536