import gradio as gr import datetime import json import os import requests 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') def process(corpus_desc, query_desc, query): 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, 'corpus': corpus, 'query_type': query_type, 'query': query, } print(json.dumps(data)) 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) return result def process_ard_cnf_multi(corpus_desc, query_desc, query, maxnum): 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, 'corpus': corpus, 'query_type': query_type, 'query': query, 'maxnum': maxnum, } print(json.dumps(data)) 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] outputs = outputs[:maxnum] while len(outputs) < 10: outputs.append([]) return output_tokens, message, outputs[0], outputs[1], outputs[2], outputs[3], outputs[4], outputs[5], outputs[6], outputs[7], outputs[8], outputs[9] with gr.Blocks() as demo: with gr.Column(): gr.HTML( '''

Infini-gram: An Engine for n-gram / ∞-gram Language Models with Trillion-Token Corpora

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.

The engine is developed by Jiacheng (Gary) Liu and documented in our paper: Infini-gram: Scaling Unbounded n-gram Language Models to a Trillion Tokens

HF Paper Page:

''' ) 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], ) with gr.Row(visible=True) as row_1: with gr.Column(): gr.HTML('

1. Count an n-gram

') gr.HTML('

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.

') gr.HTML('

Example query: natural language processing (the output is Cnt(natural language processing))

') 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.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.Row(visible=False) as row_2: with gr.Column(): gr.HTML('

2. Compute the probability of the last token in an n-gram

') gr.HTML('

This computes the n-gram probability of the last token conditioned on the previous tokens (i.e. (n-1)-gram)).

') gr.HTML('

Example query: natural language processing (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)

') gr.HTML('

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.

') 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.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.Row(visible=False) as row_3: with gr.Column(): gr.HTML('

3. Compute the next-token distribution of an (n-1)-gram

') gr.HTML('

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.

') gr.HTML('

Example query: natural language (the output is P(* | natural language), for the top-10 tokens *)

') gr.HTML(f'

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.

') 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.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.Row(visible=False) as row_4: with gr.Column(): gr.HTML('

4. Compute the ∞-gram probability of the last token

') gr.HTML('

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.

') gr.HTML('

Example query: I love natural language processing (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)

') gr.HTML('

Note: It may be possible that the effective n = 1, in which case it reduces to the uni-gram probability of the last token.

') with gr.Row(): with gr.Column(scale=1): infgram_input = gr.Textbox(placeholder='Enter a string here', label='Query', interactive=True) 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.Row(visible=False) as row_5: with gr.Column(): gr.HTML('

5. Compute the ∞-gram next-token distribution

') gr.HTML('

This is similar to Query 3, but with ∞-gram instead of n-gram.

') gr.HTML('

Example query: I love natural language (the output is P(* | natural language), for the top-10 tokens *)

') with gr.Row(): with gr.Column(scale=1): infntd_input = gr.Textbox(placeholder='Enter a string here', label='Query', interactive=True) 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.Row(visible=False) as row_6: # with gr.Column(): # gr.HTML(f'''

6. Searching for document containing n-gram(s)

#

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 CNF format, in which case the displayed document contains n-grams such that it satisfies this logical constraint.

#

Example queries:

# #

If you want another random document, simply hit the Submit button again :)

#

A few notes:

# #

❗️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 :)

# ''') # 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'''

6. Searching for documents containing n-gram(s)

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 CNF format, in which case the displayed document contains n-grams such that it satisfies this logical constraint.

Example queries:

If you want another batch of random documents, simply hit the Submit button again :)

A few notes:

❗️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 :)

''') 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.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.Row(visible=False) as row_7: with gr.Column(): gr.HTML('

7. Analyze an (AI-generated) document using ∞-gram

') gr.HTML('

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).

') gr.HTML('

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).

') 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.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.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} } ``` ''') 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, inputs=[corpus_desc, query_desc, count_input], outputs=[count_output, count_output_tokens], api_name=False) ngram_submit.click(process, inputs=[corpus_desc, query_desc, ngram_input], outputs=[ngram_output, ngram_output_tokens], api_name=False) ntd_submit.click(process, inputs=[corpus_desc, query_desc, ntd_input], outputs=[ntd_output, ntd_output_tokens], api_name=False) infgram_submit.click(process, inputs=[corpus_desc, query_desc, infgram_input], outputs=[infgram_output, infgram_output_tokens, infgram_longest_suffix], api_name=False) infntd_submit.click(process, 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_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], api_name=False) doc_analysis_submit.click(process, 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']) 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, ) for d in gr.context.Context.root_block.dependencies: d['api_name'] = False if debug: print(gr.context.Context.root_block.dependencies)