Sync changes

README.md

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

app.py

@@ -4,20 +4,16 @@ import json
 import requests
 from constants import *
 
-def process(query_type, corpus_desc, engine_desc, query, maxnum, request: gr.Request):
+def process(query_type, index_desc, **kwargs):
     timestamp = datetime.datetime.now().strftime('%Y%m%d-%H%M%S')
-    corpus = CORPUS_BY_DESC[corpus_desc]
-    engine = ENGINE_BY_DESC[engine_desc]
+    index = INDEX_BY_DESC[index_desc]
     data = {
         'source': 'hf' if not DEBUG else 'hf-dev',
         'timestamp': timestamp,
         'query_type': query_type,
-        'corpus': corpus,
-        'engine': engine,
-        'query': query,
+        'index': index,
     }
-    if maxnum is not None:
-        data['maxnum'] = maxnum
+    data.update(kwargs)
     print(json.dumps(data))
     if API_URL is None:
         raise ValueError(f'API_URL envvar is not set!')
@@ -63,8 +59,8 @@ def format_doc(doc):
     formatted += doc['spans']
     return formatted
 
-def count(corpus_desc, engine_desc, query, request: gr.Request):
-    result = process('count', corpus_desc, engine_desc, query, None, request)
+def count(index_desc, query):
+    result = process('count', index_desc, query=query)
     latency = '' if 'latency' not in result else f'{result["latency"]:.3f}'
     tokenization_info = format_tokenization_info(result)
     if 'error' in result:
@@ -73,8 +69,8 @@ def count(corpus_desc, engine_desc, query, request: gr.Request):
         count = f'{result["count"]:,}'
     return latency, tokenization_info, count
 
-def prob(corpus_desc, engine_desc, query, request: gr.Request):
-    result = process('prob', corpus_desc, engine_desc, query, None, request)
+def prob(index_desc, query):
+    result = process('prob', index_desc, query=query)
     latency = '' if 'latency' not in result else f'{result["latency"]:.3f}'
     tokenization_info = format_tokenization_info(result)
     if 'error' in result:
@@ -85,8 +81,8 @@ def prob(corpus_desc, engine_desc, query, request: gr.Request):
         prob = f'{result["prob"]:.4f} ({result["cont_cnt"]:,} / {result["prompt_cnt"]:,})'
     return latency, tokenization_info, prob
 
-def ntd(corpus_desc, engine_desc, query, request: gr.Request):
-    result = process('ntd', corpus_desc, engine_desc, query, None, request)
+def ntd(index_desc, query):
+    result = process('ntd', index_desc, query=query)
     latency = '' if 'latency' not in result else f'{result["latency"]:.3f}'
     tokenization_info = format_tokenization_info(result)
     if 'error' in result:
@@ -100,8 +96,8 @@ def ntd(corpus_desc, engine_desc, query, request: gr.Request):
             ntd = '(n-1)-gram is not found in the corpus'
     return latency, tokenization_info, ntd
 
-def infgram_prob(corpus_desc, engine_desc, query, request: gr.Request):
-    result = process('infgram_prob', corpus_desc, engine_desc, query, None, request)
+def infgram_prob(index_desc, query):
+    result = process('infgram_prob', index_desc, query=query)
     latency = '' if 'latency' not in result else f'{result["latency"]:.3f}'
     tokenization_info = format_tokenization_info(result)
     if 'error' in result:
@@ -112,8 +108,8 @@ def infgram_prob(corpus_desc, engine_desc, query, request: gr.Request):
         prob = f'{result["prob"]:.4f} ({result["cont_cnt"]:,} / {result["prompt_cnt"]:,})'
     return latency, tokenization_info, longest_suffix, prob
 
-def infgram_ntd(corpus_desc, engine_desc, query, request: gr.Request):
-    result = process('infgram_ntd', corpus_desc, engine_desc, query, None, request)
+def infgram_ntd(index_desc, query):
+    result = process('infgram_ntd', index_desc, query=query)
     latency = '' if 'latency' not in result else f'{result["latency"]:.3f}'
     tokenization_info = format_tokenization_info(result)
     if 'error' in result:
@@ -127,25 +123,21 @@ def infgram_ntd(corpus_desc, engine_desc, query, request: gr.Request):
             ntd[f'{r["token"]} ({r["cont_cnt"]} / {result["prompt_cnt"]})'] = r['prob']
     return latency, tokenization_info, longest_suffix, ntd
 
-def search_docs(corpus_desc, engine_desc, query, maxnum, request: gr.Request):
-    result = process('search_docs', corpus_desc, engine_desc, query, maxnum, request)
+def search_docs(index_desc, query, maxnum):
+    result = process('search_docs', index_desc, query=query, maxnum=maxnum)
     latency = '' if 'latency' not in result else f'{result["latency"]:.3f}'
     tokenization_info = format_tokenization_info(result)
     if 'error' in result:
         message = result['error']
-        docs = [[] for _ in range(10)]
+        docs = [[] for _ in range(MAXNUM)]
     else:
         message = result['message']
         docs = result['documents']
         docs = [format_doc(doc) for doc in docs]
     docs = docs[:maxnum]
-    while len(docs) < 10:
+    while len(docs) < MAXNUM:
         docs.append([])
-    return latency, tokenization_info, message, docs[0], docs[1], docs[2], docs[3], docs[4], docs[5], docs[6], docs[7], docs[8], docs[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', '')
+    return tuple([latency, tokenization_info, message] + docs)
 
 with gr.Blocks() as demo:
     with gr.Column():
@@ -160,10 +152,9 @@ with gr.Blocks() as demo:
         )
         with gr.Row():
             with gr.Column(scale=1):
-                corpus_desc = gr.Radio(choices=CORPUS_DESCS, label='Corpus', value=CORPUS_DESCS[0])
-                engine_desc = gr.Radio(choices=ENGINE_DESCS, label='Engine', value=ENGINE_DESCS[0])
+                index_desc = gr.Radio(choices=INDEX_DESCS, label='Corpus', value=INDEX_DESCS[0])
 
-            with gr.Column(scale=5):
+            with gr.Column(scale=7):
                 with gr.Tab('1. Count an n-gram'):
                     with gr.Column():
                         gr.HTML('<h2>1. Count an n-gram</h2>')
@@ -180,7 +171,7 @@ with gr.Blocks() as demo:
                             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)
+                    count_submit.click(count, inputs=[index_desc, count_query], outputs=[count_latency, count_tokenized, count_count], api_name=False)
 
                 with gr.Tab('2. Prob of the last token'):
                     with gr.Column():
@@ -199,14 +190,14 @@ with gr.Blocks() as demo:
                             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)
+                    prob_submit.click(prob, inputs=[index_desc, prob_query], outputs=[prob_latency, prob_tokenized, prob_probability], api_name=False)
 
                 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>')
+                        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_SUPPORT} times in the corpus, the result will be approximate.</p>')
                         with gr.Row():
                             with gr.Column(scale=1):
                                 ntd_query = gr.Textbox(placeholder='Enter a string (an (n-1)-gram) here', label='Query', interactive=True)
@@ -218,7 +209,7 @@ with gr.Blocks() as demo:
                             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)
+                    ntd_submit.click(ntd, inputs=[index_desc, ntd_query], outputs=[ntd_latency, ntd_tokenized, ntd_distribution], api_name=False)
 
                 with gr.Tab('4. ∞-gram prob'):
                     with gr.Column():
@@ -238,7 +229,7 @@ with gr.Blocks() as demo:
                             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)
+                    infgram_prob_submit.click(infgram_prob, inputs=[index_desc, infgram_prob_query], outputs=[infgram_prob_latency, infgram_prob_tokenized, infgram_prob_longest_suffix, infgram_prob_probability], api_name=False)
 
                 with gr.Tab('5. ∞-gram next-token distribution'):
                     with gr.Column():
@@ -257,7 +248,7 @@ with gr.Blocks() as demo:
                             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)
+                    infgram_ntd_submit.click(infgram_ntd, inputs=[index_desc, infgram_ntd_query], outputs=[infgram_ntd_latency, infgram_ntd_tokenized, infgram_ntd_longest_suffix, infgram_ntd_distribution], api_name=False)
 
                 with gr.Tab('6. Search documents'):
                     with gr.Column():
@@ -272,18 +263,19 @@ with gr.Blocks() as demo:
                                     <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>If the document is too long, it will be truncated to {MAX_DISP_LEN} tokens.</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>
+                                        <li>A CNF query may contain up to {MAX_CLAUSES_PER_CNF} clauses, and each clause may contain up to {MAX_TERMS_PER_CLAUSE} terms.</li>
                                         <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_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>
+                                        <li>If you query for two or more clauses, and a clause has more than {MAX_CLAUSE_FREQ} matches, 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>
                                     </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=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')
+                                search_docs_maxnum = gr.Slider(minimum=1, maximum=MAXNUM, 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)
@@ -291,44 +283,12 @@ with gr.Blocks() as demo:
                                 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.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():
-                            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)
+                                search_docs_outputs = []
+                                for i in range(MAXNUM):
+                                    with gr.Tab(label=str(i+1)):
+                                        search_docs_outputs.append(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_outputs)
+                    search_docs_submit.click(search_docs, inputs=[index_desc, search_docs_query, search_docs_maxnum], outputs=[search_docs_latency, search_docs_tokenized, search_docs_message] + search_docs_outputs, api_name=False)
 
         with gr.Row():
             gr.Markdown('''
@@ -343,14 +303,6 @@ If you find this tool useful, please kindly cite our paper:
 ```
 ''')
 
-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'])
-
 demo.queue(
     default_concurrency_limit=DEFAULT_CONCURRENCY_LIMIT,
     max_size=MAX_SIZE,
@@ -360,8 +312,3 @@ demo.queue(
     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)

constants.py

@@ -1,8 +1,8 @@
 import os
 
 # options
-CORPUS_BY_DESC = {
-    'Dolma (3.1T tokens)': 'v4_dolma-v1_6_llama',
+INDEX_BY_DESC = {
+    'Dolma-v1.6 (3.1T tokens)': 'v4_dolma-v1_6_llama',
     'RedPajama (1.4T tokens)': 'v4_rpj_llama_s4',
     'Pile-train (380B tokens)': 'v4_piletrain_llama',
     'C4-train (200B tokens)': 'v4_c4train_llama',
@@ -20,25 +20,17 @@ CORPUS_BY_DESC = {
     # 'Dolma-v1.6-cc_en_middle (650B tokens): 'v4_dolma-v1_6-cc_en_middle_llama',
     # 'Dolma-v1.6-cc_en_tail (970B tokens): 'v4_dolma-v1_6-cc_en_tail_llama',
 }
-CORPUS_DESCS = list(CORPUS_BY_DESC.keys())
-ENGINE_BY_DESC = {
-    'C++ (🚀🚀 Fast)': 'c++',
-    'Python': 'python',
-}
-ENGINE_DESCS = list(ENGINE_BY_DESC.keys())
-ENGINES = list(ENGINE_BY_DESC.values())
+INDEX_DESCS = list(INDEX_BY_DESC.keys())
 
-# engine
+# API limits
 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)) # This number is also hard-coded in app.py
-MAX_CNT_FOR_NTD = int(os.environ.get('MAX_CNT_FOR_NTD', 1000))
-MAX_CLAUSE_FREQ_PER_SHARD = int(os.environ.get('MAX_CLAUSE_FREQ_FAST_APPROX_PER_SHARD', 50000))
+MAX_CLAUSES_PER_CNF = int(os.environ.get('MAX_CLAUSES_PER_CNF', 4))
+MAX_TERMS_PER_CLAUSE = int(os.environ.get('MAX_TERMS_PER_CLAUSE', 4))
+MAX_SUPPORT = int(os.environ.get('MAX_SUPPORT', 1000))
+MAX_CLAUSE_FREQ = int(os.environ.get('MAX_CLAUSE_FREQ', 50000))
 MAX_DIFF_TOKENS = int(os.environ.get('MAX_DIFF_TOKENS', 100))
-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))
+MAXNUM = int(os.environ.get('MAXNUM', 10))
+MAX_DISP_LEN = int(os.environ.get('MAX_DISP_LEN', 5000))
 
 # HF demo
 API_URL = os.environ.get('API_URL', None)