Improve description and logging

app.py

@@ -1,4 +1,5 @@
 import gradio as gr
+import datetime
 import json
 import os
 import requests
@@ -8,12 +9,13 @@ 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)
+debug = (os.environ.get('debug', 'False') != 'False')
 
 def process(corpus_desc, query_desc, query):
     corpus = CORPUS_BY_DESC[corpus_desc]
     query_type = QUERY_TYPE_BY_DESC[query_desc]
-    print(json.dumps({'corpus': corpus, 'query_type': query_type, 'query': query}))
+    timestamp = datetime.datetime.now().strftime('%Y%m%d-%H%M%S')
+    print(json.dumps({'timestamp': timestamp, 'corpus': corpus, 'query_type': query_type, 'query': query}))
     data = {
         'corpus': corpus,
         'query_type': query_type,
@@ -26,7 +28,8 @@ def process(corpus_desc, query_desc, query):
         result = response.json()
     else:
         raise ValueError(f'Invalid response: {response.status_code}')
-    # print(result)
+    if debug:
+        print(result)
     return result
 
 with gr.Blocks() as demo:
@@ -35,12 +38,13 @@ 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 documented in our paper: <a href="">Infini-gram: Scaling Unbounded n-gram Language Models to a Trillion Tokens</a></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=4):
+            with gr.Column(scale=3):
                 query_desc = gr.Radio(
                     choices=QUERY_DESCS, label='Query Type', value=QUERY_DESCS[0],
                 )
@@ -171,6 +175,14 @@ with gr.Blocks() as demo:
                     with gr.Column(scale=1):
                         doc_analysis_output = gr.HTML(value='', label='Analysis')
 
+        with gr.Row():
+            gr.Markdown('''
+If you find this tool useful, please kindly cite our paper:
+```
+(coming soon)
+```
+''')
+
     count_clear.add([count_input, count_output, count_output_tokens])
     ngram_clear.add([ngram_input, ngram_output, ngram_output_tokens])
     a_ntd_clear.add([a_ntd_input, a_ntd_output, a_ntd_output_tokens])

constants.py

@@ -1,9 +1,9 @@
 import os
 
 CORPUS_BY_DESC = {
-    'RedPajama (LLaMA tokenizer)': 'v3_rpj_llama_c4',
-    'Pile-val (LLaMA tokenizer)': 'v3_pileval_llama',
-    'Pile-val (GPT-2 tokenizer)': 'v3_pileval',
+    'RedPajama (LLaMA tokenizer), 1.4T tokens': 'v3_rpj_llama_c4',
+    'Pile-val (LLaMA tokenizer), 790M tokens': 'v3_pileval_llama',
+    'Pile-val (GPT-2 tokenizer) 770M tokens': 'v3_pileval',
 }
 CORPUS_DESCS = list(CORPUS_BY_DESC.keys())
 QUERY_TYPE_BY_DESC = {