Update app.py

app.py

@@ -1,458 +1,528 @@
 import gradio as gr
-import spaces
-import torch
-from diffusers import AutoencoderKL, TCDScheduler
-from diffusers.models.model_loading_utils import load_state_dict
-from gradio_imageslider import ImageSlider
-from huggingface_hub import hf_hub_download
-
-from controlnet_union import ControlNetModel_Union
-from pipeline_fill_sd_xl import StableDiffusionXLFillPipeline
-
-from PIL import Image, ImageDraw
-import numpy as np
-
-config_file = hf_hub_download(
-    "xinsir/controlnet-union-sdxl-1.0",
-    filename="config_promax.json",
-)
-
-config = ControlNetModel_Union.load_config(config_file)
-controlnet_model = ControlNetModel_Union.from_config(config)
-model_file = hf_hub_download(
-    "xinsir/controlnet-union-sdxl-1.0",
-    filename="diffusion_pytorch_model_promax.safetensors",
-)
-state_dict = load_state_dict(model_file)
-model, _, _, _, _ = ControlNetModel_Union._load_pretrained_model(
-    controlnet_model, state_dict, model_file, "xinsir/controlnet-union-sdxl-1.0"
-)
-model.to(device="cuda", dtype=torch.float16)
-
-vae = AutoencoderKL.from_pretrained(
-    "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
-).to("cuda")
-
-pipe = StableDiffusionXLFillPipeline.from_pretrained(
-    "SG161222/RealVisXL_V5.0_Lightning",
-    torch_dtype=torch.float16,
-    vae=vae,
-    controlnet=model,
-    variant="fp16",
-).to("cuda")
-
-pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
-
-
-def can_expand(source_width, source_height, target_width, target_height, alignment):
-    """Checks if the image can be expanded based on the alignment."""
-    if alignment in ("Left", "Right") and source_width >= target_width:
-        return False
-    if alignment in ("Top", "Bottom") and source_height >= target_height:
-        return False
-    return True
-
-def prepare_image_and_mask(image, width, height, overlap_percentage, resize_option, custom_resize_percentage, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom):
-    target_size = (width, height)
-
-    # Calculate the scaling factor to fit the image within the target size
-    scale_factor = min(target_size[0] / image.width, target_size[1] / image.height)
-    new_width = int(image.width * scale_factor)
-    new_height = int(image.height * scale_factor)
-    
-    # Resize the source image to fit within target size
-    source = image.resize((new_width, new_height), Image.LANCZOS)
-
-    # Apply resize option using percentages
-    if resize_option == "Full":
-        resize_percentage = 100
-    elif resize_option == "50%":
-        resize_percentage = 50
-    elif resize_option == "33%":
-        resize_percentage = 33
-    elif resize_option == "25%":
-        resize_percentage = 25
-    else:  # Custom
-        resize_percentage = custom_resize_percentage
-
-    # Calculate new dimensions based on percentage
-    resize_factor = resize_percentage / 100
-    new_width = int(source.width * resize_factor)
-    new_height = int(source.height * resize_factor)
-
-    # Ensure minimum size of 64 pixels
-    new_width = max(new_width, 64)
-    new_height = max(new_height, 64)
-
-    # Resize the image
-    source = source.resize((new_width, new_height), Image.LANCZOS)
-
-    # Calculate the overlap in pixels based on the percentage
-    overlap_x = int(new_width * (overlap_percentage / 100))
-    overlap_y = int(new_height * (overlap_percentage / 100))
-
-    # Ensure minimum overlap of 1 pixel
-    overlap_x = max(overlap_x, 1)
-    overlap_y = max(overlap_y, 1)
-
-    # Calculate margins based on alignment
-    if alignment == "Middle":
-        margin_x = (target_size[0] - new_width) // 2
-        margin_y = (target_size[1] - new_height) // 2
-    elif alignment == "Left":
-        margin_x = 0
-        margin_y = (target_size[1] - new_height) // 2
-    elif alignment == "Right":
-        margin_x = target_size[0] - new_width
-        margin_y = (target_size[1] - new_height) // 2
-    elif alignment == "Top":
-        margin_x = (target_size[0] - new_width) // 2
-        margin_y = 0
-    elif alignment == "Bottom":
-        margin_x = (target_size[0] - new_width) // 2
-        margin_y = target_size[1] - new_height
-
-    # Adjust margins to eliminate gaps
-    margin_x = max(0, min(margin_x, target_size[0] - new_width))
-    margin_y = max(0, min(margin_y, target_size[1] - new_height))
-
-    # Create a new background image and paste the resized source image
-    background = Image.new('RGB', target_size, (255, 255, 255))
-    background.paste(source, (margin_x, margin_y))
-
-    # Create the mask
-    mask = Image.new('L', target_size, 255)
-    mask_draw = ImageDraw.Draw(mask)
-
-    # Calculate overlap areas
-    white_gaps_patch = 2
-
-    left_overlap = margin_x + overlap_x if overlap_left else margin_x + white_gaps_patch
-    right_overlap = margin_x + new_width - overlap_x if overlap_right else margin_x + new_width - white_gaps_patch
-    top_overlap = margin_y + overlap_y if overlap_top else margin_y + white_gaps_patch
-    bottom_overlap = margin_y + new_height - overlap_y if overlap_bottom else margin_y + new_height - white_gaps_patch
-    
-    if alignment == "Left":
-        left_overlap = margin_x + overlap_x if overlap_left else margin_x
-    elif alignment == "Right":
-        right_overlap = margin_x + new_width - overlap_x if overlap_right else margin_x + new_width
-    elif alignment == "Top":
-        top_overlap = margin_y + overlap_y if overlap_top else margin_y
-    elif alignment == "Bottom":
-        bottom_overlap = margin_y + new_height - overlap_y if overlap_bottom else margin_y + new_height
-
-
-    # Draw the mask
-    mask_draw.rectangle([
-        (left_overlap, top_overlap),
-        (right_overlap, bottom_overlap)
-    ], fill=0)
-
-    return background, mask
-
-def preview_image_and_mask(image, width, height, overlap_percentage, resize_option, custom_resize_percentage, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom):
-    background, mask = prepare_image_and_mask(image, width, height, overlap_percentage, resize_option, custom_resize_percentage, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom)
-    
-    # Create a preview image showing the mask
-    preview = background.copy().convert('RGBA')
-    
-    # Create a semi-transparent red overlay
-    red_overlay = Image.new('RGBA', background.size, (255, 0, 0, 64))  # Reduced alpha to 64 (25% opacity)
-    
-    # Convert black pixels in the mask to semi-transparent red
-    red_mask = Image.new('RGBA', background.size, (0, 0, 0, 0))
-    red_mask.paste(red_overlay, (0, 0), mask)
-    
-    # Overlay the red mask on the background
-    preview = Image.alpha_composite(preview, red_mask)
-    
-    return preview
-
-@spaces.GPU(duration=24)
-def infer(image, width, height, overlap_percentage, num_inference_steps, resize_option, custom_resize_percentage, prompt_input, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom):
-    background, mask = prepare_image_and_mask(image, width, height, overlap_percentage, resize_option, custom_resize_percentage, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom)
-    
-    if not can_expand(background.width, background.height, width, height, alignment):
-        alignment = "Middle"
-
-    cnet_image = background.copy()
-    cnet_image.paste(0, (0, 0), mask)
-
-    final_prompt = f"{prompt_input} , high quality, 4k" if prompt_input else "high quality, 4k"
-
-    # Use with torch.autocast to ensure consistent dtype
-    with torch.autocast(device_type="cuda", dtype=torch.float16):
-        (
-            prompt_embeds,
-            negative_prompt_embeds,
-            pooled_prompt_embeds,
-            negative_pooled_prompt_embeds,
-        ) = pipe.encode_prompt(final_prompt, "cuda", True)
-
-        for image in pipe(
-            prompt_embeds=prompt_embeds,
-            negative_prompt_embeds=negative_prompt_embeds,
-            pooled_prompt_embeds=pooled_prompt_embeds,
-            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
-            image=cnet_image,
-            num_inference_steps=num_inference_steps
-        ):
-            yield cnet_image, image
-
-    image = image.convert("RGBA")
-    cnet_image.paste(image, (0, 0), mask)
-
-    yield background, cnet_image
-
-def clear_result():
-    """Clears the result ImageSlider."""
-    return gr.update(value=None)
-
-def preload_presets(target_ratio, ui_width, ui_height):
-    """Updates the width and height sliders based on the selected aspect ratio."""
-    if target_ratio == "9:16":
-        changed_width = 720
-        changed_height = 1280
-        return changed_width, changed_height, gr.update()
-    elif target_ratio == "16:9":
-        changed_width = 1280
-        changed_height = 720
-        return changed_width, changed_height, gr.update()
-    elif target_ratio == "1:1":
-        changed_width = 1024
-        changed_height = 1024
-        return changed_width, changed_height, gr.update()
-    elif target_ratio == "Custom":
-        return ui_width, ui_height, gr.update(open=True)
-
-def select_the_right_preset(user_width, user_height):
-    if user_width == 720 and user_height == 1280:
-        return "9:16"
-    elif user_width == 1280 and user_height == 720:
-        return "16:9"
-    elif user_width == 1024 and user_height == 1024:
-        return "1:1"
+import datetime
+import json
+import random
+import requests
+from constants import *
+
+def process(query_type, index_desc, **kwargs):
+    timestamp = datetime.datetime.now().strftime('%Y%m%d-%H%M%S')
+    index = INDEX_BY_DESC[index_desc]
+    data = {
+        'source': 'hf' if not DEBUG else 'hf-dev',
+        'timestamp': timestamp,
+        'query_type': query_type,
+        'index': index,
+    }
+    data.update(kwargs)
+    print(json.dumps(data))
+    if API_URL is None:
+        raise ValueError(f'API_URL envvar is not set!')
+    try:
+        response = requests.post(API_URL, json=data, timeout=10)
+    except requests.exceptions.Timeout:
+        raise ValueError('Web request timed out. Please try again later.')
+    except requests.exceptions.RequestException as e:
+        raise ValueError(f'Web request error: {e}')
+    if response.status_code == 200:
+        result = response.json()
     else:
-        return "Custom"
-
-def toggle_custom_resize_slider(resize_option):
-    return gr.update(visible=(resize_option == "Custom"))
-
-def update_history(new_image, history):
-    """Updates the history gallery with the new image."""
-    if history is None:
-        history = []
-    history.insert(0, new_image)
-    return history
-
-css = """
-.gradio-container {
-    width: 1200px !important;
-}
-"""
-
-# Define the title HTML string
-title = """<h1 align="center">Diffusers Image Outpaint</h1>
-<div align="center">Drop an image you would like to extend, pick your expected ratio and hit Generate.</div>
-<div style="display: flex; justify-content: center; align-items: center; text-align: center;">
-    <p style="display: flex;gap: 6px;">
-        <a href="https://github.com/huggingface/diffusers">
-            Built with Diffusers
-        </a>
-    </p>
-</div>
-"""
-
-with gr.Blocks(css=css) as demo:
+        raise ValueError(f'HTTP error {response.status_code}: {response.json()}')
+    if DEBUG:
+        print(result)
+    return result
+
+def format_tokenization_info(result):
+    if not ('token_ids' in result and 'tokens' in result):
+        return ''
+    token_ids = result['token_ids']
+    tokens = result['tokens']
+    if type(token_ids) == list and all([type(token_id) == int for token_id in token_ids]):
+        output = '[' + " ".join(['"' + token.replace('Ġ', ' ') + '"' for token in tokens]) + '] ' + str(token_ids)
+    else:
+        ttt = []
+        for token_idss, tokenss in zip(token_ids, tokens):
+            tt = []
+            for token_ids, tokens in zip(token_idss, tokenss):
+                t = '[' + " ".join(['"' + token.replace('Ġ', ' ') + '"' for token in tokens]) + '] ' + str(token_ids)
+                tt.append(t)
+            tt = '\n'.join(tt)
+            ttt.append(tt)
+        output = '\n\n'.join(ttt)
+    return output
+def format_doc_metadata(doc):
+    formatted = f'Document #{doc["doc_ix"]}\n'
+    if doc['doc_len'] == doc['disp_len']:
+        formatted += f'Length: {doc["doc_len"]} tokens\n'
+    else:
+        formatted += f'Length: {doc["doc_len"]} tokens ({doc["disp_len"]} tokens displayed)\n'
+    metadata = doc['metadata'].strip("\n")
+    formatted += f'Metadata: {metadata}'
+    return formatted
+
+def count(index_desc, query, max_clause_freq, max_diff_tokens):
+    if ' AND ' in query or ' OR ' in query: # CNF query
+        result = process('count', index_desc, query=query, max_clause_freq=max_clause_freq, max_diff_tokens=max_diff_tokens)
+    else: # simple 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:
+        count = result['error']
+    else:
+        count = f'{result["count"]:,}'
+    return latency, tokenization_info, count
+
+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:
+        prob = result['error']
+    elif result['prompt_cnt'] == 0:
+        prob = '(n-1)-gram is not found in the corpus'
+    else:
+        prob = f'{result["prob"]:.4f} ({result["cont_cnt"]:,} / {result["prompt_cnt"]:,})'
+    return latency, tokenization_info, prob
+
+def ntd(index_desc, query, max_support):
+    result = process('ntd', index_desc, query=query, max_support=max_support)
+    latency = '' if 'latency' not in result else f'{result["latency"]:.3f}'
+    tokenization_info = format_tokenization_info(result)
+    if 'error' in result:
+        ntd = result['error']
+    else:
+        result_by_token_id = result['result_by_token_id']
+        ntd = {}
+        for token_id, r in result_by_token_id.items():
+            ntd[f'{r["token"]} ({r["cont_cnt"]} / {result["prompt_cnt"]})'] = r['prob']
+        if ntd == {}:
+            ntd = '(n-1)-gram is not found in the corpus'
+    return latency, tokenization_info, ntd
+
+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:
+        longest_suffix = ''
+        prob = result['error']
+    else:
+        longest_suffix = result['longest_suffix']
+        prob = f'{result["prob"]:.4f} ({result["cont_cnt"]:,} / {result["prompt_cnt"]:,})'
+    return latency, tokenization_info, longest_suffix, prob
+
+def infgram_ntd(index_desc, query, max_support):
+    result = process('infgram_ntd', index_desc, query=query, max_support=max_support)
+    latency = '' if 'latency' not in result else f'{result["latency"]:.3f}'
+    tokenization_info = format_tokenization_info(result)
+    if 'error' in result:
+        longest_suffix = ''
+        ntd = result['error']
+    else:
+        longest_suffix = result['longest_suffix']
+        result_by_token_id = result['result_by_token_id']
+        ntd = {}
+        for token_id, r in result_by_token_id.items():
+            ntd[f'{r["token"]} ({r["cont_cnt"]} / {result["prompt_cnt"]})'] = r['prob']
+    return latency, tokenization_info, longest_suffix, ntd
+
+def search_docs(index_desc, query, maxnum, max_disp_len, max_clause_freq, max_diff_tokens):
+    if ' AND ' in query or ' OR ' in query: # CNF query
+        result = process('search_docs', index_desc, query=query, maxnum=maxnum, max_disp_len=max_disp_len, max_clause_freq=max_clause_freq, max_diff_tokens=max_diff_tokens)
+    else: # simple query
+        result = process('search_docs', index_desc, query=query, maxnum=maxnum, max_disp_len=max_disp_len)
+    latency = '' if 'latency' not in result else f'{result["latency"]:.3f}'
+    tokenization_info = format_tokenization_info(result)
+    if 'error' in result:
+        message = result['error']
+        metadatas = ['' for _ in range(MAXNUM)]
+        docs = [[] for _ in range(MAXNUM)]
+    else:
+        message = result['message']
+        metadatas = [format_doc_metadata(doc) for doc in result['documents']]
+        docs = [doc['spans'] for doc in result['documents']]
+    metadatas = metadatas[:maxnum]
+    docs = docs[:maxnum]
+    while len(metadatas) < MAXNUM:
+        metadatas.append('')
+    while len(docs) < MAXNUM:
+        docs.append([])
+    return tuple([latency, tokenization_info, message] + metadatas + docs)
+
+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:
+        message = find_result['error']
+        idx = gr.Number(minimum=0, maximum=0, step=1, value=0, interactive=False)
+        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']
+        cnt_retrievable = sum([len(ptrs) for ptrs in ptrs_by_shard])
+        if find_result["approx"]:
+            message = f'Approximately {find_result["cnt"]} occurrences found, of which {cnt_retrievable} are retrievable'
+        else:
+            message = f'{find_result["cnt"]} occurrences found'
+    else: # simple query
+        message = f'{find_result["cnt"]} occurrences found'
+        cnt_retrievable = find_result['cnt']
+    if cnt_retrievable == 0:
+        idx = gr.Number(minimum=0, maximum=0, step=1, value=0, interactive=False)
+        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, state)
+    idx = gr.Number(minimum=0, maximum=cnt_retrievable-1, step=1, value=idx, interactive=True)
+    return latency, tokenization_info, message, idx, metadata, doc, state
+
+def clear_search_docs_new(state):
+    state = None
+    idx = gr.Number(minimum=0, maximum=0, step=1, value=0, interactive=False)
+    return idx, state
+
+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
+    if find_result['type'] == 'cnf':
+        ptrs_by_shard = find_result['ptrs_by_shard']
+        cnt_by_shard = [len(ptrs) for ptrs in ptrs_by_shard]
+        s = 0
+        while idx >= cnt_by_shard[s]:
+            idx -= cnt_by_shard[s]
+            s += 1
+        ptr = ptrs_by_shard[s][idx]
+        result = process('get_doc_by_ptr', index_desc, s=s, ptr=ptr, max_disp_len=max_disp_len, query_ids=find_result['token_ids'])
+    else: # simple query
+        segment_by_shard = find_result['segment_by_shard']
+        cnt_by_shard = [end - start for (start, end) in segment_by_shard]
+        s = 0
+        while idx >= cnt_by_shard[s]:
+            idx -= cnt_by_shard[s]
+            s += 1
+        rank = segment_by_shard[s][0] + idx
+        result = process('get_doc_by_rank', index_desc, s=s, rank=rank, max_disp_len=max_disp_len, query_ids=find_result['token_ids'])
+    if 'error' in result:
+        metadata = result['error']
+        doc = []
+        return metadata, doc
+    metadata = format_doc_metadata(result)
+    doc = result['spans']
+    return metadata, doc
+
+with gr.Blocks() as demo:
     with gr.Column():
-        gr.HTML(title)
-
+        gr.HTML(
+            '''<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 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>
+            '''
+        )
         with gr.Row():
-            with gr.Column():
-                input_image = gr.Image(
-                    type="pil",
-                    label="Input Image"
-                )
-
-                with gr.Row():
-                    with gr.Column(scale=2):
-                        prompt_input = gr.Textbox(label="Prompt (Optional)")
-                    with gr.Column(scale=1):
-                        run_button = gr.Button("Generate")
-
-                with gr.Row():
-                    target_ratio = gr.Radio(
-                        label="Expected Ratio",
-                        choices=["9:16", "16:9", "1:1", "Custom"],
-                        value="9:16",
-                        scale=2
-                    )
-                    
-                    alignment_dropdown = gr.Dropdown(
-                        choices=["Middle", "Left", "Right", "Top", "Bottom"],
-                        value="Middle",
-                        label="Alignment"
-                    )
+            with gr.Column(scale=1, min_width=240):
+                index_desc = gr.Radio(choices=INDEX_DESCS, label='Corpus', value=INDEX_DESCS[0])
 
-                with gr.Accordion(label="Advanced settings", open=False) as settings_panel:
+            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>')
+                        with gr.Accordion(label='Click to view instructions', open=False):
+                            gr.HTML(f'''<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. You can also make more complex queries by connecting multiple n-gram terms with the AND/OR operators, in the <a href="https://en.wikipedia.org/wiki/Conjunctive_normal_form">CNF format</a>.</p>
+                                        <br>
+                                        <p style="font-size: 16px;">Example queries:</p>
+                                        <ul style="font-size: 16px;">
+                                            <li><b>natural language processing</b> (the output is number of occurrences of "natural language processing")</li>
+                                            <li><b>natural language processing AND deep learning</b> (the output is the number of co-occurrences of "natural language processing" and "deep learning")</li>
+                                            <li><b>natural language processing OR artificial intelligence AND deep learning OR machine learning</b> (the output is the number of co-occurrences of [one of "natural language processing" / "artificial intelligence"] and [one of "deep learning" / "machine learning"])</li>
+                                        </ul>
+                                        <br>
+                                        <p style="font-size: 16px;">Notes on CNF queries:</p>
+                                        <ul style="font-size: 16px;">
+                                            <li>A CNF query may contain up to {MAX_CLAUSES_PER_CNF} clauses, and each clause may contain up to {MAX_TERMS_PER_CLAUSE} n-gram 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>In AND queries, we can only examine co-occurrences where adjacent clauses are separated by no more than {max_diff_tokens} tokens. This value can be adjusted within range [1, {MAX_DIFF_TOKENS}] in "Advanced options".</li>
+                                            <li>In AND queries, if a clause has more than {max_clause_freq} matches, we will estimate the count by examining a random subset of {max_clause_freq} occurrences of clause. This value can be adjusted within range [1, {MAX_CLAUSE_FREQ}] in "Advanced options".</li>
+                                            <li>The above subsampling mechanism might cause a zero count on co-occurrences of some simple n-grams (e.g., <b>birds AND oil</b>).</li>
+                                        </ul>
+                                    ''')
                         with gr.Row():
-                            width_slider = gr.Slider(
-                                label="Target Width",
-                                minimum=720,
-                                maximum=1536,
-                                step=8,
-                                value=720,  # Set a default value
-                            )
-                            height_slider = gr.Slider(
-                                label="Target Height",
-                                minimum=720,
-                                maximum=1536,
-                                step=8,
-                                value=1280,  # Set a default value
-                            )
-                        
-                        num_inference_steps = gr.Slider(label="Steps", minimum=4, maximum=12, step=1, value=8)
-                        with gr.Group():
-                            overlap_percentage = gr.Slider(
-                                label="Mask overlap (%)",
-                                minimum=1,
-                                maximum=50,
-                                value=10,
-                                step=1
-                            )
-                            with gr.Row():
-                                overlap_top = gr.Checkbox(label="Overlap Top", value=True)
-                                overlap_right = gr.Checkbox(label="Overlap Right", value=True)
-                            with gr.Row():
-                                overlap_left = gr.Checkbox(label="Overlap Left", value=True)
-                                overlap_bottom = gr.Checkbox(label="Overlap Bottom", value=True)
+                            with gr.Column(scale=1):
+                                count_query = gr.Textbox(placeholder='Enter a string (an n-gram) here', label='Query', interactive=True)
+                                with gr.Accordion(label='Advanced options', open=False):
+                                    with gr.Row():
+                                        count_max_clause_freq = gr.Slider(minimum=1, maximum=MAX_CLAUSE_FREQ, value=max_clause_freq, step=1, label='max_clause_freq')
+                                        count_max_diff_tokens = gr.Slider(minimum=1, maximum=MAX_DIFF_TOKENS, value=max_diff_tokens, step=1, label='max_diff_tokens')
+                                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=1, 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=[index_desc, count_query, count_max_clause_freq, count_max_diff_tokens], outputs=[count_latency, count_tokenized, count_count], api_name=False)
+
+                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>')
+                        with gr.Accordion(label='Click to view instructions', open=False):
+                            gr.HTML(f'''<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>
+                                        <br>
+                                        <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>
+                                        <br>
+                                        <p style="font-size: 16px;">Notes:</p>
+                                        <ul style="font-size: 16px;">
+                                            <li>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.</li>
+                                        </ul>
+                                    ''')
                         with gr.Row():
-                            resize_option = gr.Radio(
-                                label="Resize input image",
-                                choices=["Full", "50%", "33%", "25%", "Custom"],
-                                value="Full"
-                            )
-                            custom_resize_percentage = gr.Slider(
-                                label="Custom resize (%)",
-                                minimum=1,
-                                maximum=100,
-                                step=1,
-                                value=50,
-                                visible=False
-                            )
-                        
-                        with gr.Column():
-                            preview_button = gr.Button("Preview alignment and mask")
-                            
-                            
-                gr.Examples(
-                    examples=[
-                        ["./examples/example_2.jpg", 1440, 810, "Left"],
-                        ["./examples/example_3.jpg", 1024, 1024, "Top"],
-                        ["./examples/example_3.jpg", 1024, 1024, "Bottom"],
-                    ],
-                    inputs=[input_image, width_slider, height_slider, alignment_dropdown],
-                )
-
-                
-
-            with gr.Column():
-                result = ImageSlider(
-                    interactive=False,
-                    label="Generated Image",
-                )
-                use_as_input_button = gr.Button("Use as Input Image", visible=False)
-
-                history_gallery = gr.Gallery(label="History", columns=6, object_fit="contain", interactive=False)
-                preview_image = gr.Image(label="Preview")
-
-        
-
-    def use_output_as_input(output_image):
-        """Sets the generated output as the new input image."""
-        return gr.update(value=output_image[1])
-
-    use_as_input_button.click(
-        fn=use_output_as_input,
-        inputs=[result],
-        outputs=[input_image]
-    )
-    
-    target_ratio.change(
-        fn=preload_presets,
-        inputs=[target_ratio, width_slider, height_slider],
-        outputs=[width_slider, height_slider, settings_panel],
-        queue=False
-    )
-
-    width_slider.change(
-        fn=select_the_right_preset,
-        inputs=[width_slider, height_slider],
-        outputs=[target_ratio],
-        queue=False
-    )
-
-    height_slider.change(
-        fn=select_the_right_preset,
-        inputs=[width_slider, height_slider],
-        outputs=[target_ratio],
-        queue=False
-    )
-
-    resize_option.change(
-        fn=toggle_custom_resize_slider,
-        inputs=[resize_option],
-        outputs=[custom_resize_percentage],
-        queue=False
-    )
-    
-    run_button.click(  # Clear the result
-        fn=clear_result,
-        inputs=None,
-        outputs=result,
-    ).then(  # Generate the new image
-        fn=infer,
-        inputs=[input_image, width_slider, height_slider, overlap_percentage, num_inference_steps,
-                resize_option, custom_resize_percentage, prompt_input, alignment_dropdown,
-                overlap_left, overlap_right, overlap_top, overlap_bottom],
-        outputs=result,
-    ).then(  # Update the history gallery
-        fn=lambda x, history: update_history(x[1], history),
-        inputs=[result, history_gallery],
-        outputs=history_gallery,
-    ).then(  # Show the "Use as Input Image" button
-        fn=lambda: gr.update(visible=True),
-        inputs=None,
-        outputs=use_as_input_button,
-    )
-
-    prompt_input.submit(  # Clear the result
-        fn=clear_result,
-        inputs=None,
-        outputs=result,
-    ).then(  # Generate the new image
-        fn=infer,
-        inputs=[input_image, width_slider, height_slider, overlap_percentage, num_inference_steps,
-                resize_option, custom_resize_percentage, prompt_input, alignment_dropdown,
-                overlap_left, overlap_right, overlap_top, overlap_bottom],
-        outputs=result,
-    ).then(  # Update the history gallery
-        fn=lambda x, history: update_history(x[1], history),
-        inputs=[result, history_gallery],
-        outputs=history_gallery,
-    ).then(  # Show the "Use as Input Image" button
-        fn=lambda: gr.update(visible=True),
-        inputs=None,
-        outputs=use_as_input_button,
-    )
+                            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=1, 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=[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>')
+                        with gr.Accordion(label='Click to view instructions', open=False):
+                            gr.HTML(f'''<p style="font-size: 16px;">This is an extension of the Query Type 2: It interprets your input as the (n-1)-gram and gives you the full next-token distribution.</p>
+                                        <br>
+                                        <p style="font-size: 16px;">Example query: <b>natural language</b> (the output is P(* | natural language), for the top-10 tokens *)</p>
+                                        <br>
+                                        <p style="font-size: 16px;">Notes:</p>
+                                        <ul style="font-size: 16px;">
+                                            <li>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.</li>
+                                            <li>If the (n-1)-gram appears more than {max_support} times in the corpus, the result will be approximate: we will estimate the distribution by examining a subset of {max_support} occurrences of the (n-1)-gram. This value can be adjusted within range [1, {MAX_SUPPORT}] in "Advanced options".</li>
+                                        </ul>
+                                    ''')
 
-    preview_button.click(
-        fn=preview_image_and_mask,
-        inputs=[input_image, width_slider, height_slider, overlap_percentage, resize_option, custom_resize_percentage, alignment_dropdown,
-                overlap_left, overlap_right, overlap_top, overlap_bottom],
-        outputs=preview_image,
-        queue=False
-    )
+                        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)
+                                with gr.Accordion(label='Advanced options', open=False):
+                                    ntd_max_support = gr.Slider(minimum=1, maximum=MAX_SUPPORT, value=MAX_SUPPORT, step=1, label='max_support')
+                                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=1, 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=[index_desc, ntd_query, ntd_max_support], outputs=[ntd_latency, ntd_tokenized, ntd_distribution], api_name=False)
+
+                with gr.Tab('4. ∞-gram prob'):
+                    with gr.Column():
+                        gr.HTML('<h2>4. Compute the ∞-gram probability of the last token</h2>')
+                        with gr.Accordion(label='Click to view instructions', open=False):
+                            gr.HTML(f'''<p style="font-size: 16px;">This computes the ∞-gram probability of the last token conditioned on the previous tokens. Compared to Query Type 2 (which uses your entire input for n-gram modeling), here we take the longest suffix that we can find in the corpus.</p>
+                                        <br>
+                                        <p style="font-size: 16px;">Example query: <b>I love natural language processing</b> (if "natural language" appears in the corpus but "love natural language" doesn't, the output is P(processing | natural language); in this case the effective n = 3)</p>
+                                        <br>
+                                        <p style="font-size: 16px;">Notes:</p>
+                                        <ul style="font-size: 16px;">
+                                            <li>It may be possible that the effective n = 1, i.e. longest found suffix is empty, in which case it reduces to the uni-gram probability of the last token.</li>
+                                        </ul>
+                                    ''')
+                        with gr.Row():
+                            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=1, 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=[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():
+                        gr.HTML('<h2>5. Compute the ∞-gram next-token distribution</h2>')
+                        with gr.Accordion(label='Click to view instructions', open=False):
+                            gr.HTML(f'''<p style="font-size: 16px;">This is similar to Query Type 3, but with ∞-gram instead of n-gram.</p>
+                                        <br>
+                                        <p style="font-size: 16px;">Example query: <b>I love natural language</b> (if "natural language" appears in the corpus but "love natural language" doesn't, the output is P(* | natural language), for the top-10 tokens *)</p>
+                                    ''')
+                        with gr.Row():
+                            with gr.Column(scale=1):
+                                infgram_ntd_query = gr.Textbox(placeholder='Enter a string here', label='Query', interactive=True)
+                                with gr.Accordion(label='Advanced options', open=False):
+                                    infgram_ntd_max_support = gr.Slider(minimum=1, maximum=MAX_SUPPORT, value=MAX_SUPPORT, step=1, label='max_support')
+                                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=1, 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=[index_desc, infgram_ntd_query, infgram_ntd_max_support], outputs=[infgram_ntd_latency, infgram_ntd_tokenized, infgram_ntd_longest_suffix, infgram_ntd_distribution], api_name=False)
+
+                with gr.Tab('6. Search documents', visible=False):
+                    with gr.Column():
+                        gr.HTML(f'''<h2>6. Search for documents containing n-gram(s)</h2>''')
+                        with gr.Accordion(label='Click to view instructions', open=False):
+                            gr.HTML(f'''<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>
+                                        <br>
+                                        <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>
+                                        <br>
+                                        <p style="font-size: 16px;">If you want another batch of random documents, simply hit the Submit button again :)</p>
+                                        <br>
+                                        <p style="font-size: 16px;">Notes on CNF queries:</p>
+                                        <ul style="font-size: 16px;">
+                                            <li>A CNF query may contain up to {MAX_CLAUSES_PER_CNF} clauses, and each clause may contain up to {MAX_TERMS_PER_CLAUSE} n-gram 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>In AND queries, we can only examine co-occurrences where adjacent clauses are separated by no more than {max_diff_tokens} tokens. This value can be adjusted within range [1, {MAX_DIFF_TOKENS}] in "Advanced options".</li>
+                                            <li>In AND queries, if a clause has more than {max_clause_freq} matches, we will estimate the count by examining a random subset of {max_clause_freq} documents out of all documents containing that clause. This value can be adjusted within range [1, {MAX_CLAUSE_FREQ}] in "Advanced options".</li>
+                                            <li>The above subsampling mechanism might cause a zero count on co-occurrences of some simple n-grams (e.g., <b>birds AND oil</b>).</li>
+                                        </ul>
+                                        <br>
+                                        <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):
+                                search_docs_query = gr.Textbox(placeholder='Enter a query here', label='Query', interactive=True)
+                                search_docs_maxnum = gr.Slider(minimum=1, maximum=MAXNUM, value=maxnum, step=1, label='Number of documents to display')
+                                search_docs_max_disp_len = gr.Slider(minimum=1, maximum=MAX_DISP_LEN, value=max_disp_len, step=1, label='Number of tokens to display')
+                                with gr.Accordion(label='Advanced options', open=False):
+                                    with gr.Row():
+                                        search_docs_max_clause_freq = gr.Slider(minimum=1, maximum=MAX_CLAUSE_FREQ, value=max_clause_freq, step=1, label='max_clause_freq')
+                                        search_docs_max_diff_tokens = gr.Slider(minimum=1, maximum=MAX_DIFF_TOKENS, value=max_diff_tokens, step=1, label='max_diff_tokens')
+                                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=1, interactive=False)
+                            with gr.Column(scale=2):
+                                search_docs_message = gr.Label(label='Message', num_top_classes=0)
+                                search_docs_metadatas = []
+                                search_docs_outputs = []
+                                for i in range(MAXNUM):
+                                    with gr.Tab(label=str(i+1)):
+                                        search_docs_metadatas.append(gr.Textbox(label='Metadata', lines=3, interactive=False))
+                                        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_metadatas + search_docs_outputs)
+                    search_docs_submit.click(search_docs, inputs=[index_desc, search_docs_query, search_docs_maxnum, search_docs_max_disp_len, search_docs_max_clause_freq, search_docs_max_diff_tokens], outputs=[search_docs_latency, search_docs_tokenized, search_docs_message] + search_docs_metadatas + search_docs_outputs, api_name=False)
+
+                with gr.Tab('6. Search documents'):
+                    with gr.Column():
+                        gr.HTML(f'''<h2>6. Search for documents containing n-gram(s)</h2>''')
+                        with gr.Accordion(label='Click to view instructions', open=False):
+                            gr.HTML(f'''<p style="font-size: 16px;">This displays the 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>
+                                        <br>
+                                        <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>
+                                        <br>
+                                        <p style="font-size: 16px;">Notes on CNF queries:</p>
+                                        <ul style="font-size: 16px;">
+                                            <li>A CNF query may contain up to {MAX_CLAUSES_PER_CNF} clauses, and each clause may contain up to {MAX_TERMS_PER_CLAUSE} n-gram 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>In AND queries, we can only examine co-occurrences where adjacent clauses are separated by no more than {max_diff_tokens} tokens. This value can be adjusted within range [1, {MAX_DIFF_TOKENS}] in "Advanced options".</li>
+                                            <li>In AND queries, if a clause has more than {max_clause_freq} matches, we will estimate the count by examining a random subset of {max_clause_freq} occurrences of that clause. This value can be adjusted within range [1, {MAX_CLAUSE_FREQ}] in "Advanced options".</li>
+                                            <li>The above subsampling mechanism might cause a zero count on co-occurrences of some simple n-grams (e.g., <b>birds AND oil</b>).</li>
+                                        </ul>
+                                        <br>
+                                        <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):
+                                search_docs_new_query = gr.Textbox(placeholder='Enter a query here', label='Query', interactive=True)
+                                search_docs_new_max_disp_len = gr.Slider(minimum=1, maximum=MAX_DISP_LEN, value=max_disp_len, step=1, label='Number of tokens to display')
+                                with gr.Accordion(label='Advanced options', open=False):
+                                    with gr.Row():
+                                        search_docs_new_max_clause_freq = gr.Slider(minimum=1, maximum=MAX_CLAUSE_FREQ, value=max_clause_freq, step=1, label='max_clause_freq')
+                                        search_docs_new_max_diff_tokens = gr.Slider(minimum=1, maximum=MAX_DIFF_TOKENS, value=max_diff_tokens, step=1, label='max_diff_tokens')
+                                with gr.Row():
+                                    search_docs_new_clear = gr.ClearButton(value='Clear', variant='secondary', visible=True)
+                                    search_docs_new_submit = gr.Button(value='Submit', variant='primary', visible=True)
+                                search_docs_new_latency = gr.Textbox(label='Latency (milliseconds)', interactive=False, lines=1)
+                                search_docs_new_tokenized = gr.Textbox(label='Tokenized', lines=1, interactive=False)
+                            with gr.Column(scale=2):
+                                search_docs_new_message = gr.Label(label='Message', num_top_classes=0)
+                                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=[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]
+                    )
 
-demo.queue(max_size=12).launch(share=False)
+        with gr.Row():
+            gr.Markdown('''
+If you find this tool useful, please kindly cite our paper:
+```bibtex
+@article{Liu2024InfiniGram,
+  title={Infini-gram: Scaling Unbounded n-gram Language Models to a Trillion Tokens},
+  author={Liu, Jiacheng and Min, Sewon and Zettlemoyer, Luke and Choi, Yejin and Hajishirzi, Hannaneh},
+  journal={arXiv preprint arXiv:2401.17377},
+  year={2024}
+}
+```
+''')
+
+demo.queue(
+    default_concurrency_limit=DEFAULT_CONCURRENCY_LIMIT,
+    max_size=MAX_SIZE,
+    api_open=False,
+).launch(
+    max_threads=MAX_THREADS,
+    debug=DEBUG,
+    show_api=False,
+)