import gradio as gr import json #from llmlingua import PromptCompressor from utils_llmlingua2_test import PromptCompressor import tiktoken compressors = { "xlm-roberta": PromptCompressor( #model_name="microsoft/llmlingua-2-xlm-roberta-large-meetingbank", #model_name="qminh369/token-classification-llmlingua2-xlm-roberta-10k_merge_10_epoch_paper", #model_name='qminh369/token-classification-llmlingua2-xlm-roberta-42k_merge_1_epoch', model_name='qminh369/token-classification-llmlingua2-xlm-roberta-42k_merge_10_epoch', use_llmlingua2=True, device_map="cpu" ) } tokenizer = tiktoken.encoding_for_model("gpt-4") def compress(original_prompt, compression_rate, base_model="xlm-roberta", force_tokens = ['. ', ', '], chunk_end_tokens=['.', '\n']): if '\\n' in force_tokens: idx = force_tokens.index('\\n') force_tokens[idx] = '\n' compressor = compressors.get(base_model, compressors["xlm-roberta"]) results = compressor.compress_prompt_llmlingua2( original_prompt, rate=compression_rate, force_tokens=force_tokens, chunk_end_tokens=chunk_end_tokens, return_word_label=True, drop_consecutive=True, force_reserve_digit=True, ) compressed_prompt = results["compressed_prompt"] n_word_compressed = len(tokenizer.encode(compressed_prompt)) word_sep = "\t\t|\t\t" label_sep = " " lines = results["fn_labeled_original_prompt"].split(word_sep) preserved_tokens = [] for line in lines: word, label = line.split(label_sep) preserved_tokens.append((word, '+') if label == '1' else (word, None)) return compressed_prompt, preserved_tokens, n_word_compressed title = "LLMLingua-2" header = """# LLMLingua-2 """ theme = "soft" css = """#anno-img .mask {opacity: 0.5; transition: all 0.2s ease-in-out;} #anno-img .mask.active {opacity: 0.7}""" original_prompt_text = """""" with gr.Blocks(title=title, css=css) as app: gr.Markdown(header) with gr.Row(): with gr.Column(scale=3): original_prompt = gr.Textbox(value=original_prompt_text, label="Original Prompt", lines=10, max_lines=10, interactive=True) compressed_prompt = gr.Textbox(value='', label="Compressed Prompt", lines=10, max_lines=10, interactive=False) with gr.Column(scale=1): base_model = gr.Radio(["xlm-roberta"], label="Base Model", value="xlm-roberta", interactive=True) force_tokens = gr.Dropdown(['\\n', '.', '!', '?', ','], label="Tokens to Preserve", value=['\\n', '.', '!', '?', ','], multiselect=True, interactive=True) compression_rate = gr.Slider(minimum=0.1, maximum=1.0, step=0.1, value=0.7, label="Compression rate", info="after compr. / befor compr.", interactive=True) n_word_original = gr.Textbox(lines=1, label="Original (GPT-4 Tokens)", interactive=False, value=len(tokenizer.encode(original_prompt_text))) n_word_compressed = gr.Textbox(lines=1, label="Compressed (GPT-4 Tokens)", interactive=False) button = gr.Button("⚡Click to Compress") with gr.Accordion(label="Compression Details", open=False): diff_text = gr.HighlightedText(label="Diff", combine_adjacent=False, show_legend=True, color_map={"+": "green"}) original_prompt.change(lambda x: len(tokenizer.encode(x)), inputs=[original_prompt], outputs=[n_word_original]) original_prompt.change(lambda x: ("", "", []), inputs=[original_prompt], outputs=[compressed_prompt, n_word_compressed, diff_text]) button.click(fn=compress, inputs=[original_prompt, compression_rate, base_model, force_tokens], outputs=[compressed_prompt, diff_text, n_word_compressed]) app.queue(max_size=10, api_open=False).launch(show_api=False)